-
Notifications
You must be signed in to change notification settings - Fork 253
[Klaud Cold][agentic experiment][Variant E] Kimi-K3 B200 agg TP8xPP2 agentic — direct vllm serve + SimpleCPU KV offload / Kimi-K3 B200 聚合式 TP8xPP2 智能体实验——直接 vllm serve + SimpleCPU KV 卸载 #2370
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
functionstackx
wants to merge
21
commits into
main
from
klaud/kimik3-b200-agentic-direct-vllm-simplecpu
Closed
Changes from all commits
Commits
Show all changes
21 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 8fd319c
feat: agentic experiment D — direct vllm serve via srt-slurm PR #278
functionstackx c0ace4d
docs: point changelog and MODELS rows at experiment PR #2359
functionstackx 862024d
fix: drop gpu-memory-utilization to 0.90 (flashinfer MoE workspace OOM)
functionstackx 4370988
fix: expandable_segments allocator, drop NCCL_CUMEM_ENABLE (prefill OOM)
functionstackx 0c5fe11
Revert "fix: expandable_segments allocator, drop NCCL_CUMEM_ENABLE (p…
functionstackx e675dd2
feat: add VLLM_PREFIX_CACHE_RETENTION_INTERVAL=32768
functionstackx 4b0c3a4
feat: widen agentic conc list to 1/8/16/32
functionstackx 479b74b
fix: remove VLLM_PREFIX_CACHE_RETENTION_INTERVAL (K3 scheduler block)
functionstackx 77988f7
feat: agentic experiment E — SimpleCPUOffloadConnector CPU DRAM KV of…
functionstackx 833d4e8
docs: point changelog and MODELS rows at experiment PR #2370
functionstackx 05eae6f
fix: disable cross-layer blocks in SimpleCPUOffloadConnector (K3 hybrid)
functionstackx b99c9ce
fix: expandable_segments on the SimpleCPU-offload variant (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 |
62 changes: 62 additions & 0 deletions
62
benchmarks/multi_node/srt-slurm-recipes/patches/srt-slurm-pr278-direct-vllm-multinode.patch
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,62 @@ | ||
| diff --git a/src/srtctl/backends/vllm.py b/src/srtctl/backends/vllm.py | ||
| index 74f673b..377606a 100644 | ||
| --- a/src/srtctl/backends/vllm.py | ||
| +++ b/src/srtctl/backends/vllm.py | ||
| @@ -716,25 +716,31 @@ class VLLMProtocol: | ||
| if frontend_type == "vllm": | ||
| if mode != "agg": | ||
| raise ValueError("frontend.type: vllm supports aggregate vLLM jobs only") | ||
| - if is_multi_node: | ||
| - raise ValueError("frontend.type: vllm currently supports single-node aggregate jobs only") | ||
|
|
||
| config.pop("host", None) | ||
| config.pop("port", None) | ||
| config.pop("connector", None) | ||
| config.setdefault("served-model-name", served_model_name) | ||
|
|
||
| - cmd.extend( | ||
| - [ | ||
| - "vllm", | ||
| - "serve", | ||
| - model_arg, | ||
| - "--host", | ||
| - "0.0.0.0", | ||
| - "--port", | ||
| - str(runtime.frontend_port), | ||
| - ] | ||
| - ) | ||
| + node_rank = endpoint_nodes.index(process.node) | ||
| + cmd.extend(["vllm", "serve", model_arg]) | ||
| + if node_rank == 0: | ||
| + cmd.extend(["--host", "0.0.0.0", "--port", str(runtime.frontend_port)]) | ||
| + if is_multi_node: | ||
| + # vLLM-native multi-node serve (torchrun-style): the leader owns | ||
| + # the OpenAI server; other node ranks run headless engine workers. | ||
| + cmd.extend( | ||
| + [ | ||
| + "--master-addr", | ||
| + leader_ip, | ||
| + "--nnodes", | ||
| + str(len(endpoint_nodes)), | ||
| + "--node-rank", | ||
| + str(node_rank), | ||
| + ] | ||
| + ) | ||
| + if node_rank > 0: | ||
| + cmd.append("--headless") | ||
| if not self.set_cuda_visible_devices: | ||
| device_ids = ",".join(str(i) for i in sorted(process.gpu_indices)) | ||
| if device_ids: | ||
| diff --git a/src/srtctl/core/schema.py b/src/srtctl/core/schema.py | ||
| index 1263ddc..0ef7ae4 100644 | ||
| --- a/src/srtctl/core/schema.py | ||
| +++ b/src/srtctl/core/schema.py | ||
| @@ -1587,8 +1587,6 @@ class SrtConfig: | ||
| raise ValidationError("frontend.type: vllm supports aggregate jobs only, not disaggregated layouts") | ||
| if self.resources.num_agg < 1: | ||
| raise ValidationError("frontend.type: vllm requires resources.agg_workers >= 1") | ||
| - if (self.resources.agg_nodes or 1) != 1: | ||
| - raise ValidationError("frontend.type: vllm currently supports single-node aggregate jobs only") | ||
|
|
||
| def _validate_het_jobs(self): | ||
| """When ``resources.het_jobs`` is set to True, enforce supported shape. |
165 changes: 165 additions & 0 deletions
165
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,165 @@ | ||
| name: "kimik3-vllm-agg-b200-tp8pp2-simplecpu-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" | ||
|
|
||
| # Direct vLLM serving (frontend.type: vllm, srt-slurm PR #278 + the | ||
| # InferenceX multinode patch): `vllm serve` owns the OpenAI port itself, so | ||
| # no Dynamo frontend/worker is involved and no dynamo install is needed. | ||
| dynamo: | ||
| install: false | ||
|
|
||
| # 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 | ||
|
|
||
| environment: | ||
| # SimpleCPUOffloadConnector: identical prefixes must hash to identical | ||
| # block keys across ranks. | ||
| PYTHONHASHSEED: "42" | ||
|
|
||
| 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: | ||
| # Direct vLLM OpenAI server (srt-slurm PR #278): the vllm serve leader owns | ||
| # the public port; rank-1 runs a headless engine worker (vLLM-native | ||
| # multi-node TP8xPP2 via --master-addr/--nnodes/--node-rank, enabled by | ||
| # patches/srt-slurm-pr278-direct-vllm-multinode.patch). | ||
| type: vllm | ||
| 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" | ||
| # Evidence from THIS variant (conc16, SimpleCPU offload): a 2.36 GiB MLA | ||
| # prefill transient OOM'd while 2.63 GiB sat reserved-but-unallocated — | ||
| # allocator fragmentation, exactly what this mode fixes (and what the | ||
| # torch OOM message recommends). The GPU-resident variant D runs green | ||
| # without it; the offload connector's staging shifts the layout enough | ||
| # to fragment. NCCL_CUMEM_ENABLE stays unchanged. | ||
| PYTORCH_CUDA_ALLOC_CONF: "expandable_segments:True" | ||
| # No VLLM_PREFIX_CACHE_RETENTION_INTERVAL: the GB200/GB300 AgentX value | ||
| # (32768) hard-fails engine init on Kimi-K3 — the KDA hybrid gives it a | ||
| # scheduler_block_size of 3145728 and the interval must be a multiple of | ||
| # it ("VLLM_PREFIX_CACHE_RETENTION_INTERVAL (32768) must be non-negative | ||
| # and a multiple of scheduler_block_size (3145728)"). Default retention | ||
| # served fine in earlier runs. | ||
| NCCL_CUMEM_ENABLE: "1" | ||
| 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 | ||
| # SimpleCPUOffloadConnector (variant E): offload KV blocks to CPU DRAM. | ||
| # enable_cross_layers_blocks false: with true, KV-cache init crashes on | ||
| # K3's KDA hybrid — "shape '[64980, 64, 576]' is invalid for input of | ||
| # size 199618560", a 12x mismatch equal to the MLA layers per PP stage | ||
| # (24/2); the cross-layer block folding gets the hybrid geometry wrong. | ||
| # The GB200 recipes also run connectors with cross-layer blocks off. | ||
| # cpu_bytes_to_use_per_rank = TOTAL_CPU_DRAM_GB * 1e9 / GPU_COUNT with | ||
| # the framework's own budget math for cluster:b200-dgxc — DRAM | ||
| # min(3,095,781 MiB, 2,861,022 MiB cap) * 0.80 utilization = 2399 GB | ||
| # per node (matrix total-cpu-dram-gb), / 8 ranks = 299,875,000,000 | ||
| # bytes per rank. | ||
| kv-transfer-config: '{"kv_connector":"SimpleCPUOffloadConnector","kv_role":"kv_both","kv_connector_extra_config":{"cpu_bytes_to_use_per_rank":299875000000,"lazy_offload":false,"enable_cross_layers_blocks":false}}' | ||
| # Explicit (default-on) — offloaded blocks are only useful if prefix | ||
| # caching reuses them across agentic turns. | ||
| enable-prefix-caching: 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 (seen on the dynamo-frontend variants). 0.90 | ||
| # matches the GB200/GB300 agentic recipes. | ||
| gpu-memory-utilization: 0.90 | ||
| no-enable-flashinfer-autotune: true | ||
| # kimi_k3 parsers via the native vllm serve OpenAI-frontend flags — | ||
| # legitimate here because this recipe serves directly with vllm serve | ||
| # (frontend.type: vllm), not through the dynamo worker entrypoint that | ||
| # rejects them. | ||
| enable-auto-tool-choice: true | ||
| tool-call-parser: kimi_k3 | ||
| 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
Oops, something went wrong.
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.
🟡 MODELS.md and MODELS_zh.md are edited from a stale pre-#2359 base, so this PR's diff repoints the Kimi-K3 'Date added' link from #2359 (the PR that actually onboarded the model) to #2370 (this PR, which is just a follow-on offloading experiment on the already-added model). This misattributes when/where Kimi-K3 entered the benchmark suite in both files; the link should stay on #2359.
Extended reasoning...
What the bug is: Both
MODELS.mdline 21 andMODELS_zh.mdline 21 change the Kimi-K3 row's 'Date added' link. The diff's before context is a bare, unlinked2026-07-27, and the after is2026-07-27 ([#2370](...)). But on currentmain(verified viagit show HEAD:MODELS.md/HEAD:MODELS_zh.md), that row already reads2026-07-27 ([#2359](...))— this was set by an earlier commit that landed after this PR's branch point. Because this PR's working tree is stale relative to main, its diff is computed against the old bare-date version, and applying it as-is would overwrite the correct #2359 link with an incorrect #2370 link.\n\nWhy #2359 vs #2370 matters: the PR's own description spells out the lineage explicitly: '#2359' is 'Variant D (base): direct vllm serve, GPU-resident KV', and this PR (#2370) is 'Variant E — Variant D + SimpleCPUOffloadConnector CPU DRAM offload'. In other words, #2359 is the PR that actually onboarded Kimi-K3 to the benchmark suite (added the model, container, recipes, master-config entry for the first time), while #2370 is a follow-on experiment layered on top of an already-added model — just a KV-offloading variant of an existing config. The 'Date added' column's established convention (see GLM-5.2 → #2268, MiniMax-M3 → #1724 in the same table) is to link the PR that onboarded the model, not every subsequent experiment PR that touches its config.\n\nWhy nothing else in the PR catches this: the change lands as an incidental side effect of stale-branch diffing, not a deliberate edit — there is no test or lint that checks documentation attribution links, so it would merge silently (or, more likely, surface as a same-line merge conflict against main's #2359 edit, since both sides changed the same line from the same bare-date ancestor — but a human resolving that conflict could easily default to 'ours'/#2370 without noticing the semantic error).\n\nStep-by-step proof:\n1. Rungit show HEAD:MODELS.md | sed -n '"21p"'on current main → outputs| Kimi-K3 | \kimik3` | 2026-07-27 (#2359) | Agentic coding | |.\n2. This PR's diff hunk for that same line shows a-side of2026-07-27 | Agentic coding(no link at all) and a+side of2026-07-27 (#2370)— proving the diff was authored against a version of the file that predates the #2359 attribution commit.\n3. Read this PR's own description: '#2359' is labeled 'Variant D (base)' and described as the PR onboarding Kimi-K3 with direct vllm serve; this PR (#2370) is labeled 'Variant E', an experiment adding CPU DRAM KV offload on top of the Variant D config — i.e., #2370 doesn't onboard the model, it experiments on an already-onboarded one.\n4. Therefore replacing #2359 with #2370 in the Date-added column is factually wrong regardless of merge mechanics, and merging this PR as authored would regress the currently-correct attribution.\n\n**Fix:** keep (or restore) the link to #2359 in bothMODELS.mdandMODELS_zh.md` line 21 — i.e., drop this PR's edit to that line entirely, or rebase onto current main so the line is left untouched.