diff --git a/inferencex-e2e/benchmarks/single_node/srt-slurm-recipes/kimik3/atom/mi355x-fp4-mtp/agentic.yaml b/inferencex-e2e/benchmarks/single_node/srt-slurm-recipes/kimik3/atom/mi355x-fp4-mtp/agentic.yaml index 4646809ea3..6f40b5636a 100644 --- a/inferencex-e2e/benchmarks/single_node/srt-slurm-recipes/kimik3/atom/mi355x-fp4-mtp/agentic.yaml +++ b/inferencex-e2e/benchmarks/single_node/srt-slurm-recipes/kimik3/atom/mi355x-fp4-mtp/agentic.yaml @@ -1,14 +1,16 @@ -# Kimi-K3 MXFP4 AgentX on MI355X with ATOM DSpark: the interactive band, TP8 -# with GPU-resident KV and the deepest published draft (seven tokens). The -# 1.56 TB checkpoint only fits at TP8. The DCP8 LMCache bands were removed with -# the legacy script (#3461): srtctl reserves ATOM's kv-transfer-config for -# disaggregated workers. +# Kimi-K3 MXFP4 AgentX on MI355X with ATOM DSpark, tracking ROCm/ATOM +# recipes/Agentic-Kimi-K3.md. TP8 only: the 1.56 TB checkpoint is ~195 GB/GPU. +# interactive (1, 4) DCP1, DSpark 7, GPU-resident KV. +# mid (14, 16) DCP8, DSpark 3, ReplaySSM, LMCache 128 GB/rank. +# throughput (48, 56, 72) DCP8, no draft, LMCache 128 or 192 GB/rank. +# LMCache is ATOM's in-process lmcache_offload connector +# (runners/srt-slurm/patches/507-lmcache-server-atom-sglang.patch). base: schema: 2 name: kimik3-fp4-mi355x-atom-agentic model: path: hf:moonshotai/Kimi-K3 - container: rocm/atom-dev:ubuntu24.04_py3.12_pytorch_release_2.10.0_kimi_k3_agentic_0911 + container: rocm/atom-dev:nightly_202609251613 precision: fp4 resources: gpu_type: mi355x @@ -68,6 +70,8 @@ base: # Anchor-only state checkpointing. ATOM_STATE_CHECKPOINT_DEMAND: '0' ATOM_GDN_SSM_DTYPE: fp16 + # FlyDSL FP8 prefill attention; ATOM defaults it off. + ATOM_USE_FLYDSL_FP8_PREFILL_ATTN: '1' benchmark: type: custom command: bash /infmax-workspace/benchmarks/srt_agentic.sh @@ -85,7 +89,7 @@ base: AIPERF_APPLY_CHAT_TEMPLATE: 'true' # One variant per point. Full graphs cover every verify batch up to 2x CONC -# times the eight-token verify window. +# times the verify window (1 + draft tokens). override_tp8_c1: roles: agg: @@ -103,3 +107,136 @@ override_tp8_c4: benchmark: env: CONC: '4' + +# Mid band: DCP8 shards the KV read across all eight GPUs, ReplaySSM rebuilds +# KDA state from the checkpoint ring, and the paged KV rides the LMCache CPU +# tier. From concurrency 16 up, a ready prefill waits four decode passes. +override_tp8_dcp8_c14_lmcache: + roles: + agg: + args: + decode-context-parallel-size: 8 + num-speculative-tokens: 3 + cudagraph-capture-sizes: '[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50,51,52,53,54,55,56,57,58,59,60,61,62,63,64,65,66,67,68,69,70,71,72,73,74,75,76,77,78,79,80,81,82,83,84,85,86,87,88,89,90,91,92,93,94,95,96,97,98,99,100,101,102,103,104,105,106,107,108,109,110,111,112]' + extra-kv-connectors: &lmcache_offload + - kv_connector: lmcache_offload + kv_role: offload + env: + <<: &lmcache_env + PYTHONHASHSEED: '0' + LMCACHE_LOCAL_CPU: 'True' + # The offload hash block is block-size 128 x DCP 8, so the KV and the + # state-checkpoint grids coincide. + LMCACHE_CHUNK_SIZE: '1024' + # K3 is a hybrid; without this the paged KV never reaches the CPU tier. + OFFLOAD_KV_FOR_HYBRID: '1' + # Pin the eight ranks to the two sockets explicitly. + LMCACHE_NUMA_MODE: auto + ATOM_NUMA_BIND: '1' + ATOM_NUMA_NODE: 0,0,0,0,1,1,1,1 + ATOM_AUTO_NUMA_BIND: '0' + OFFLOAD_PROFILE: '1' + # One K3 state entry is 54.78 MiB; the default 8 MiB staging buffer + # cannot hold it. + OFFLOAD_GPU_STAGING_CHUNKS: '32' + LMCACHE_MAX_LOCAL_CPU_SIZE: '128' + ATOM_ENABLE_REPLAYSSM: '1' + benchmark: + env: + CONC: '14' + KV_OFFLOADING: dram + TOTAL_CPU_DRAM_GB: '1028' + +override_tp8_dcp8_c16_lmcache: + roles: + agg: + args: + decode-context-parallel-size: 8 + num-speculative-tokens: 3 + cudagraph-capture-sizes: '[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50,51,52,53,54,55,56,57,58,59,60,61,62,63,64,65,66,67,68,69,70,71,72,73,74,75,76,77,78,79,80,81,82,83,84,85,86,87,88,89,90,91,92,93,94,95,96,97,98,99,100,101,102,103,104,105,106,107,108,109,110,111,112,113,114,115,116,117,118,119,120,121,122,123,124,125,126,127,128]' + extra-kv-connectors: *lmcache_offload + env: + <<: *lmcache_env + LMCACHE_MAX_LOCAL_CPU_SIZE: '128' + ATOM_ENABLE_REPLAYSSM: '1' + ATOM_PREFILL_DECODE_INTERVAL: '4' + ATOM_PREFILL_DELAYER_MAX_QUEUE_MS: '5000' + benchmark: + env: + CONC: '16' + KV_OFFLOADING: dram + TOTAL_CPU_DRAM_GB: '1028' + +# Throughput band: past the knee the draft no longer pays for itself, and +# max-num-seqs tracks 2x CONC. 56 and 72 need the 192 GB/rank pool and reuse +# identical AITER communicator groups. +override_tp8_dcp8_c48_lmcache: + roles: + agg: + args: + decode-context-parallel-size: 8 + method: null + draft-model: null + num-speculative-tokens: null + max-num-seqs: 96 + cudagraph-capture-sizes: '[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50,51,52,53,54,55,56,57,58,59,60,61,62,63,64,65,66,67,68,69,70,71,72,73,74,75,76,77,78,79,80,81,82,83,84,85,86,87,88,89,90,91,92,93,94,95,96]' + extra-kv-connectors: *lmcache_offload + env: + <<: *lmcache_env + LMCACHE_MAX_LOCAL_CPU_SIZE: '128' + ATOM_PREFILL_DECODE_INTERVAL: '4' + ATOM_PREFILL_DELAYER_MAX_QUEUE_MS: '5000' + benchmark: + env: + CONC: '48' + SPEC_DECODING: mtp + KV_OFFLOADING: dram + TOTAL_CPU_DRAM_GB: '1028' + +override_tp8_dcp8_c56_lmcache: + roles: + agg: + args: + decode-context-parallel-size: 8 + method: null + draft-model: null + num-speculative-tokens: null + max-num-seqs: 112 + cudagraph-capture-sizes: '[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50,51,52,53,54,55,56,57,58,59,60,61,62,63,64,65,66,67,68,69,70,71,72,73,74,75,76,77,78,79,80,81,82,83,84,85,86,87,88,89,90,91,92,93,94,95,96,97,98,99,100,101,102,103,104,105,106,107,108,109,110,111,112]' + extra-kv-connectors: *lmcache_offload + env: + <<: *lmcache_env + LMCACHE_MAX_LOCAL_CPU_SIZE: '192' + AITER_REUSE_IDENTICAL_COMM_GROUPS: '1' + ATOM_PREFILL_DECODE_INTERVAL: '4' + ATOM_PREFILL_DELAYER_MAX_QUEUE_MS: '5000' + benchmark: + env: + CONC: '56' + SPEC_DECODING: mtp + KV_OFFLOADING: dram + TOTAL_CPU_DRAM_GB: '1538' + +override_tp8_dcp8_c72_lmcache: + roles: + agg: + args: + decode-context-parallel-size: 8 + method: null + draft-model: null + num-speculative-tokens: null + max-num-seqs: 144 + cudagraph-capture-sizes: '[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50,51,52,53,54,55,56,57,58,59,60,61,62,63,64,65,66,67,68,69,70,71,72,73,74,75,76,77,78,79,80,81,82,83,84,85,86,87,88,89,90,91,92,93,94,95,96,97,98,99,100,101,102,103,104,105,106,107,108,109,110,111,112,113,114,115,116,117,118,119,120,121,122,123,124,125,126,127,128,129,130,131,132,133,134,135,136,137,138,139,140,141,142,143,144]' + extra-kv-connectors: *lmcache_offload + env: + <<: *lmcache_env + LMCACHE_MAX_LOCAL_CPU_SIZE: '192' + AITER_REUSE_IDENTICAL_COMM_GROUPS: '1' + ATOM_PREFILL_DECODE_INTERVAL: '4' + ATOM_PREFILL_DELAYER_MAX_QUEUE_MS: '5000' + benchmark: + env: + CONC: '72' + SPEC_DECODING: mtp + KV_OFFLOADING: dram + TOTAL_CPU_DRAM_GB: '1538' diff --git a/inferencex-e2e/configs/amd-master.yaml b/inferencex-e2e/configs/amd-master.yaml index 8d12f99c3c..b12197a349 100644 --- a/inferencex-e2e/configs/amd-master.yaml +++ b/inferencex-e2e/configs/amd-master.yaml @@ -674,7 +674,7 @@ dsr1-fp8-mi355x-sglang-disagg-mtp: # wider in-flight window needs the deeper pool to keep the paged KV # resident. kimik3-fp4-mi355x-atom-agentic-mtp: - image: rocm/atom-dev:ubuntu24.04_py3.12_pytorch_release_2.10.0_kimi_k3_agentic_0911 + image: rocm/atom-dev:nightly_202609251613 model: moonshotai/Kimi-K3 model-prefix: kimik3 runner: cluster:mi355x-amds @@ -686,6 +686,10 @@ kimik3-fp4-mi355x-atom-agentic-mtp: - dram-utilization: 0.343 search-space: - { tp: 8, kv-offloading: none, conc-list: [1, 4], spec-decoding: mtp, srt-recipe: benchmarks/single_node/srt-slurm-recipes/kimik3/atom/mi355x-fp4-mtp/agentic.yaml } + - { tp: 8, dcp-size: 8, kv-offloading: dram, kv-offload-backend: { name: lmcache, version: "0.4.5" }, conc-list: [14, 16, 48], spec-decoding: mtp, srt-recipe: benchmarks/single_node/srt-slurm-recipes/kimik3/atom/mi355x-fp4-mtp/agentic.yaml } + - dram-utilization: 0.513 + search-space: + - { tp: 8, dcp-size: 8, kv-offloading: dram, kv-offload-backend: { name: lmcache, version: "0.4.5" }, conc-list: [56, 72], spec-decoding: mtp, srt-recipe: benchmarks/single_node/srt-slurm-recipes/kimik3/atom/mi355x-fp4-mtp/agentic.yaml } minimaxm3-fp4-mi355x-atom-agentic-mtp: image: rocm/atom-dev:nightly_202609171455 diff --git a/inferencex-e2e/infx/srt_slurm/single_node.py b/inferencex-e2e/infx/srt_slurm/single_node.py index 1e6baece4b..0c758d1f45 100644 --- a/inferencex-e2e/infx/srt_slurm/single_node.py +++ b/inferencex-e2e/infx/srt_slurm/single_node.py @@ -127,8 +127,8 @@ def validate_recipe(recipe: dict[str, Any], environment: Mapping[str, str]) -> N if engine == "atom": # Native ATOM derives -tp from the aggregate worker's GPU allocation. expected["ATOM TP"] = (role["gpus"], int(environment["TP"])) - # vLLM shards decode KV across its tensor-parallel ranks. - dcp = str(args.get("decode-context-parallel-size", 1)) if engine == "vllm" else "1" + # vLLM and ATOM shard decode KV across their tensor-parallel ranks. + dcp = str(args.get("decode-context-parallel-size", 1)) if engine in {"vllm", "atom"} else "1" for name, value in {"PP_SIZE": "1", "DCP_SIZE": dcp, "PCP_SIZE": "1"}.items(): expected[name] = (environment[name], value) for name, (actual, wanted) in expected.items(): diff --git a/inferencex-e2e/infx/tests/srt_slurm/test_srt_single_node.py b/inferencex-e2e/infx/tests/srt_slurm/test_srt_single_node.py index 09e7b4590a..306aea6b6f 100644 --- a/inferencex-e2e/infx/tests/srt_slurm/test_srt_single_node.py +++ b/inferencex-e2e/infx/tests/srt_slurm/test_srt_single_node.py @@ -232,9 +232,15 @@ def test_atom_binding_uses_allocation_tp_and_native_mtp_arguments(point): ({"EP_SIZE": "1"}, "enable-expert-parallel"), ({"TP": "8", "EP_SIZE": "8"}, "ATOM TP"), ({"DP_ATTENTION": "false"}, "DP_ATTENTION"), + ({"DCP_SIZE": "8"}, "DCP_SIZE"), ]: with pytest.raises(ValueError, match=error): runtime_arguments(f"{path}:base", {**env, **changes}) + recipe["roles"]["agg"]["args"]["decode-context-parallel-size"] = 8 + path.write_text(yaml.safe_dump({"base": recipe})) + runtime_arguments(f"{path}:base", {**env, "DCP_SIZE": "8"}) + with pytest.raises(ValueError, match="DCP_SIZE"): + runtime_arguments(f"{path}:base", env) @pytest.mark.parametrize("record,expected", [ diff --git a/inferencex-e2e/perf-changelog.yaml b/inferencex-e2e/perf-changelog.yaml index 9fa528a7cb..cdd9f07f34 100644 --- a/inferencex-e2e/perf-changelog.yaml +++ b/inferencex-e2e/perf-changelog.yaml @@ -9092,3 +9092,15 @@ - "在 Dynamo 前端启用 MiniMax-M3 推理解析器，使推理内容按一致方式分类，而不会作为普通内容返回。" - "在每个 TensorRT-LLM 引擎配置中设置 reasoning_parser: minimax_m3，使前端与引擎的解析器设置保持一致。" pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/3195 + +- config-keys: + - kimik3-fp4-mi355x-atom-agentic-mtp + scenario-type: + - agentic-coding + description: + - "Move the MI355X Kimi-K3 FP4 ATOM AgentX submission to image rocm/atom-dev:nightly_202609251613, tracking recipes/Agentic-Kimi-K3.md as retuned in ROCm/ATOM#2382." + - "Enable FlyDSL FP8 prefill attention (ATOM_USE_FLYDSL_FP8_PREFILL_ATTN=1) on every concurrency; ATOM defaults it off." + - "Hold a ready prefill for four decode passes from concurrency 16 up (ATOM_PREFILL_DECODE_INTERVAL=4, ATOM_PREFILL_DELAYER_MAX_QUEUE_MS=5000). Concurrency 1, 4 and 14 are unchanged." + - "Restore the DCP8 LMCache bands on the native srt-slurm recipe: concurrency 14 and 16 (DSpark 3, ReplaySSM) and 48, 56 and 72 (no draft) run ATOM's in-process lmcache_offload connector through roles.agg.args.extra-kv-connectors (srt-slurm patch 507), with 128 GB/rank up to 48 and 192 GB/rank at 56 and 72. Concurrency 1 and 4 stay GPU-resident." + - "No change to the Inferact/Kimi-K3-DSpark draft's precision: online_quant_config still excludes every draft linear (layers.*, context_proj), so its weights and activations stay BF16, and it keeps the target's FP8 KV cache (kv_cache_dtype fp8). FlyDSL FP8 prefill attention applies only to the target, since the draft runs its block pass as decode attention." + pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/3407