Repository navigation
[KimiK3][AgentX]: GB200 DSpark and Simple CPU KV offload #2404
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
Merged
Merged
Changes from 30 commits
Commits
Show all changes
69 commits
Select commit
Hold shift + click to select a range
245dcda
feat(agentx): add Kimi K3 GB200 day-0 support
cquil11 3e2b960
chore: link GB200 Kimi K3 changelog
cquil11 4aa6b8e
Merge origin/main into GB200 Kimi K3 AgentX
cquil11 18b1af1
fix(agentx): align Kimi K3 DEP metadata
cquil11 acfd9f3
perf(gb200): load Kimi K3 from local NVMe
cquil11 fd088cd
chore: simplify Kimi K3 performance config
cquil11 6c335b5
perf(gb200): size Kimi K3 AgentX sweep from HBM
cquil11 dc90e33
fix(gb200): drop unsupported Kimi K3 worker parsers
cquil11 d3b3f05
fix(gb200): let srt-slurm select collective interfaces
cquil11 f513599
fix(gb200): launch Kimi K3 TP2 DP8 per node
cquil11 bc21aac
fix(gb200): rebalance Kimi K3 EP16 for HBM
cquil11 819af41
fix(gb200): use native srt-slurm TPDP support
cquil11 458d6f6
docs(perf): correct Kimi K3 throughput topology
cquil11 6d10fbc
fix(gb200): use official Kimi K3 Dynamo runtime
cquil11 4158041
fix(gb200): align Kimi K3 runtime metadata
cquil11 a32ee0b
fix(gb200): reserve Kimi K3 prefill workspace
cquil11 86f4e70
fix(config): report Kimi K3 TP16 as EP16
cquil11 ed0f177
merge: sync main into Kimi K3 GB200 support
cquil11 1f052d2
fix(gb200): stage srt AgentX aggregates
cquil11 64de559
fix(agentx): burst Kimi K3 saturation starts
cquil11 ff5fb06
refactor(kimik3): depend on shared phase starts
cquil11 2ddb6d5
merge: sync main into Kimi K3 GB200 support
cquil11 fc99700
Merge origin/main into agent/kimik3-gb200-agentx
cquil11 235aa5e
fix(agentx): extend Kimi K3 DEP warmup drain
cquil11 26e6697
merge: sync main into Kimi K3 GB200 support
cquil11 403e0c1
merge: sync latest main into Kimi K3 GB200 support
cquil11 1187a30
chore(runners): register fourth GB200 worker
cquil11 275eef2
merge: sync latest main into Kimi K3 GB200 offload
cquil11 d29685c
feat(agentx): add GB200 Kimi K3 CPU KV offload
cquil11 5bb3034
merge: sync latest main into Kimi K3 GB200 offload
cquil11 4321a95
Merge branch 'main' into agent/kimik3-gb200-simple-offload
cquil11 a1e7500
feat(kimik3): use upstream vLLM with DSpark on GB200
cquil11 072a242
Merge branch 'main' into agent/kimik3-gb200-simple-offload
cquil11 8400af0
fix(kimik3): pin K3-compatible Dynamo preview
cquil11 dba74e0
Merge origin/main into agent/kimik3-gb200-simple-offload
cquil11 74b7aec
merge: sync latest main into Kimi K3 GB200 offload
cquil11 9eaa372
fix(kimik3): bridge DSpark mask metadata for Dynamo
cquil11 389141c
fix(kimik3): use model runner v2 for DSpark
cquil11 d4791fc
Merge origin/main into Kimi K3 GB200 AgentX
cquil11 07d0a1c
Merge origin/main into Kimi K3 GB200 AgentX
cquil11 35c272d
fix(kimik3): preserve full context on DEP
cquil11 e141992
fix(kimik3): reserve DEP graph memory
cquil11 c6841ed
fix(kimik3): use piecewise DEP graphs
cquil11 07775b0
merge: sync latest main
cquil11 608e9c0
fix(kimik3): stabilize DSpark DEP MoE
cquil11 d598886
fix(agentx): reduce Kimi K3 DEP load memory
cquil11 4595493
fix(agentx): right-size Kimi K3 DEP graphs
cquil11 dce6567
fix(agentx): report pinned Dynamo version
cquil11 4192722
fix(agentx): reduce offload graph residency
cquil11 2eeac85
chore: merge main into Kimi K3 GB200 branch
cquil11 f002be9
fix(kimik3): extend offload sweep allocation
cquil11 447fcae
perf(kimik3): report prompt cache reads
cquil11 54c8a7c
merge: sync latest main
cquil11 d238ef7
fix(agentx): align GB200 replay timing
cquil11 4491f28
fix(agentx): let Dynamo report cached prompt tokens
cquil11 43f2a8b
merge: sync current main into GB200 AgentX
cquil11 636897b
fix(agentx): align AIPerf with Qwen frontiers
cquil11 2ba037b
fix(agentx): separate live and final error gates
cquil11 76f6954
merge: sync latest main into GB200 AgentX
cquil11 c37ef59
fix(agentx): allow full-context sweeps to finish
cquil11 a972ff5
chore(agentx): bump AIPerf harness
cquil11 015e89d
perf(gb200): reuse Lustre-backed runner workspace
cquil11 ab9f53e
fix(gb200): isolate InferenceX Slurm jobs
cquil11 48cf6ea
fix(gb200): preserve namespaced Slurm job name
cquil11 1a4a4fb
chore(agentx): bump AIPerf timing harness
cquil11 ca38c5b
chore: merge origin/main into GB200 Kimi K3 branch
cquil11 2ad4ef3
fix(agentx): remove duplicate trace idle cap plumbing
cquil11 1c5681c
Merge branch 'main' into agent/kimik3-gb200-simple-offload
cquil11 947c9b3
Update perf-changelog.yaml
cquil11 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
132 changes: 132 additions & 0 deletions
132
...multi_node/srt-slurm-recipes/vllm/kimi-k3/agentic/agg-gb200-dep16-throughput-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,132 @@ | ||
| name: "kimi-k3-vllm-agg-gb200-dep16-throughput-agentic" | ||
|
|
||
| # Day-0 GB200 translation of the official throughput-oriented multi_node_dep | ||
| # profile. TP4 x DP4 gives EP16 across four four-GPU GB200 nodes. | ||
| # https://recipes.vllm.ai/moonshotai/Kimi-K3?hardware=gb200&nodes=4&strategy=multi_node_dep | ||
|
|
||
| model: | ||
| path: "kimi-k3" | ||
| container: "nvcr.io/nvidia/ai-dynamo/vllm-runtime:1.4.0-kimi-k3-dev.1" | ||
| precision: "fp4" | ||
|
|
||
| identity: | ||
| model: | ||
| repo: "moonshotai/Kimi-K3" | ||
| container: | ||
| image: "nvcr.io/nvidia/ai-dynamo/vllm-runtime:1.4.0-kimi-k3-dev.1" | ||
| frameworks: | ||
| dynamo: "1.3.0" | ||
|
|
||
| dynamo: | ||
| install: false | ||
|
|
||
| environment: | ||
| ETCD_LEASE_TTL: "7200" | ||
|
|
||
| slurm: | ||
| time_limit: "8:00:00" | ||
|
|
||
| health_check: | ||
| max_attempts: 2160 | ||
| interval_seconds: 10 | ||
|
|
||
| resources: | ||
| gpu_type: "gb200" | ||
| gpus_per_node: 4 | ||
| agg_nodes: 4 | ||
| 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 | ||
| args: | ||
| dyn-chat-processor: "dynamo" | ||
| router-mode: "kv" | ||
| router-kv-events: true | ||
| router-temperature: "0" | ||
| router-min-initial-workers: 1 | ||
| kv-cache-block-size: 64 | ||
|
|
||
| backend: | ||
| type: vllm | ||
| connector: null | ||
| aggregated_environment: | ||
| HF_HUB_CACHE: "/hf_hub_cache" | ||
| HUGGINGFACE_HUB_CACHE: "/hf_hub_cache" | ||
| TRANSFORMERS_CACHE: "/hf_hub_cache" | ||
| VLLM_ENGINE_READY_TIMEOUT_S: "7200" | ||
| VLLM_RPC_TIMEOUT: "600000" | ||
| VLLM_LOG_STATS_INTERVAL: "1" | ||
| VLLM_USE_V2_MODEL_RUNNER: "0" | ||
| VLLM_ENABLE_K3_LATENT_MOE_TAIL_FUSION: "1" | ||
| VLLM_ALLREDUCE_USE_FLASHINFER: "1" | ||
| VLLM_FLASHINFER_ALLREDUCE_BACKEND: "mnnvl" | ||
| NCCL_CUMEM_ENABLE: "1" | ||
| NCCL_MNNVL_ENABLE: "1" | ||
| NCCL_NVLS_ENABLE: "1" | ||
| UCX_MEMTYPE_CACHE: "n" | ||
| UCX_NET_DEVICES: "mlx5_0:1,mlx5_1:1,mlx5_2:1,mlx5_3:1" | ||
| UCX_TLS: "rc,cuda_copy" | ||
| NCCL_IB_HCA: "mlx5_0,mlx5_1,mlx5_2,mlx5_3" | ||
| NCCL_P2P_LEVEL: "NVL" | ||
| NVIDIA_GDRCOPY: "1" | ||
| PYTORCH_ALLOC_CONF: "expandable_segments:True" | ||
| DG_JIT_CACHE_DIR: "/tmp/dg-cache-kimi-k3-gb200-dep16-{job_id}" | ||
| vllm_config: | ||
| aggregated: | ||
| served-model-name: "moonshotai/Kimi-K3" | ||
| tensor-parallel-size: 4 | ||
| pipeline-parallel-size: 1 | ||
| data-parallel-size: 4 | ||
| data-parallel-rpc-port: 13345 | ||
| enable-expert-parallel: true | ||
| trust-remote-code: true | ||
| load-format: "fastsafetensors" | ||
| safetensors-load-strategy: "lazy" | ||
| kv-cache-dtype: "fp8" | ||
| attention-backend: "FLASHINFER_MLA" | ||
| attention-config: '{"mla_prefill_backend":"TRTLLM_RAGGED","use_prefill_query_quantization":true}' | ||
| moe-backend: "deep_gemm_mega_moe" | ||
| kda-prefill-backend: "flashkda" | ||
| kernel-config: '{"enable_cutedsl_warmup":true}' | ||
| all2all-backend: "flashinfer_nvlink_one_sided" | ||
| gpu-memory-utilization: 0.92 | ||
| max-num-seqs: 128 | ||
| max-num-batched-tokens: 16384 | ||
| block-size: 64 | ||
| language-model-only: true | ||
| disable-custom-all-reduce: true | ||
| enable-prefix-caching: true | ||
| scheduler-cls: "vllm.v1.core.sched.async_scheduler.AsyncScheduler" | ||
| dyn-tool-call-parser: "kimi_k3" | ||
| reasoning-parser: "kimi_k3" | ||
| dyn-reasoning-parser: "kimi_k3" | ||
| no-enable-flashinfer-autotune: true | ||
|
|
||
| sbatch_directives: | ||
| cpus-per-task: "144" | ||
| mem: "0" | ||
|
|
||
| srun_options: | ||
| container-remap-root: "" | ||
|
|
||
| benchmark: | ||
| type: custom | ||
| aiperf_server_metrics: true | ||
| command: bash /infmax-workspace/benchmarks/multi_node/agentic_srt.sh | ||
| env: | ||
| INFMAX_CONTAINER_WORKSPACE: "/infmax-workspace" | ||
| RESULT_DIR: "/logs/agentic" | ||
| PORT: "8000" | ||
| IS_MULTINODE: "true" | ||
| AIPERF_USE_DYNAMO_CONV_AWARE_ROUTING: "0" | ||
| AIPERF_DYNAMO_SESSION_TIMEOUT_SECONDS: "14400" | ||
| AGENTIC_WARMUP_GRACE_PERIOD: "3600" | ||
| AIPERF_DATASET_MMAP_CACHE_DIR: "/aiperf_mmap_cache" | ||
| HF_HUB_CACHE: "/hf_hub_cache" | ||
| WEKA_LOADER_OVERRIDE: "semianalysis_cc_traces_weka_062126" | ||
135 changes: 135 additions & 0 deletions
135
...-recipes/vllm/kimi-k3/agentic/agg-gb200-dep16-throughput-vllm-simple-offload-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,135 @@ | ||
| name: "kimi-k3-vllm-agg-gb200-dep16-throughput-vllm-simple-offload-agentic" | ||
|
|
||
| # High-concurrency host-DRAM KV-offload variant of the official throughput- | ||
| # oriented multi_node_dep profile. TP4 x DP4 gives EP16 across four four-GPU | ||
| # GB200 nodes. Each TP rank receives a 128 GiB CPU KV pool (512 GiB per node). | ||
| # https://recipes.vllm.ai/moonshotai/Kimi-K3?hardware=gb200&nodes=4&strategy=multi_node_dep | ||
|
|
||
| model: | ||
| path: "kimi-k3" | ||
| container: "nvcr.io/nvidia/ai-dynamo/vllm-runtime:1.4.0-kimi-k3-dev.1" | ||
| precision: "fp4" | ||
|
|
||
| identity: | ||
| model: | ||
| repo: "moonshotai/Kimi-K3" | ||
| container: | ||
| image: "nvcr.io/nvidia/ai-dynamo/vllm-runtime:1.4.0-kimi-k3-dev.1" | ||
| frameworks: | ||
| dynamo: "1.3.0" | ||
|
|
||
| dynamo: | ||
| install: false | ||
|
|
||
| environment: | ||
| ETCD_LEASE_TTL: "7200" | ||
|
|
||
| slurm: | ||
| time_limit: "8:00:00" | ||
|
|
||
| health_check: | ||
| max_attempts: 2160 | ||
| interval_seconds: 10 | ||
|
|
||
| resources: | ||
| gpu_type: "gb200" | ||
| gpus_per_node: 4 | ||
| agg_nodes: 4 | ||
| 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 | ||
| args: | ||
| dyn-chat-processor: "dynamo" | ||
| router-mode: "kv" | ||
| router-kv-events: true | ||
| router-temperature: "0" | ||
| router-min-initial-workers: 1 | ||
| kv-cache-block-size: 64 | ||
|
|
||
| backend: | ||
| type: vllm | ||
| connector: null | ||
| aggregated_environment: | ||
| HF_HUB_CACHE: "/hf_hub_cache" | ||
| HUGGINGFACE_HUB_CACHE: "/hf_hub_cache" | ||
| TRANSFORMERS_CACHE: "/hf_hub_cache" | ||
| VLLM_ENGINE_READY_TIMEOUT_S: "7200" | ||
| VLLM_RPC_TIMEOUT: "600000" | ||
| VLLM_LOG_STATS_INTERVAL: "1" | ||
| VLLM_USE_V2_MODEL_RUNNER: "0" | ||
| VLLM_ENABLE_K3_LATENT_MOE_TAIL_FUSION: "1" | ||
| VLLM_ALLREDUCE_USE_FLASHINFER: "1" | ||
| VLLM_FLASHINFER_ALLREDUCE_BACKEND: "mnnvl" | ||
| NCCL_CUMEM_ENABLE: "1" | ||
| NCCL_MNNVL_ENABLE: "1" | ||
| NCCL_NVLS_ENABLE: "1" | ||
| UCX_MEMTYPE_CACHE: "n" | ||
| UCX_NET_DEVICES: "mlx5_0:1,mlx5_1:1,mlx5_2:1,mlx5_3:1" | ||
| UCX_TLS: "rc,cuda_copy" | ||
| NCCL_IB_HCA: "mlx5_0,mlx5_1,mlx5_2,mlx5_3" | ||
| NCCL_P2P_LEVEL: "NVL" | ||
| NVIDIA_GDRCOPY: "1" | ||
| PYTHONHASHSEED: "42" | ||
| PYTORCH_ALLOC_CONF: "expandable_segments:True" | ||
| DG_JIT_CACHE_DIR: "/tmp/dg-cache-kimi-k3-gb200-dep16-offload-{job_id}" | ||
| vllm_config: | ||
| aggregated: | ||
| served-model-name: "moonshotai/Kimi-K3" | ||
| tensor-parallel-size: 4 | ||
| pipeline-parallel-size: 1 | ||
| data-parallel-size: 4 | ||
| data-parallel-rpc-port: 13345 | ||
| enable-expert-parallel: true | ||
| trust-remote-code: true | ||
| load-format: "fastsafetensors" | ||
| safetensors-load-strategy: "lazy" | ||
| kv-cache-dtype: "fp8" | ||
| attention-backend: "FLASHINFER_MLA" | ||
| attention-config: '{"mla_prefill_backend":"TRTLLM_RAGGED","use_prefill_query_quantization":true}' | ||
| moe-backend: "deep_gemm_mega_moe" | ||
| kda-prefill-backend: "flashkda" | ||
| kernel-config: '{"enable_cutedsl_warmup":true}' | ||
| all2all-backend: "flashinfer_nvlink_one_sided" | ||
| gpu-memory-utilization: 0.92 | ||
| max-num-seqs: 128 | ||
| max-num-batched-tokens: 16384 | ||
| block-size: 64 | ||
| language-model-only: true | ||
| disable-custom-all-reduce: true | ||
| enable-prefix-caching: true | ||
| kv-transfer-config: '{"kv_connector":"SimpleCPUOffloadConnector","kv_role":"kv_both","kv_connector_extra_config":{"cpu_bytes_to_use":549755813888,"cpu_bytes_to_use_per_rank":137438953472,"lazy_offload":false}}' | ||
| scheduler-cls: "vllm.v1.core.sched.async_scheduler.AsyncScheduler" | ||
| dyn-tool-call-parser: "kimi_k3" | ||
| reasoning-parser: "kimi_k3" | ||
| dyn-reasoning-parser: "kimi_k3" | ||
| no-enable-flashinfer-autotune: true | ||
|
|
||
| sbatch_directives: | ||
| cpus-per-task: "144" | ||
| mem: "0" | ||
|
|
||
| srun_options: | ||
| container-remap-root: "" | ||
|
|
||
| benchmark: | ||
| type: custom | ||
| aiperf_server_metrics: true | ||
| command: bash /infmax-workspace/benchmarks/multi_node/agentic_srt.sh | ||
| env: | ||
| INFMAX_CONTAINER_WORKSPACE: "/infmax-workspace" | ||
| RESULT_DIR: "/logs/agentic" | ||
| PORT: "8000" | ||
| IS_MULTINODE: "true" | ||
| AIPERF_USE_DYNAMO_CONV_AWARE_ROUTING: "0" | ||
| AIPERF_DYNAMO_SESSION_TIMEOUT_SECONDS: "14400" | ||
| AGENTIC_WARMUP_GRACE_PERIOD: "3600" | ||
| AIPERF_DATASET_MMAP_CACHE_DIR: "/aiperf_mmap_cache" | ||
| HF_HUB_CACHE: "/hf_hub_cache" | ||
| WEKA_LOADER_OVERRIDE: "semianalysis_cc_traces_weka_062126" |
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.
do u have vllm upstream image
can u add dspark, there is non-mtp for agentx