-
Notifications
You must be signed in to change notification settings - Fork 249
[Klaud Cold][agentic experiment][Variant A] feat: add Kimi-K3 MXFP4 B200 aggregated TP8xPP2 Dynamo-vLLM agentic bring-up / 新增 Kimi-K3 MXFP4 B200 聚合式 TP8xPP2 Dynamo-vLLM 智能体编码基准测试（bring-up） #2355
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Closed
Closed
Changes from all commits
Commits
Show all changes
11 commits
Select commit
Hold shift + click to select a range
b4fd077
feat: add Kimi-K3 MXFP4 B200 aggregated TP8xPP2 Dynamo-vLLM agentic r…
functionstackx 4dbbdc8
docs: link PR #2355 in changelog entry and MODELS rows
functionstackx e8d42a7
fix: pin agentic srt-slurm to NVIDIA v1.0.36, dynamo 1.2.1, conc-8 sm…
functionstackx c1e2a56
fix: drop OpenAI-frontend tool-choice flags from dynamo-vllm worker args
functionstackx ef35fd1
fix: try dynamo wheel 1.2.0.dev20260426 for Kimi-K3 frontend tokenizer
functionstackx be6c56e
fix: pin dynamo to day-zero Kimi-K3 commit, restore kimi_k3 parser flags
functionstackx c6917e6
fix: use dynamo namespaced --dyn-* kimi_k3 parser args on the worker
functionstackx a26853a
fix: disable aiperf conv-aware routing (session_control 400-rejected)
functionstackx f61eafb
fix: patch kimi-k3 image mamba_hybrid index_fill_ dtype via setup_script
functionstackx fff449c
fix: drop gpu-memory-utilization to 0.90 (flashinfer MoE workspace OOM)
functionstackx 73f7e63
fix: expandable_segments allocator, drop NCCL_CUMEM_ENABLE (prefill OOM)
functionstackx File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
39 changes: 39 additions & 0 deletions
39
benchmarks/multi_node/srt-slurm-recipes/configs/kimi-k3-container-deps.sh
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,39 @@ | ||
| #!/bin/bash | ||
| # Setup script for the Kimi-K3 vLLM bring-up image (vllm/vllm-openai:kimi-k3). | ||
| # srt-slurm runs this in every worker container before dynamo install and | ||
| # worker startup (recipe field: setup_script). | ||
|
|
||
| set -euo pipefail | ||
|
|
||
| # The image's first decode step crashes in the KDA hybrid-state postprocess: | ||
| # vllm/v1/worker/gpu/model_states/mamba_hybrid.py, postprocess_state: | ||
| # IndexError: index_fill_(): Expected dtype int64 for index. | ||
| # torch's index_fill_ requires an int64 index tensor, but the runner passes | ||
| # the int32 idx_mapping (hit by moonshotai/Kimi-K3 agentic bring-up, first | ||
| # decode step, engine v0.1.dev19262+gb6bbf29dd). Coerce the index to int64. | ||
| # Idempotent: exits 0 if the patch is already applied. | ||
| python3 - <<'PY' | ||
| import pathlib | ||
| import re | ||
|
|
||
| import vllm.v1.worker.gpu.model_states.mamba_hybrid as mh | ||
|
|
||
| path = pathlib.Path(mh.__file__) | ||
| src = path.read_text() | ||
| if "idx_mapping.long()" in src: | ||
| print(f"mamba_hybrid index_fill_ patch already applied: {path}") | ||
| raise SystemExit(0) | ||
|
|
||
| new, n = re.subn( | ||
| r"index_fill_\(\s*0,\s*idx_mapping,", | ||
| "index_fill_(0, idx_mapping.long(),", | ||
| src, | ||
| ) | ||
| if n != 1: | ||
| raise SystemExit( | ||
| f"expected exactly one index_fill_(0, idx_mapping, ...) call in " | ||
| f"{path}, found {n} — image layout changed, refusing to patch" | ||
| ) | ||
| path.write_text(new) | ||
| print(f"Patched mamba_hybrid index_fill_ index dtype: {path}") | ||
| PY |
140 changes: 140 additions & 0 deletions
140
benchmarks/multi_node/srt-slurm-recipes/vllm/kimi-k3/agentic/agg-b200-tp8pp2-agentic.yaml
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,140 @@ | ||
| name: "kimik3-vllm-agg-b200-tp8pp2-agentic" | ||
|
|
||
| # Kimi-K3 MXFP4 B200 AGGREGATED TP8 x PP2 agentic recipe (2 nodes / 16 GPUs). | ||
| # The native MXFP4 checkpoint (2.8T total params, ~1.4TB of weights) does not | ||
| # fit one 8xB200 node, so TP8 shards attention/dense (/8) and PP2 splits the | ||
| # 93 layers (/2) across 16 GPUs. Plain TP (NOT TEP): expert parallelism is | ||
| # deliberately off, so the 896 routed experts are TP-sharded inside each | ||
| # pipeline stage. Node allocation = tp*pp/gpus_per_node = 8*2/8 = 2 nodes. | ||
| # Aggregated (single worker, decode num-worker 0) — no P/D split, no NIXL. | ||
| # VLLM_ENABLE_K3_LATENT_MOE_TAIL_FUSION fuses the K3 LatentMoE tail path in | ||
| # the kimi-k3 bring-up image. | ||
| model: | ||
| path: "kimik3" | ||
| container: "vllm/vllm-openai:kimi-k3" | ||
| precision: "fp4" | ||
|
|
||
| identity: | ||
| model: | ||
| repo: "moonshotai/Kimi-K3" | ||
| container: | ||
| image: "vllm/vllm-openai:kimi-k3" | ||
| frameworks: | ||
| dynamo: "1.4.0-kimi-k3-dev.1" | ||
|
|
||
| dynamo: | ||
| install: true | ||
| # Day-zero Kimi-K3 dynamo ("feat: Added support for Kimi-K3", tag | ||
| # v1.4.0-kimi-k3-dev.1 == this commit): adds the kimi_k3 tiktoken tokenizer | ||
| # to the rust frontend (dynamo <=1.2.1 only knows kimi/kimi_k2/kimi_k25/ | ||
| # deepseek_v3, so the model never registers and every request 404s) and the | ||
| # kimi_k3 tool-call/reasoning parser worker args. | ||
| hash: "ba83080ecd31c1ce918559e576d3c5bc9e092ff1" | ||
|
|
||
| # Patches the image's mamba_hybrid postprocess_state: torch index_fill_ | ||
| # requires an int64 index but the runner passes the int32 idx_mapping, | ||
| # crashing the first decode step (IndexError: Expected dtype int64 for index). | ||
| setup_script: kimi-k3-container-deps.sh | ||
|
|
||
| slurm: | ||
| time_limit: "8:00:00" | ||
|
|
||
| health_check: | ||
| interval_seconds: 10 | ||
| max_attempts: 1440 | ||
|
|
||
| resources: | ||
| gpu_type: "b200" | ||
| gpus_per_node: 8 | ||
| agg_nodes: 2 | ||
| agg_workers: 1 | ||
| gpus_per_agg: 16 | ||
|
|
||
| infra: | ||
| etcd_nats_dedicated_node: false | ||
| nats_max_payload_mb: 32 | ||
|
|
||
| frontend: | ||
| type: dynamo | ||
| enable_multiple_frontends: false | ||
|
|
||
| backend: | ||
| type: vllm | ||
| connector: null | ||
| aggregated_environment: | ||
| VLLM_ENABLE_K3_LATENT_MOE_TAIL_FUSION: "1" | ||
| VLLM_SERVER_DEV_MODE: "1" | ||
| # ~1.4TB of MXFP4 weights off shared Lustre: keep the engine-ready window | ||
| # generous, and let one long AgentX request hold a PP stage beyond vLLM's | ||
| # 300-second model-execution default. | ||
| VLLM_ENGINE_READY_TIMEOUT_S: "3600" | ||
| VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: "1800" | ||
| # expandable_segments: the eighth sweep attempt OOM'd on a 2.92 GiB MLA | ||
| # long-context prefill transient (kv_b_proj in _compute_prefill_context) | ||
| # while 3.39 GiB sat reserved-but-unallocated — fragmentation, exactly | ||
| # the case this allocator mode fixes (and what the torch OOM message | ||
| # recommends; the DSv4 recipes set it too). NCCL_CUMEM_ENABLE dropped at | ||
| # the same time to trim NCCL's share of non-PyTorch device memory. | ||
| PYTORCH_CUDA_ALLOC_CONF: "expandable_segments:True" | ||
| TILELANG_CLEANUP_TEMP_FILES: "1" | ||
| UCX_MEMTYPE_CACHE: "n" | ||
| UCX_MEMTYPE_REG_WHOLE: "n" | ||
| UCX_NET_DEVICES: "mlx5_0:1,mlx5_1:1,mlx5_2:1,mlx5_3:1,mlx5_4:1,mlx5_5:1,mlx5_10:1,mlx5_11:1" | ||
| HF_HUB_CACHE: "/hf_hub_cache" | ||
| HUGGINGFACE_HUB_CACHE: "/hf_hub_cache" | ||
| vllm_config: | ||
| aggregated: | ||
| served-model-name: "moonshotai/Kimi-K3" | ||
| tensor-parallel-size: 8 | ||
| pipeline-parallel-size: 2 | ||
| trust-remote-code: true | ||
| load-format: fastsafetensors | ||
| moe-backend: auto | ||
| # 0.90, not 0.95: the flashinfer trtllm MXFP4 MoE kernel allocates a | ||
| # ~1.6 GiB runtime workspace OUTSIDE vLLM's memory pool on the first | ||
| # forward; at 0.95 a 178 GiB B200 has only ~1.35 GiB free and the first | ||
| # warmup request OOMs (seventh sweep attempt). 0.90 matches the | ||
| # GB200/GB300 agentic recipes. | ||
| gpu-memory-utilization: 0.90 | ||
| no-enable-flashinfer-autotune: true | ||
| # kimi_k3 parsers via dynamo's namespaced worker args (--dyn-*), which | ||
| # the day-zero K3 dynamo build pinned above provides. The vLLM | ||
| # OpenAI-frontend spellings (--enable-auto-tool-choice / | ||
| # --tool-call-parser) are rejected by the dynamo worker entrypoint as | ||
| # unrecognized arguments (different arg parser than `vllm serve`). | ||
| dyn-tool-call-parser: kimi_k3 | ||
| reasoning-parser: kimi_k3 | ||
| dyn-reasoning-parser: kimi_k3 | ||
| # No explicit max-model-len: let vLLM derive the native 1M window from | ||
| # the model config (agentic trajectories blow past any small cap, and | ||
| # K3's KDA layers keep per-token KV small — only the 24 gated-MLA | ||
| # layers hold cache). Prefix caching stays on (default) for trajectory | ||
| # reuse. Cap prefill chunks so a single long request cannot OOM a | ||
| # pipeline stage; let vLLM pick max-num-seqs. | ||
| max-num-batched-tokens: 8192 | ||
|
|
||
| sbatch_directives: | ||
| segment: "1" | ||
|
|
||
| srun_options: | ||
| container-remap-root: "" | ||
|
|
||
| benchmark: | ||
| type: custom | ||
| command: bash /infmax-workspace/benchmarks/multi_node/agentic_srt.sh | ||
| env: | ||
| INFMAX_CONTAINER_WORKSPACE: "/infmax-workspace" | ||
| RESULT_DIR: "/logs/agentic" | ||
| PORT: "8000" | ||
| # Keep the aggregate worker in the multinode result schema so ingestion | ||
| # uses the zero decode-worker count instead of duplicating TP into P and D. | ||
| IS_MULTINODE: "true" | ||
| # aiperf's conv-aware routing emits nvext.session_control, a removed POC | ||
| # field this dynamo build 400-rejects at warmup (schema moved to | ||
| # router/routing_constraints/agent_hints). Same opt-out as the GB300 | ||
| # aggregate AgentX recipes — and with a single aggregate worker there is | ||
| # no P/D routing to bind anyway. | ||
| AIPERF_USE_DYNAMO_CONV_AWARE_ROUTING: "0" | ||
| AIPERF_DATASET_MMAP_CACHE_DIR: "/aiperf_mmap_cache" | ||
| HF_HUB_CACHE: "/hf_hub_cache" | ||
| WEKA_LOADER_OVERRIDE: "semianalysis_cc_traces_weka_062126" | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
🔴 The recipe deliberately sets
health_check.max_attempts: 1440(4h) for the ~1.4TB Kimi-K3 MXFP4 checkpoint, butrunners/launch_b200-dgxc.sh:291unconditionally runssed -i 's/^ max_attempts: [0-9]*/ max_attempts: 720/'on the copied config right beforesrtctl apply, silently clobbering it back to 720 attempts (2h) — the same budget sized for DSR1-FP8 at roughly half this checkpoint's weight size. The recipe file will look correctly configured but the wider window never actually takes effect at runtime; either make the sed a floor (only raise, never lower) or special-case kimik3 like the other model-prefix branches nearby.Extended reasoning...
The bug:
runners/launch_b200-dgxc.shline 291 runs:right before
srtctl apply -f "$CONFIG_FILE"inside theIS_MULTINODEbranch. This is an unconditional hard-set, not a max()/floor operation — whatever numeric value followsmax_attempts:(two-space indent) in the config file gets overwritten to exactly720, no matter what it was before.The interaction this PR introduces: the new recipe
benchmarks/multi_node/srt-slurm-recipes/vllm/kimi-k3/agentic/agg-b200-tp8pp2-agentic.yamlsets:with a two-space indent that matches the sed's regex exactly. The recipe's own comments make clear this 1440 (= 4h at 10s/attempt) was deliberately sized up from the launcher's usual value, specifically because the native MXFP4 checkpoint is ~1.4TB and has to be pulled over shared Lustre across 2 nodes. But because the sed is unconditional,
srtctl applynever sees 1440 — it sees the file after the sed has already rewritten it to 720 (= 2h). The comment directly above the sed ("Bump recipe health-check timeout from 360×10s to 720×10s so large-model loads ... finish in time") shows the sed was written with the mental model of bumping the common 360 default up to 720; it never anticipated a recipe explicitly setting a value above 720, so for this recipe the sed is actually a reduction, not a bump.Why nothing else catches this: there's no validation step between writing the recipe's
health_checkblock andsrtctl applythat would surface the discrepancy — the sed runs silently, will echo the sed invocation but not diff the before/after content of the file, and the recipe file committed to the repo (and reviewed in this PR) legitimately shows 1440. Someone reading the recipe or this diff would have no way to know the value gets overwritten at runtime.Step-by-step proof:
CONFIG_FILEresolves torecipes/vllm/kimi-k3/agentic/agg-b200-tp8pp2-agentic.yaml(per theconfigs/nvidia-master.yamlentry'sadditional-settings: CONFIG_FILE=...), a copy of this PR's new recipe overlaid into the srt-slurm-nv clone.max_attempts: 1440.sed -i "s/^name:.*/name: \"${RUNNER_NAME}\"/"(unrelated), then line 291 doessed -i 's/^ max_attempts: [0-9]*/ max_attempts: 720/' "${CONFIG_FILE%%:*}".^ max_attempts: [0-9]*matchesmax_attempts: 1440(2-space indent, then digits), so the line becomesmax_attempts: 720.srtctl apply -f "$CONFIG_FILE"runs immediately after, reading the now-mutated file — it appliesmax_attempts: 720,interval_seconds: 10→ a 7200s (2h) health-check window, not the 14400s (4h) the recipe declares.Impact: if 2h genuinely isn't enough to pull ~1.4TB of MXFP4 weights and bring the engine up (workers loading in parallel off shared Lustre, contended by whatever else is running on the same filesystem), the health check gives up and fails the entire multinode job — the exact scenario this recipe's author was trying to prevent by widening the window, on the largest checkpoint in the fleet and its first-ever bring-up. Because the recipe file looks correct, anyone debugging a health-check failure here would have to know to check the launcher script rather than trusting the recipe as documentation of what will actually run — a debugging detour on a bring-up PR that's likely to already have plenty of other failure modes to sort through.
Fix: make the sed a floor instead of a hard-set, e.g. only replace when the existing value is less than 720 (or compare-and-max in the shell), or skip the sed entirely for
kimik3/any recipe that already sets a larger value, similar to how other model prefixes get their own branch in this same script.