Add MoonEP expert-parallel arm of the Kimi-K3 B300 agentic sweep / 新增 Kimi-K3 B300 agentic sweep 的 MoonEP 专家并行分支 - #2453
Add MoonEP expert-parallel arm of the Kimi-K3 B300 agentic sweep / 新增 Kimi-K3 B300 agentic sweep 的 MoonEP 专家并行分支#2453Oseltamivir wants to merge 2 commits into
Conversation
…rk sweep Mirrors kimik3-fp4-b300-vllm-agentic-dspark (same image, model, draft head, spec config, KV arms and concurrency ladder) with the MoE layers served expert-parallel through MoonEP at TP8/EP8 instead of tensor-parallel. Both keys resolve to the same benchmark script, so it now branches on EP_SIZE: unset/1 keeps the pure-TP8 profile byte-for-byte, EP_SIZE=8 takes the MoonEP path. The stock image has neither the MoonEP package nor a vLLM that knows the backend, so the recipe builds MoonEP at a pinned commit and overlays a checked-in vLLM diff inside the container, failing the job on any rejected hunk. Two settings differ from the -dspark sibling, both measured on b300: --load-format auto (fastsafetensors' staging buffers OOM against MoonEP's out-of-band VMM expert allocations) and gpu-memory-utilization 0.96 (at 0.90 the engine measures -2.36 GiB of available KV and refuses to start; 0.96 yields 508,586 KV tokens). Validated end to end before this change: Kimi-K3 serves on the backend at TP8/EP8, cudagraph capture and replay work, and greedy output matches a plain TP8 baseline. 中文：新增 kimik3-fp4-b300-vllm-agentic-dspark 的 MoonEP 专家并行(EP)分支，镜像原 配置的镜像、模型、draft head、投机解码配置、KV 方案与并发梯度，区别在于 MoE 层以 TP8/EP8 的专家并行方式服务，而非张量并行。两个 config key 解析到同一个基准测试脚 本，脚本按 EP_SIZE 分支：未设置或为 1 时完全保持原有纯 TP8 路径不变，EP_SIZE=8 时 走 MoonEP 路径。由于官方镜像既不含 MoonEP 包，也不含支持 moonep 后端的 vLLM，脚本 在容器内按固定 commit 构建 MoonEP 并叠加仓库内的 vLLM 补丁，任一 hunk 应用失败即 让任务失败。相对原配置有两处经 b300 实测的差异：--load-format auto（fastsafetensors 的大暂存缓冲区会与 MoonEP 在 torch 分配器之外的 VMM 专家权重争抢显存导致 OOM），以及 gpu-memory-utilization 0.96（0.90 时引擎测得可用 KV 为 -2.36 GiB 无法启动，0.96 可得 508,586 KV tokens）。本次改动前已完成端到端验证：Kimi-K3 在该后端上以 TP8/EP8 正常 服务，cudagraph 捕获与重放正常，贪心解码输出与纯 TP8 基线一致。
|
Thanks for the contribution! Please reach out to respective companies' CODEOWNER to fill in the latest PR_REVIEW_CHECKLIST.md before pinging core maintainer on Slack for review. In order for the signoff PR check bot to trigger, you must follow the PR_REVIEW_CHECKLIST.md template correctly, including the phrase For PR verification, add the PR authors are responsible for ensuring that after merging, all GitHub Action jobs fully pass. A lot of the time, failures are just flakes and simply re-running the failed jobs will fix it. See GitHub's docs on re-running failed jobs 感谢你的贡献！请联系相应公司的 CODEOWNER 填写最新的 PR_REVIEW_CHECKLIST.md，然后再在 Slack 上联系核心维护者进行审阅。为了触发 signoff PR 检查机器人，你必须正确遵循 PR_REVIEW_CHECKLIST.md 模板，包括保留英文语句 如需进行 PR 验证，请为此 PR 添加 PR 作者有责任确保合并后所有 GitHub Action 任务完全通过。 很多时候失败只是偶发抖动（flake），重新运行失败的任务即可解决。参见 GitHub 关于重新运行失败任务的文档 |
1 similar comment
|
Thanks for the contribution! Please reach out to respective companies' CODEOWNER to fill in the latest PR_REVIEW_CHECKLIST.md before pinging core maintainer on Slack for review. In order for the signoff PR check bot to trigger, you must follow the PR_REVIEW_CHECKLIST.md template correctly, including the phrase For PR verification, add the PR authors are responsible for ensuring that after merging, all GitHub Action jobs fully pass. A lot of the time, failures are just flakes and simply re-running the failed jobs will fix it. See GitHub's docs on re-running failed jobs 感谢你的贡献！请联系相应公司的 CODEOWNER 填写最新的 PR_REVIEW_CHECKLIST.md，然后再在 Slack 上联系核心维护者进行审阅。为了触发 signoff PR 检查机器人，你必须正确遵循 PR_REVIEW_CHECKLIST.md 模板，包括保留英文语句 如需进行 PR 验证，请为此 PR 添加 PR 作者有责任确保合并后所有 GitHub Action 任务完全通过。 很多时候失败只是偶发抖动（flake），重新运行失败的任务即可解决。参见 GitHub 关于重新运行失败任务的文档 |
中文：补全 perf-changelog 中的 PR 链接。
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Deferring given the scope of this change (new MoonEP EP backend, VMM-backed symmetric expert weights, and a runtime git-clone/pip-install/patch flow) — worth a human look beyond the inline pr-link finding. Beyond that finding, I also checked whether the MoonEP MXFP4 path rebinds w13_weight/w2_weight to the globalized symmetric buffer after publish_converted, and whether the launcher's EP_SIZE != TP guard could silently pass for a non-canonical numeric string — neither held up as a real issue.
Extended reasoning...
This PR is large and touches production-critical serving code: a brand-new all2all backend (MoonEP), CUDA VMM symmetric memory for cross-rank expert weights, and a container-runtime flow that clones an external repo, builds it from source, and patches the installed vLLM package before serving. That scope alone warrants a human look independent of any specific bugs. The bug-hunting pass for this run confirmed one real issue (the appended perf-changelog.yaml entry has pr-link set to the placeholder pull/XXXX, which fails the repo's own changelog validator — flagged inline). It also raised and then refuted two other candidate correctness concerns in the MoonEP integration: whether layer.w13_weight/w2_weight get correctly rebound to the globalized symmetric buffer after publish_converted, and whether the shell guard requiring EP_SIZE == TP could be bypassed by a non-canonical numeric string. Neither of those held up under verification.
|
see unofficial run visualizer at https://inferencex.semianalysis.com/inference?unofficialRun=30753806837 |
There was a problem hiding this comment.
No new bugs found this run, but I'm deferring to a human reviewer given the scope of this change: it adds an entirely new vLLM all2all backend (MoonEP) via a ~1,700-line patch touching core MoE dispatch/weight-loading paths, plus CUDA VMM symmetric-memory allocation and custom Triton kernels, applied by patching an installed package at container runtime.
What was reviewed: the EP_SIZE branch in the benchmark script (byte-for-byte fallback to the pure-TP8 path when unset/1), the MoonEP all2all manager, prepare/finalize dispatch/combine logic, and the MXFP4 symmetric-weight allocation/publishing path in the patch. The candidate concern that MoonEP's symmetric weights are never bound onto the layer was checked and ruled out -- process_weights_after_loading rebinds layer.w13_weight/w2_weight to the full symmetric mappings after the per-rank scale transform. Separately, the perf-changelog pr-link placeholder flagged in an earlier pass has since been fixed (commit 665638b).
Extended reasoning...
Overview
This PR adds a MoonEP expert-parallel arm of the Kimi-K3 B300 agentic sweep. Three of the four changed files are benchmark plumbing (shell script EP_SIZE branch, master-config entry, changelog), but the fourth is a new checked-in .patch file that adds an entirely new vLLM all2all backend: a MoonEPAll2AllManager, a MoonEPPrepareAndFinalize prepare/finalize implementation with a custom Triton kernel for building DeepGEMM m_indices, a new moonep_weights.py module doing CUDA VMM symmetric-memory allocation across ranks, and a new MXFP4 MoE method that publishes/binds those symmetric weights. This patch is applied at container runtime via patch -p1 onto the installed vLLM package inside the image, with a fail-fast on any rejected hunk.
Security risks
No injection/auth/credential-exposure surface here -- this is single-node benchmark infrastructure, not a production request path. The main correctness risk is in the memory-safety class: cross-rank NVLink/VMM symmetric memory reads (a rank routinely reads a peer's HBM for experts it does not own), and a runtime source patch with a fail-fast gate. Both are already defended by asserts and hard failures in the diff (patch-apply verified via has_moonep()/import checks, VMM granularity asserts, stride asserts on the scale round-trip), which is reassuring but does not substitute for a second set of eyes given the amount of new low-level code.
Level of scrutiny
This warrants a higher bar than most single-node recipe PRs. It is scoped to a benchmark-only surface (not upstream vLLM, not this repo's production serving path), and the PR description documents substantial validation (correct completions, cudagraph capture/replay bit-identical to eager, greedy match against a TP8 baseline, MoonEP's own kernel suite passing, cross-rank weight bitwise-identity check). But the sheer size and novelty of the patch -- a new memory-sharing scheme, a new grouped-GEMM entry path, and a first-of-its-kind checked-in .patch convention for single-node recipes -- are exactly the kind of large-scale design decisions this review process should route to a human rather than rubber-stamp.
Other factors
The bug-hunting system found nothing this run, and the one issue it previously flagged (an unfilled pull/XXXX changelog placeholder) has been fixed in a follow-up commit. The PR author's own "Known caveats" section is candid about two loose ends (exp-name collision with the -dspark sibling, and ep: 8 doubling as an implicit backend selector with no real schema field) that a maintainer may want to weigh in on before merge.
1 participant
Summary
Adds
kimik3-fp4-b300-vllm-agentic-dspark-moonep, a MoonEP expert-parallel arm of the existingkimik3-fp4-b300-vllm-agentic-dsparksweep. Same image, model, DSpark draft head, spec config, KV arms and concurrency ladder as run 30425131777; the MoE layers are served expert-parallel at TP8/EP8 instead of tensor-parallel.MoonEP is Moonshot's expert-parallel dispatch/combine library from the Kimi-K3 tech report (§5.2.1). It holds every rank at exactly
S·Ktokens by reassigning surplus from overloaded expert owners, which makes per-layer GEMM shapes static and removes the per-layer host sync. It communicates over CUDA VMM plus NVSwitch multicast, so it is single-NVLink-domain only — which one B300 node is.What changes
benchmarks/single_node/agentic/kimik3_fp4_b300_vllm_mtp.shEP_SIZE;unset/1keeps the pure-TP8 profile byte-for-byte,8takes the MoonEP pathbenchmarks/single_node/agentic/patches/kimik3_moonep_vllm.patchmoonepbackend, overlaid into the container at runtimeconfigs/nvidia-master.yamlconc 1/2/4/8/16×kv none/kv dram) withep: 8perf-changelog.yamlBoth config keys resolve to the same benchmark script — the launcher derives the script name from model-prefix/precision/framework/spec-decoding, none of which differ — hence the
EP_SIZEbranch rather than a second script.The stock image contains neither the MoonEP package nor a vLLM that knows the backend, so the recipe builds MoonEP at a pinned commit inside the container (with a one-line torch-2.13 fix, upstreamed as MoonshotAI/MoonEP#25) and applies the checked-in vLLM diff, failing the job on any rejected hunk rather than silently serving an unpatched engine.
Two settings differ from the
-dsparksibling, both measured--load-format autoinstead of fastsafetensors — its large staging buffers OOM against MoonEP's out-of-band VMM expert-weight allocations.autoloads the ~1.5 TB checkpoint in ~231 s off the staged mount.--gpu-memory-utilization 0.96on this arm only — MoonEP's expert weights are CUDA VMM allocations outside the torch caching allocator. At 0.90 the engine measures −2.36 GiB of available KV and refuses to start; 0.96 yields 508,586 KV tokens.Validation before opening this
Kimi-K3 was brought up on the backend on b300 and serves correctly at TP8/EP8:
SERVER READY, correct completions ("List the first five prime numbers:"→" 2, 3, 5, 7, 11. The sum of these primes is 28…")generate_sweep_configs.py test-configemits exactly the 10 expected cells.Known caveats
ep, so this config's exp-names collide with the-dsparksibling's. Result artifacts stay distinct (RESULT_FILENAMEembedsep8,process_result.pyrecordsep), but GitHub job display names are ambiguous between the two.ep: 8is what selects the MoonEP branch, because the config schema has noall2all-backendfield. Fine while there is one EP backend; a second would need a real selector..patchunderbenchmarks/single_node/agentic/patches/is a new convention for single-node recipes (multi-node AMD has precedent). The alternative is a custom image.中文说明
新增
kimik3-fp4-b300-vllm-agentic-dspark-moonep，是现有kimik3-fp4-b300-vllm-agentic-dspark的 MoonEP 专家并行(EP)分支。镜像、模型、DSpark draft head、投机解码配置、KV 方案与并发梯度均与 run 30425131777 一致，区别在于 MoE 层以 TP8/EP8 的专家并行方式服务，而非张量并行。MoonEP 是 Moonshot 在 Kimi-K3 技术报告(§5.2.1)中提出的专家并行 dispatch/combine 库。它通过将过载专家宿主上的溢出 token 重新分配给负载不足的 rank，使每个 rank 恰好收到
S·K个 token，从而让每层的 GEMM 形状保持静态，并消除每层的 host 同步。其通信基于 CUDA VMM 与 NVSwitch multicast，因此仅支持单一 NVLink 域——单台 B300 节点正好满足。改动内容： 基准测试脚本按
EP_SIZE分支（未设置或为 1 时完全保持原纯 TP8 路径不变，为 8 时走 MoonEP 路径）；新增仓库内 vLLM 补丁，在容器内运行时叠加；master config 新增 key，共 10 个 cell；perf-changelog 追加条目。两个 config key 解析到同一脚本（launcher 依据 model-prefix/precision/framework/spec-decoding 推导脚本名，而这些字段完全相同），因此采用EP_SIZE分支而非新建脚本。官方镜像既无 MoonEP 包，也无支持该后端的 vLLM，脚本因此在容器内按固定 commit 构建 MoonEP（含一处 torch 2.13 修复，已上游至 MoonshotAI/MoonEP#25）并应用补丁，任一 hunk 应用失败即让任务失败，避免静默使用未打补丁的引擎。两处实测差异：
--load-format auto（fastsafetensors 的大暂存缓冲区会与 MoonEP 在 torch 分配器之外的 VMM 专家权重争抢显存导致 OOM，auto 从预置挂载点约 231 秒加载完 ~1.5 TB 权重）；--gpu-memory-utilization 0.96（仅此分支——MoonEP 专家权重是 torch 缓存分配器之外的 CUDA VMM 分配，0.90 时引擎测得可用 KV 为 −2.36 GiB 无法启动，0.96 可得 508,586 KV tokens）。开 PR 前的验证： Kimi-K3 已在 b300 上以 TP8/EP8 正常服务，输出正确；cudagraph 捕获与重放正常且重放输出与 eager 逐字节一致；贪心解码输出与纯 TP8 基线一致；MoonEP 自带 kernel 测试（planning/dispatch/combine）在 8×B300、镜像自带 cutlass-dsl 4.6.0 下全部通过；对称内存专家权重在跨 rank 读取对端 HBM 时验证为逐位一致。
generate_sweep_configs.py test-config生成的 cell 数量与预期的 10 个完全一致。已知问题： 单节点 agentic 的 exp-name 不编码
ep，与-dspark同名（结果文件仍可区分，但 GitHub job 名称有歧义）；ep: 8承担了后端选择器的角色，因为 config schema 没有all2all-backend字段，若将来引入第二个 EP 后端则需要真正的选择器；在benchmarks/single_node/agentic/patches/下提交.patch对单节点 recipe 是新约定（多节点 AMD 已有先例），另一种方案是构建定制镜像。