Kimi K3 B200 Day 0 AgentX Pipeline Parallelism (non-DSpark as it doesn't work with Pipeline yet, offloading & TP16 and DEP8PP2 to be done in followup PR) #2391
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,140 @@ | ||
| name: "kimik3-vllm-agg-b200-tp8pp2-retention0-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 | ||
| # multinode extension on functionstackx/srt-slurm-nv branch | ||
| # klaud/direct-vllm-multinode): `vllm serve` owns the OpenAI port itself, so | ||
| # no Dynamo frontend/worker is involved and no dynamo install is needed. | ||
| dynamo: | ||
| install: false | ||
|
|
||
| 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, carried by the | ||
| # functionstackx/srt-slurm-nv klaud/direct-vllm-multinode branch). | ||
| 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" | ||
| # Prefix-cache retention (variant G): 0, on the otherwise-unchanged | ||
| # GPU-resident variant D config. Any positive value must be a multiple of | ||
| # Kimi-K3's KDA-hybrid scheduler_block_size (3145728; the GB recipes' | ||
| # 32768 is hard-rejected at engine init — verified on this PR family), so | ||
| # 0 is the only value below one 3.1M-token scheduler block. | ||
| VLLM_PREFIX_CACHE_RETENTION_INTERVAL: "0" | ||
| # 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. | ||
|
Comment on lines
+66
to
+78
Contributor
There was a problem hiding this comment. 🟡 The comment block right after Extended reasoning...In Immediately after that, lines 72-77 open with "No VLLM_PREFIX_CACHE_RETENTION_INTERVAL: the GB200/GB300 AgentX value (32768) hard-fails engine init..." — i.e., a comment asserting the key is absent from the config and that default retention is used instead. That is flatly false for this file: the key is set, right above, to How this happened: per the PR history (perf-changelog.yaml), this recipe evolved through several variants. In an earlier revision (Variant D era), Step-by-step proof:
Impact: none on runtime behavior — YAML comments are not parsed, so the engine only sees Fix: delete the stale block (lines 72-77) since the correct, up-to-date rationale is already captured in lines 66-70. No functional change needed, just documentation cleanup. This is an explicitly non-merge probe PR (Variant K, meant only to produce an empirical answer about whether an image patch is still needed), so there's no urgency to block on this — flagging as a nit for whenever the recipe is next touched. |
||
| 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 | ||
| 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" | ||
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 sets health_check.max_attempts: 1440 to give the ~1.4TB MXFP4 checkpoint load extra headroom, but the pre-existing sed at runners/launch_b200-dgxc.sh:296 unconditionally rewrites any
max_attempts: Nline to 720, silently halving the recipe's intentional margin back to 7200s. This is a pre-existing launcher behavior (not introduced by this PR) that applies to every recipe, and since this config is otherwise identical to the fully-green Variant G (#2374) which ran under the same 720 override, the 1440 value appears to be dead/misleading rather than functionally necessary — worth reconciling but not blocking.Extended reasoning...
The bug:
benchmarks/multi_node/srt-slurm-recipes/vllm/kimi-k3/agentic/agg-b200-tp8pp2-agentic.yamlsets:at 2-space indentation, intended to give the ~1.4TB MXFP4 Kimi-K3 checkpoint load (off shared Lustre) double the usual health-check window (14400s vs the launcher's usual 7200s).
The code path:
runners/launch_b200-dgxc.sh(line ~296, inside theIS_MULTINODEbranch that this recipe hits sincemultinode: true) runs, right beforesrtctl apply:The regex anchors on exactly two leading spaces before
max_attempts:, which is precisely how the recipe indents it underhealth_check:. This is an unconditional overwrite, not a floor/max clamp — the accompanying comment ('Bump recipe health-check timeout from 360x10s=3600s to 720x10s=7200s ... so large-model loads finish in time') describes it as always raising from a smaller default, and doesn't anticipate a recipe deliberately setting a larger value. So the sed always forces the file to 720 regardless of what was there, silently overwriting the recipe's 1440 down to 720 beforesrtctl applyever reads the config.Why nothing in the codebase prevents this: there's no validation or diffing between the recipe's declared value and what the launcher writes, and the sed runs on every CONFIG_FILE unconditionally for every recipe launched through this script — this recipe has no special-casing to opt out.
Step-by-step proof:
max_attempts: 1440(2-space indent, matches PR diff lines 32-35).sed -i 's/^ max_attempts: [0-9]*/ max_attempts: 720/' "${CONFIG_FILE%%:*}"on that same file.max_attempts: 1440matches^ max_attempts: [0-9]*(2 leading spaces + literal text + digits), so sed rewrites it in place tomax_attempts: 720.srtctl apply -f "$CONFIG_FILE"is invoked immediately after, reading the now-mutated 720 value.Impact: if the 1.4TB MXFP4 weight load off shared Lustre ever exceeds 7200s, the health check would exhaust and kill the job before it ever reaches the decode step the probe is meant to test. However, the impact here appears to be nil in practice: (1) this sed is pre-existing launcher behavior, unmodified by this PR, and applies identically to every recipe that goes through this launcher; (2) per the PR description, this config is otherwise identical to the already-merged, fully-green Variant G (#2374), which ran successfully under the exact same 1440-in-recipe/720-effective override — proving the model load completes comfortably within 7200s and the extra 1440 margin was never actually exercised. So the recipe's 1440 is effectively a dead/misleading value the launcher always clobbers, rather than a value that changes observed behavior.
Fix: either make the launcher's sed only raise (not lower) the value, e.g. compare against the existing number and skip the rewrite if it's already >= 720, or drop the recipe's max_attempts override entirely (since it currently has no effect) and rely on the launcher's fixed 720, with a comment noting that. Given this is an explicitly not-to-merge probe PR whose entire purpose is a one-off empirical check, this is not blocking, but the author should reconcile the recipe's declared value with the launcher's clobbering behavior before this pattern is copied into a real (mergeable) recipe.