From 055dfc2ddcd77841d127ef29e77adb04c7802ddd Mon Sep 17 00:00:00 2001 From: Samuel Shen Date: Wed, 12 Aug 2026 19:58:31 +0000 Subject: [PATCH 01/31] kimik3-fp4-mi355x-vllm-agentic-mtp: add LMCache DRAM KV-offload arm Add an lmcache kv-offload-backend point at TP8 conc 10 on top of the existing DSpark MTP serving stack, mirroring the vllm-simple offload arm for a direct backend comparison. The benchmark script gains an lmcache case arm that installs the LMCache 0.5.4rc1 ROCm wheel (torch/ROCm stack untouched), starts one MP server per the Kimi-K3 recipe (chunk size 768 = K3 unified block size at 8 GPUs, --separate-object-groups for the hybrid KDA/MLA two-group KV layout, --enable-extra-logging, --max-cpu-workers 8 --max-gpu-workers 1), and wires vLLM to it via LMCacheMPConnector. --- .../agentic/kimik3_fp4_mi355x_mtp.sh | 75 +++++++++++++++++++ configs/amd-master.yaml | 4 + 2 files changed, 79 insertions(+) diff --git a/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh b/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh index ec2eaf3c3..0f1355dfb 100644 --- a/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh +++ b/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh @@ -109,12 +109,14 @@ SERVER_LOG="$RESULT_DIR/server.log" mkdir -p "$RESULT_DIR" SERVER_PID="" +LMCACHE_PID="" cleanup_agentic_services() { local exit_code=$? trap - EXIT INT TERM set +e stop_background_process_tree "$SERVER_PID" "vLLM server" 60 + stop_background_process_tree "$LMCACHE_PID" "LMCache server" exit "$exit_code" } trap cleanup_agentic_services EXIT @@ -143,6 +145,79 @@ case "${KV_OFFLOAD_BACKEND:-}" in ) echo "SimpleCPUOffloadConnector: ${CPU_BYTES_PER_RANK} B/rank x ${TP} ranks, lazy_offload=$SIMPLE_LAZY_OFFLOAD" ;; + lmcache) + require_agentic_kv_offload_backend "$KV_OFFLOAD_BACKEND" + + # Keep the image's tested torch/ROCm stack and install only LMCache's + # missing runtime dependencies, same as the MiniMax-M3 lmcache arm. + LMCACHE_VERSION="0.5.4rc1" + LMCACHE_ROCM_INDEX="https://github.com/LMCache/LMCache/releases/expanded_assets/v${LMCACHE_VERSION}-rocm" + agentic_pip_install --quiet --no-cache-dir --no-deps \ + "sortedcontainers==2.4.0" \ + "opentelemetry-exporter-prometheus==0.61b0" \ + "cupy-rocm-7-0==14.1.1" \ + "lmcache==${LMCACHE_VERSION}" --find-links "$LMCACHE_ROCM_INDEX" + python3 -c \ + "import cupy; import lmcache.integration.vllm.lmcache_mp_connector; import opentelemetry.exporter.prometheus" \ + >/dev/null + + # One MP server for the node, per the Kimi-K3 recipe + # (docs.lmcache.ai/recipes/kimi_k3.html). --chunk-size must equal the + # K3 unified block size N=768 at 8 GPUs, and the hybrid KDA/MLA layout + # (two KV-cache groups under MTP) requires one object group per + # sliding-window size: --separate-object-groups. + LMCACHE_PORT=6555 + LMCACHE_HTTP_PORT=8090 + LMCACHE_LOG="$RESULT_DIR/lmcache_server.log" + + # The whole generated node-DRAM budget backs the single server's L1, + # which lives in /dev/shm; fail early if it cannot fit. + LMCACHE_L1_SIZE_GB="$TOTAL_CPU_DRAM_GB" + SHM_FREE_GB=$(df -BG --output=avail /dev/shm 2>/dev/null | tail -1 | tr -dc '0-9') + if [ -n "$SHM_FREE_GB" ] && [ "$SHM_FREE_GB" -gt 0 ]; then + SHM_CAP_GB=$((SHM_FREE_GB * 90 / 100)) + if [ "$LMCACHE_L1_SIZE_GB" -gt "$SHM_CAP_GB" ]; then + echo "Error: LMCache L1 ${LMCACHE_L1_SIZE_GB} GB exceeds 90% of free /dev/shm (${SHM_CAP_GB} GB)." >&2 + exit 1 + fi + fi + + LMCACHE_CMD=( + lmcache server + --host 127.0.0.1 + --port "$LMCACHE_PORT" + --http-host 127.0.0.1 + --http-port "$LMCACHE_HTTP_PORT" + --l1-size-gb "$LMCACHE_L1_SIZE_GB" + --l1-init-size-gb 10 + --chunk-size 768 + --separate-object-groups + --enable-extra-logging + --max-cpu-workers 8 + --max-gpu-workers 1 + --eviction-policy LRU + ) + append_command "$RESULT_DIR/lmcache_command.txt" "${LMCACHE_CMD[@]}" + "${LMCACHE_CMD[@]}" > "$LMCACHE_LOG" 2>&1 & + LMCACHE_PID=$! + wait_for_ready \ + --endpoint "http://127.0.0.1:${LMCACHE_HTTP_PORT}/healthcheck" \ + --log "$LMCACHE_LOG" \ + --pid "$LMCACHE_PID" \ + --sleep-interval 1 \ + --timeout 600 + + # 100k-330k-token agentic prefixes make single retrieves large; use the + # same MQ timeout headroom as the MiniMax-M3 arm. + OFFLOAD_ARGS=( + --kv-transfer-config + "{\"kv_connector\":\"LMCacheMPConnector\",\"kv_connector_module_path\":\"lmcache.integration.vllm.lmcache_mp_connector\",\"kv_role\":\"kv_both\",\"kv_connector_extra_config\":{\"lmcache.mp.port\":$LMCACHE_PORT,\"lmcache.mp.mq_timeout\":6000.0}}" + ) + ;; + *) + echo "Error: unsupported KV_OFFLOAD_BACKEND='$KV_OFFLOAD_BACKEND' (expected vllm-simple or lmcache)" >&2 + exit 1 + ;; esac fi diff --git a/configs/amd-master.yaml b/configs/amd-master.yaml index c94029484..3aa757c0a 100644 --- a/configs/amd-master.yaml +++ b/configs/amd-master.yaml @@ -650,6 +650,10 @@ kimik3-fp4-mi355x-vllm-agentic-mtp: search-space: - { tp: 8, kv-offloading: none, conc-list: [1, 4, 8] , spec-decoding: mtp} - { tp: 8, ep: 1, kv-offloading: dram, kv-offload-backend: { name: vllm-simple }, conc-list: [10], spec-decoding: mtp } + # LMCache MP-server DRAM offload on top of the same DSpark MTP serving + # stack, at the same concurrency as the vllm-simple arm for a direct + # offload-backend comparison. + - { tp: 8, ep: 1, kv-offloading: dram, kv-offload-backend: { name: lmcache, version: "0.5.4rc1" }, conc-list: [10], spec-decoding: mtp } dsr1-fp4-mi355x-sglang-disagg: image: lmsysorg/sglang-rocm:v0.5.12-rocm720-mi35x-20260519 From 328836b6f1aea4da24ac60651cdf4b7614920a25 Mon Sep 17 00:00:00 2001 From: Samuel Shen Date: Wed, 12 Aug 2026 19:59:05 +0000 Subject: [PATCH 02/31] perf-changelog: select kimik3-fp4-mi355x-vllm-agentic-mtp LMCache arm --- perf-changelog.yaml | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/perf-changelog.yaml b/perf-changelog.yaml index f118b99aa..da67c0138 100644 --- a/perf-changelog.yaml +++ b/perf-changelog.yaml @@ -5909,3 +5909,12 @@ - "Increase MTP speculative steps from 3 to 5 (num-draft-tokens 4→6) and raise SGLANG_SIMULATE_ACC_LEN from 2.99 to 3.61 to reflect higher measured acceptance rate, targeting ~15% throughput improvement" - "Replace the TP8/EP8 low-concurrency arm (conc [1,2,4]) hicache offload with kv-offloading: none to reduce per-request latency at low load; conc [1,2,4] are now tested on both tp=4+hicache and tp=8+no-offload so SemiAnalysis can select the Pareto-optimal point per concurrency" pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2570 + +- config-keys: + - kimik3-fp4-mi355x-vllm-agentic-mtp + scenario-type: + - agentic-coding + description: + - "Add an LMCache 0.5.4rc1 DRAM KV-offload arm at TP8 conc 10 on top of the DSpark MTP stack, mirroring the vllm-simple offload point for a direct backend comparison." + - "Run one LMCache MP server per node with chunk size 768 (the K3 unified block size at 8 GPUs) and --separate-object-groups for the hybrid KDA/MLA two-group KV layout." + pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2583 From 4584d37ab75aa38fd8af0060b3d6a429a6b82154 Mon Sep 17 00:00:00 2001 From: Samuel Shen Date: Wed, 12 Aug 2026 21:20:39 +0000 Subject: [PATCH 03/31] Move the LMCache arm to a dedicated config key at conc 4/8/16 A separate kimik3-fp4-mi355x-vllm-agentic-mtp-lmcache key lets the changelog select only the LMCache points instead of re-running the resident and vllm-simple arms of the base key. The base key returns to its upstream shape. --- configs/amd-master.yaml | 22 ++++++++++++++++++---- perf-changelog.yaml | 4 ++-- 2 files changed, 20 insertions(+), 6 deletions(-) diff --git a/configs/amd-master.yaml b/configs/amd-master.yaml index 3aa757c0a..5afd30620 100644 --- a/configs/amd-master.yaml +++ b/configs/amd-master.yaml @@ -650,10 +650,24 @@ kimik3-fp4-mi355x-vllm-agentic-mtp: search-space: - { tp: 8, kv-offloading: none, conc-list: [1, 4, 8] , spec-decoding: mtp} - { tp: 8, ep: 1, kv-offloading: dram, kv-offload-backend: { name: vllm-simple }, conc-list: [10], spec-decoding: mtp } - # LMCache MP-server DRAM offload on top of the same DSpark MTP serving - # stack, at the same concurrency as the vllm-simple arm for a direct - # offload-backend comparison. - - { tp: 8, ep: 1, kv-offloading: dram, kv-offload-backend: { name: lmcache, version: "0.5.4rc1" }, conc-list: [10], spec-decoding: mtp } + +# LMCache MP-server DRAM offload on top of the same DSpark MTP serving stack as +# kimik3-fp4-mi355x-vllm-agentic-mtp (same image, script, and topology). A +# dedicated key so LMCache points can be selected and swept without re-running +# the resident and vllm-simple arms of the base key. +kimik3-fp4-mi355x-vllm-agentic-mtp-lmcache: + image: vllm/vllm-openai-rocm:nightly-cb8104839c141609d99f1254459ef3a4f1bd4263 + model: moonshotai/Kimi-K3 + model-prefix: kimik3 + runner: cluster:mi355x-amds + precision: fp4 + framework: vllm + multinode: false + scenarios: + agentic-coding: + - dram-utilization: 0.50 + search-space: + - { tp: 8, ep: 1, kv-offloading: dram, kv-offload-backend: { name: lmcache, version: "0.5.4rc1" }, conc-list: [4, 8, 16], spec-decoding: mtp } dsr1-fp4-mi355x-sglang-disagg: image: lmsysorg/sglang-rocm:v0.5.12-rocm720-mi35x-20260519 diff --git a/perf-changelog.yaml b/perf-changelog.yaml index da67c0138..b5e9172a3 100644 --- a/perf-changelog.yaml +++ b/perf-changelog.yaml @@ -5911,10 +5911,10 @@ pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2570 - config-keys: - - kimik3-fp4-mi355x-vllm-agentic-mtp + - kimik3-fp4-mi355x-vllm-agentic-mtp-lmcache scenario-type: - agentic-coding description: - - "Add an LMCache 0.5.4rc1 DRAM KV-offload arm at TP8 conc 10 on top of the DSpark MTP stack, mirroring the vllm-simple offload point for a direct backend comparison." + - "Add a dedicated LMCache 0.5.4rc1 DRAM KV-offload key at TP8 conc 4/8/16 on top of the unchanged kimik3-fp4-mi355x-vllm-agentic-mtp DSpark MTP stack." - "Run one LMCache MP server per node with chunk size 768 (the K3 unified block size at 8 GPUs) and --separate-object-groups for the hybrid KDA/MLA two-group KV layout." pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2583 From aba14d1b1f300be44daaf97ecf0693e4f5dcdd1e Mon Sep 17 00:00:00 2001 From: Samuel Shen Date: Wed, 12 Aug 2026 21:57:33 +0000 Subject: [PATCH 04/31] Lower the lmcache key's dram-utilization to fit the shm-backed L1 The LMCache MP server's L1 lives in /dev/shm and the script rejects budgets above 90% of free shm. mi355x-amds nodes mount ~1.5 TB of shm (cap ~1360 GB), so 0.50's 1499 GB budget failed the check in run 31644286169. 0.40 generates ~1199 GB, which fits with margin. --- configs/amd-master.yaml | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/configs/amd-master.yaml b/configs/amd-master.yaml index 5afd30620..96d3b2986 100644 --- a/configs/amd-master.yaml +++ b/configs/amd-master.yaml @@ -665,7 +665,11 @@ kimik3-fp4-mi355x-vllm-agentic-mtp-lmcache: multinode: false scenarios: agentic-coding: - - dram-utilization: 0.50 + # 0.40, not the base key's 0.50: the LMCache L1 is /dev/shm-backed and the + # script refuses budgets above 90% of free shm. mi355x-amds nodes mount + # ~1.5 TB of shm (cap ~1360 GB), so 0.50's 1499 GB budget fails the check + # (run 31644286169); 0.40 -> ~1199 GB fits with margin. + - dram-utilization: 0.40 search-space: - { tp: 8, ep: 1, kv-offloading: dram, kv-offload-backend: { name: lmcache, version: "0.5.4rc1" }, conc-list: [4, 8, 16], spec-decoding: mtp } From 6b95b7b830beec89a2417cc2d2478cc5a233443e Mon Sep 17 00:00:00 2001 From: Samuel Shen Date: Wed, 12 Aug 2026 22:10:05 +0000 Subject: [PATCH 05/31] Match the LMCache chunk size to this stack's 1536-token block size vLLM sizes the K3 unified attention block to 1536 tokens on the MI355X fp8-KV TRITON_MLA path (attention page >= mamba page), and the MP connector asserts chunk %% block == 0, so the recipe's CUDA-path 768 fails engine init (run 31644990546). --- .../single_node/agentic/kimik3_fp4_mi355x_mtp.sh | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh b/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh index 0f1355dfb..eba514b26 100644 --- a/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh +++ b/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh @@ -163,9 +163,13 @@ case "${KV_OFFLOAD_BACKEND:-}" in # One MP server for the node, per the Kimi-K3 recipe # (docs.lmcache.ai/recipes/kimi_k3.html). --chunk-size must equal the - # K3 unified block size N=768 at 8 GPUs, and the hybrid KDA/MLA layout - # (two KV-cache groups under MTP) requires one object group per - # sliding-window size: --separate-object-groups. + # unified attention block size, which THIS stack (fp8 KV, TP8, no + # mamba-cache-mode align) sets to 1536 -- "Setting attention block size + # to 1536 tokens to ensure that attention page size is >= mamba page + # size" (run 31644990546); the recipe's 768 is the CUDA-path value and + # fails the connector's chunk %% block == 0 assert here. The hybrid + # KDA/MLA layout (two KV-cache groups under MTP) requires one object + # group per sliding-window size: --separate-object-groups. LMCACHE_PORT=6555 LMCACHE_HTTP_PORT=8090 LMCACHE_LOG="$RESULT_DIR/lmcache_server.log" @@ -190,7 +194,7 @@ case "${KV_OFFLOAD_BACKEND:-}" in --http-port "$LMCACHE_HTTP_PORT" --l1-size-gb "$LMCACHE_L1_SIZE_GB" --l1-init-size-gb 10 - --chunk-size 768 + --chunk-size 1536 --separate-object-groups --enable-extra-logging --max-cpu-workers 8 From a8cda1afc0cfa50c85515fb7632dd5b7f027662c Mon Sep 17 00:00:00 2001 From: Samuel Shen Date: Wed, 12 Aug 2026 22:21:13 +0000 Subject: [PATCH 06/31] Raise the LMCache chunk size to 3072 for the KDA state group The connector requires the chunk to be a multiple of every engine KV group's tokens_per_block. On this stack the hybrid layout registers attention groups at 1536 and a KDA state group at 3072 (run 31645828378), so 1536 fails registration; 3072 is the minimum valid chunk. --- .../agentic/kimik3_fp4_mi355x_mtp.sh | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh b/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh index eba514b26..b9e58ee03 100644 --- a/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh +++ b/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh @@ -162,14 +162,15 @@ case "${KV_OFFLOAD_BACKEND:-}" in >/dev/null # One MP server for the node, per the Kimi-K3 recipe - # (docs.lmcache.ai/recipes/kimi_k3.html). --chunk-size must equal the - # unified attention block size, which THIS stack (fp8 KV, TP8, no - # mamba-cache-mode align) sets to 1536 -- "Setting attention block size - # to 1536 tokens to ensure that attention page size is >= mamba page - # size" (run 31644990546); the recipe's 768 is the CUDA-path value and - # fails the connector's chunk %% block == 0 assert here. The hybrid - # KDA/MLA layout (two KV-cache groups under MTP) requires one object - # group per sliding-window size: --separate-object-groups. + # (docs.lmcache.ai/recipes/kimi_k3.html), with --chunk-size sized for + # THIS stack rather than the recipe's CUDA-path 768: the connector + # requires the chunk to be a multiple of every engine KV group's + # tokens_per_block, and the hybrid KDA/MLA layout here registers + # attention groups at 1536 ("Setting attention block size to 1536", + # run 31644990546) plus a KDA state group at 3072 (run 31645828378), + # so 3072 is the minimum valid chunk. The multi-group layout also + # requires one object group per sliding-window size: + # --separate-object-groups. LMCACHE_PORT=6555 LMCACHE_HTTP_PORT=8090 LMCACHE_LOG="$RESULT_DIR/lmcache_server.log" @@ -194,7 +195,7 @@ case "${KV_OFFLOAD_BACKEND:-}" in --http-port "$LMCACHE_HTTP_PORT" --l1-size-gb "$LMCACHE_L1_SIZE_GB" --l1-init-size-gb 10 - --chunk-size 1536 + --chunk-size 3072 --separate-object-groups --enable-extra-logging --max-cpu-workers 8 From afee680555d81c8783ee45cfd44047f0a434fadf Mon Sep 17 00:00:00 2001 From: Samuel Shen Date: Wed, 12 Aug 2026 22:44:35 +0000 Subject: [PATCH 07/31] Pin the LMCache MP server to the lmcache_driven transfer path Auto mode loads both transfer paths; pin server-driven STORE/RETRIEVE (as the MiniMax-M3 arm does) so the benchmark measures one deterministic path. The L1 stays shm-backed either way, so the /dev/shm capacity check still applies. --- benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh b/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh index b9e58ee03..aa2a5ae89 100644 --- a/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh +++ b/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh @@ -201,6 +201,12 @@ case "${KV_OFFLOAD_BACKEND:-}" in --max-cpu-workers 8 --max-gpu-workers 1 --eviction-policy LRU + # Pin the server-driven STORE/RETRIEVE path (same as the MiniMax-M3 + # arm) so the benchmark measures one deterministic transfer path + # instead of the auto-mode pair. The L1 stays /dev/shm-backed either + # way (shm_name defaults on), which is why the capacity check above + # applies in this mode too. + --supported-transfer-mode lmcache_driven ) append_command "$RESULT_DIR/lmcache_command.txt" "${LMCACHE_CMD[@]}" "${LMCACHE_CMD[@]}" > "$LMCACHE_LOG" 2>&1 & From 20b4fdae81f816d7b9741f2c03d66b801f53b7dc Mon Sep 17 00:00:00 2001 From: Samuel Shen Date: Wed, 12 Aug 2026 23:21:43 +0000 Subject: [PATCH 08/31] Hold LMCache L1 read locks for the job duration The default 300s read-lock TTL expires under a single GPU worker serializing huge K3 transfers: run 31648224111 logged 57k finish-read-on-non-read-locked-key warnings starting exactly at warmup+300s, followed by a GPU illegal-access crash mid-profile. Match the MiniMax-M3 arm's 7200s read TTL. --- benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh b/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh index aa2a5ae89..55944729b 100644 --- a/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh +++ b/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh @@ -195,6 +195,13 @@ case "${KV_OFFLOAD_BACKEND:-}" in --http-port "$LMCACHE_HTTP_PORT" --l1-size-gb "$LMCACHE_L1_SIZE_GB" --l1-init-size-gb 10 + # Read locks default to a 300s TTL, and with a single GPU worker + # serializing 100k-300k-token K3 transfers, queued reads routinely + # outlive it: run 31648224111 logged 57k "finish read on + # non-read-locked key" warnings starting exactly warmup+300s before + # a GPU illegal-access crash. Outlast the job like the MiniMax-M3 + # arm does. + --l1-read-ttl-seconds 7200 --chunk-size 3072 --separate-object-groups --enable-extra-logging From 3c1908b33c48ba9a37a6efd7e031a672b5a7ae33 Mon Sep 17 00:00:00 2001 From: Samuel Shen Date: Wed, 12 Aug 2026 23:36:49 +0000 Subject: [PATCH 09/31] Revert "Hold LMCache L1 read locks for the job duration" This reverts commit 20b4fdae81f816d7b9741f2c03d66b801f53b7dc. --- benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh | 7 ------- 1 file changed, 7 deletions(-) diff --git a/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh b/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh index 55944729b..aa2a5ae89 100644 --- a/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh +++ b/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh @@ -195,13 +195,6 @@ case "${KV_OFFLOAD_BACKEND:-}" in --http-port "$LMCACHE_HTTP_PORT" --l1-size-gb "$LMCACHE_L1_SIZE_GB" --l1-init-size-gb 10 - # Read locks default to a 300s TTL, and with a single GPU worker - # serializing 100k-300k-token K3 transfers, queued reads routinely - # outlive it: run 31648224111 logged 57k "finish read on - # non-read-locked key" warnings starting exactly warmup+300s before - # a GPU illegal-access crash. Outlast the job like the MiniMax-M3 - # arm does. - --l1-read-ttl-seconds 7200 --chunk-size 3072 --separate-object-groups --enable-extra-logging From f2968297fcd26725a55ef1f84ed0bce902874f32 Mon Sep 17 00:00:00 2001 From: Sirra Date: Thu, 13 Aug 2026 12:39:37 +0530 Subject: [PATCH 10/31] [AMD] [WIP] [AGENTX] KIMI-K3 Perf Tuning Signed-off-by: Sirra --- benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh | 2 +- configs/amd-master.yaml | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh b/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh index ec2eaf3c3..41c261475 100644 --- a/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh +++ b/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh @@ -147,7 +147,7 @@ esac fi # ---- LLM server ------------------------------------------------------------ -bash "$(dirname "$0")/apply_k3_container_patches.sh" + # ---- Parallelism ------------------------------------------------------------ EP_ARGS=() diff --git a/configs/amd-master.yaml b/configs/amd-master.yaml index c94029484..da31d0bb4 100644 --- a/configs/amd-master.yaml +++ b/configs/amd-master.yaml @@ -637,7 +637,7 @@ dsr1-fp8-mi355x-sglang-disagg-mtp: - "DECODE_MTP_SIZE=2" kimik3-fp4-mi355x-vllm-agentic-mtp: - image: vllm/vllm-openai-rocm:nightly-cb8104839c141609d99f1254459ef3a4f1bd4263 + image: vllm/vllm-openai-rocm:nightly-3ee2df30337a301164c46ae444b76ee67e71c106 model: moonshotai/Kimi-K3 model-prefix: kimik3 runner: cluster:mi355x-amds @@ -648,7 +648,7 @@ kimik3-fp4-mi355x-vllm-agentic-mtp: agentic-coding: - dram-utilization: 0.50 search-space: - - { tp: 8, kv-offloading: none, conc-list: [1, 4, 8] , spec-decoding: mtp} + # - { tp: 8, kv-offloading: none, conc-list: [1, 4, 8] , spec-decoding: mtp} - { tp: 8, ep: 1, kv-offloading: dram, kv-offload-backend: { name: vllm-simple }, conc-list: [10], spec-decoding: mtp } dsr1-fp4-mi355x-sglang-disagg: From d28c525ce5eff6cf29fcd6f2d0eda32d5787eac6 Mon Sep 17 00:00:00 2001 From: Sirra Date: Thu, 13 Aug 2026 13:17:28 +0530 Subject: [PATCH 11/31] [AMD] [WIP] [AGENTX] KIMI-K3 Perf Tuning Signed-off-by: Sirra --- .../single_node/agentic/kimik3_fp4_mi355x_mtp.sh | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh b/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh index 41c261475..4a8215d3d 100644 --- a/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh +++ b/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh @@ -82,13 +82,13 @@ resolve_trace_source install_agentic_deps # ---- Reference env block ---------------------------------------------------- -export VLLM_ROCM_AITER_MLA_ASM_PADDING=asm +# export VLLM_ROCM_AITER_MLA_ASM_PADDING=asm export VLLM_ROCM_USE_AITER=1 -export SAFETENSORS_FAST_GPU=1 -export VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4=1 -export AITER_BF16_FP8_MOE_BOUND=0 -# REQUIRED on ROCm per the upstream recipe: the build auto-enables this to 1. -export VLLM_USE_BREAKABLE_CUDAGRAPH=0 +# export SAFETENSORS_FAST_GPU=1 +# export VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4=1 +# export AITER_BF16_FP8_MOE_BOUND=0 +# # REQUIRED on ROCm per the upstream recipe: the build auto-enables this to 1. +# export VLLM_USE_BREAKABLE_CUDAGRAPH=0 # Workaround for MEC FW <177 RCCL memory reclaim issue (shared with the other # gfx950 recipes in this tree). From 38faee4121442955997d6d531f56de8a66283bb6 Mon Sep 17 00:00:00 2001 From: Sirra Date: Thu, 13 Aug 2026 14:43:49 +0530 Subject: [PATCH 12/31] [AMD] [WIP] [AGENTX] KIMI-K3 Perf Tuning Signed-off-by: Sirra --- .../agentic/kimik3_fp4_mi355x_mtp.sh | 38 ++++++++++++++++--- 1 file changed, 32 insertions(+), 6 deletions(-) diff --git a/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh b/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh index 4a8215d3d..fe7ba1320 100644 --- a/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh +++ b/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh @@ -82,13 +82,18 @@ resolve_trace_source install_agentic_deps # ---- Reference env block ---------------------------------------------------- -# export VLLM_ROCM_AITER_MLA_ASM_PADDING=asm +# Keep ALL of these. Commenting them out does not avoid the AITER FMHA crash: +# that crash is gated on VLLM_ROCM_USE_AITER alone (AiterFlashAttnPrefillBackend +# .is_available() consults only rocm_aiter_ops.is_enabled()), so disabling the +# others just loses the MoE kernels while keeping the failure. The crash is +# avoided by pinning the MLA prefill backend in VLLM_CMD below. +export VLLM_ROCM_AITER_MLA_ASM_PADDING=asm export VLLM_ROCM_USE_AITER=1 -# export SAFETENSORS_FAST_GPU=1 -# export VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4=1 -# export AITER_BF16_FP8_MOE_BOUND=0 -# # REQUIRED on ROCm per the upstream recipe: the build auto-enables this to 1. -# export VLLM_USE_BREAKABLE_CUDAGRAPH=0 +export SAFETENSORS_FAST_GPU=1 +export VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4=1 +export AITER_BF16_FP8_MOE_BOUND=0 +# REQUIRED on ROCm per the upstream recipe: the build auto-enables this to 1. +export VLLM_USE_BREAKABLE_CUDAGRAPH=0 # Workaround for MEC FW <177 RCCL memory reclaim issue (shared with the other # gfx950 recipes in this tree). @@ -171,6 +176,26 @@ else ) fi +# ---- MLA prefill backend ----------------------------------------------------- +# On ROCm the prefill priority is [ROCM_AITER_FA, FLASH_ATTN]. ROCM_AITER_FA +# JIT-builds module_fmha_fwd_bf16_opus at runtime; that module registers its own +# aiter_tensor_t, distinct from the one in the prebuilt module_aiter_core, so the +# first call dies with: +# TypeError: fmha_fwd_bf16_opus_fwd(): incompatible function arguments +# during compile_or_warm_up_model -> _dummy_run, before the server binds. +# Pinning FLASH_ATTN keeps every AITER MoE kernel (and its throughput) while +# skipping only the broken FMHA prefill path. +# Set MLA_PREFILL_BACKEND=ROCM_AITER_FA to restore stock behaviour once the +# AITER packaging issue is fixed upstream. +MLA_PREFILL_BACKEND="${MLA_PREFILL_BACKEND:-FLASH_ATTN}" +MLA_PREFILL_ARGS=() +if [ -n "$MLA_PREFILL_BACKEND" ]; then + MLA_PREFILL_ARGS=( + --attention-config + "{\"mla_prefill_backend\":\"$MLA_PREFILL_BACKEND\"}" + ) +fi + # ---- HIP graph ------------------------------------------------------------ MAX_NUM_SEQS=20 MAX_CUDAGRAPH_CAPTURE_SIZE=60 @@ -202,6 +227,7 @@ VLLM_CMD=( --max-model-len 1048576 --enable-prefix-caching --kv-cache-dtype "fp8" + "${MLA_PREFILL_ARGS[@]}" "${COMPILATION_CONFIG_ARGS[@]}" "${SPEC_ARGS[@]}" "${OFFLOAD_ARGS[@]}" From 0811fb0c6962785d07607a9b1cffb621fe01473c Mon Sep 17 00:00:00 2001 From: ApostaC Date: Thu, 13 Aug 2026 08:50:20 -0700 Subject: [PATCH 13/31] trigger lmcache 0.5.4rc2 Signed-off-by: ApostaC --- configs/amd-master.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/configs/amd-master.yaml b/configs/amd-master.yaml index 96d3b2986..4b95cd413 100644 --- a/configs/amd-master.yaml +++ b/configs/amd-master.yaml @@ -671,7 +671,7 @@ kimik3-fp4-mi355x-vllm-agentic-mtp-lmcache: # (run 31644286169); 0.40 -> ~1199 GB fits with margin. - dram-utilization: 0.40 search-space: - - { tp: 8, ep: 1, kv-offloading: dram, kv-offload-backend: { name: lmcache, version: "0.5.4rc1" }, conc-list: [4, 8, 16], spec-decoding: mtp } + - { tp: 8, ep: 1, kv-offloading: dram, kv-offload-backend: { name: lmcache, version: "0.5.4rc2" }, conc-list: [4, 8, 16], spec-decoding: mtp } dsr1-fp4-mi355x-sglang-disagg: image: lmsysorg/sglang-rocm:v0.5.12-rocm720-mi35x-20260519 From ac6ad7d514a4a5346293e62a5378204d19e32f55 Mon Sep 17 00:00:00 2001 From: ApostaC Date: Thu, 13 Aug 2026 11:02:46 -0700 Subject: [PATCH 14/31] update lmcache configus Signed-off-by: ApostaC --- .../agentic/kimik3_fp4_mi355x_mtp.sh | 19 +++---------------- configs/amd-master.yaml | 10 +++++----- perf-changelog.yaml | 4 ++-- 3 files changed, 10 insertions(+), 23 deletions(-) diff --git a/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh b/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh index aa2a5ae89..0a3fd4318 100644 --- a/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh +++ b/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh @@ -150,7 +150,7 @@ case "${KV_OFFLOAD_BACKEND:-}" in # Keep the image's tested torch/ROCm stack and install only LMCache's # missing runtime dependencies, same as the MiniMax-M3 lmcache arm. - LMCACHE_VERSION="0.5.4rc1" + LMCACHE_VERSION="0.5.4rc2" LMCACHE_ROCM_INDEX="https://github.com/LMCache/LMCache/releases/expanded_assets/v${LMCACHE_VERSION}-rocm" agentic_pip_install --quiet --no-cache-dir --no-deps \ "sortedcontainers==2.4.0" \ @@ -175,17 +175,7 @@ case "${KV_OFFLOAD_BACKEND:-}" in LMCACHE_HTTP_PORT=8090 LMCACHE_LOG="$RESULT_DIR/lmcache_server.log" - # The whole generated node-DRAM budget backs the single server's L1, - # which lives in /dev/shm; fail early if it cannot fit. LMCACHE_L1_SIZE_GB="$TOTAL_CPU_DRAM_GB" - SHM_FREE_GB=$(df -BG --output=avail /dev/shm 2>/dev/null | tail -1 | tr -dc '0-9') - if [ -n "$SHM_FREE_GB" ] && [ "$SHM_FREE_GB" -gt 0 ]; then - SHM_CAP_GB=$((SHM_FREE_GB * 90 / 100)) - if [ "$LMCACHE_L1_SIZE_GB" -gt "$SHM_CAP_GB" ]; then - echo "Error: LMCache L1 ${LMCACHE_L1_SIZE_GB} GB exceeds 90% of free /dev/shm (${SHM_CAP_GB} GB)." >&2 - exit 1 - fi - fi LMCACHE_CMD=( lmcache server @@ -198,15 +188,12 @@ case "${KV_OFFLOAD_BACKEND:-}" in --chunk-size 3072 --separate-object-groups --enable-extra-logging + --extra-logging-interval 30 --max-cpu-workers 8 --max-gpu-workers 1 --eviction-policy LRU - # Pin the server-driven STORE/RETRIEVE path (same as the MiniMax-M3 - # arm) so the benchmark measures one deterministic transfer path - # instead of the auto-mode pair. The L1 stays /dev/shm-backed either - # way (shm_name defaults on), which is why the capacity check above - # applies in this mode too. --supported-transfer-mode lmcache_driven + --shm-name "" ) append_command "$RESULT_DIR/lmcache_command.txt" "${LMCACHE_CMD[@]}" "${LMCACHE_CMD[@]}" > "$LMCACHE_LOG" 2>&1 & diff --git a/configs/amd-master.yaml b/configs/amd-master.yaml index 4b95cd413..64b156590 100644 --- a/configs/amd-master.yaml +++ b/configs/amd-master.yaml @@ -665,11 +665,11 @@ kimik3-fp4-mi355x-vllm-agentic-mtp-lmcache: multinode: false scenarios: agentic-coding: - # 0.40, not the base key's 0.50: the LMCache L1 is /dev/shm-backed and the - # script refuses budgets above 90% of free shm. mi355x-amds nodes mount - # ~1.5 TB of shm (cap ~1360 GB), so 0.50's 1499 GB budget fails the check - # (run 31644286169); 0.40 -> ~1199 GB fits with margin. - - dram-utilization: 0.40 + # 0.50 matches the base key: the LMCache server runs with --shm-name "" + # so its L1 lives in regular process memory instead of /dev/shm, and the + # budget is no longer capped by the ~1.5 TB shm mount (which forced 0.40 + # before, run 31644286169). + - dram-utilization: 0.50 search-space: - { tp: 8, ep: 1, kv-offloading: dram, kv-offload-backend: { name: lmcache, version: "0.5.4rc2" }, conc-list: [4, 8, 16], spec-decoding: mtp } diff --git a/perf-changelog.yaml b/perf-changelog.yaml index d8619dec9..9b0682f01 100644 --- a/perf-changelog.yaml +++ b/perf-changelog.yaml @@ -5924,6 +5924,6 @@ scenario-type: - agentic-coding description: - - "Add a dedicated LMCache 0.5.4rc1 DRAM KV-offload key at TP8 conc 4/8/16 on top of the unchanged kimik3-fp4-mi355x-vllm-agentic-mtp DSpark MTP stack." - - "Run one LMCache MP server per node with chunk size 768 (the K3 unified block size at 8 GPUs) and --separate-object-groups for the hybrid KDA/MLA two-group KV layout." + - "Add a dedicated LMCache 0.5.4rc2 DRAM KV-offload key at TP8 conc 4/8/16 on top of the unchanged kimik3-fp4-mi355x-vllm-agentic-mtp DSpark MTP stack, with the version pinned in the master config and consumed by the script via KV_OFFLOAD_BACKEND_METADATA." + - "Run one LMCache MP server per node with chunk size 3072 (the minimum multiple of the hybrid KDA/MLA group block sizes) and --separate-object-groups, keeping the L1 in process memory (--shm-name \"\") so the DRAM budget is not capped by /dev/shm." pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2583 From e890a2b44a770a2cfec61a9b025fbad6fcbe8e69 Mon Sep 17 00:00:00 2001 From: ApostaC Date: Thu, 13 Aug 2026 11:25:39 -0700 Subject: [PATCH 15/31] update conc Signed-off-by: ApostaC --- configs/amd-master.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/configs/amd-master.yaml b/configs/amd-master.yaml index 64b156590..44075c764 100644 --- a/configs/amd-master.yaml +++ b/configs/amd-master.yaml @@ -671,7 +671,7 @@ kimik3-fp4-mi355x-vllm-agentic-mtp-lmcache: # before, run 31644286169). - dram-utilization: 0.50 search-space: - - { tp: 8, ep: 1, kv-offloading: dram, kv-offload-backend: { name: lmcache, version: "0.5.4rc2" }, conc-list: [4, 8, 16], spec-decoding: mtp } + - { tp: 8, ep: 1, kv-offloading: dram, kv-offload-backend: { name: lmcache, version: "0.5.4rc2" }, conc-list: [4, 8, 10, 12], spec-decoding: mtp } dsr1-fp4-mi355x-sglang-disagg: image: lmsysorg/sglang-rocm:v0.5.12-rocm720-mi35x-20260519 From 96783731b0e849a118768ccb38b62cb23393f04d Mon Sep 17 00:00:00 2001 From: Samuel Shen Date: Thu, 13 Aug 2026 18:17:37 -0700 Subject: [PATCH 16/31] perf-changelog: point the LMCache MI355X entry at #2598 --- perf-changelog.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/perf-changelog.yaml b/perf-changelog.yaml index 9b0682f01..cb082b6c4 100644 --- a/perf-changelog.yaml +++ b/perf-changelog.yaml @@ -5926,4 +5926,4 @@ description: - "Add a dedicated LMCache 0.5.4rc2 DRAM KV-offload key at TP8 conc 4/8/16 on top of the unchanged kimik3-fp4-mi355x-vllm-agentic-mtp DSpark MTP stack, with the version pinned in the master config and consumed by the script via KV_OFFLOAD_BACKEND_METADATA." - "Run one LMCache MP server per node with chunk size 3072 (the minimum multiple of the hybrid KDA/MLA group block sizes) and --separate-object-groups, keeping the L1 in process memory (--shm-name \"\") so the DRAM budget is not capped by /dev/shm." - pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2583 + pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2598 From c0119e6be379f35102386fccd5500e90398dd3ad Mon Sep 17 00:00:00 2001 From: Sirra Date: Fri, 14 Aug 2026 07:51:03 +0530 Subject: [PATCH 17/31] [AMD] [AGENTX] KIMI-K3 Perf Tuning Signed-off-by: Sirra --- .../agentic/apply_aiter_pybind11_fix.sh | 118 ++++++++++++++++++ .../agentic/apply_triton_mla_cudagraph_fix.sh | 65 ++++++++++ .../agentic/kimik3_fp4_mi355x_mtp.sh | 34 ++++- 3 files changed, 212 insertions(+), 5 deletions(-) create mode 100644 benchmarks/single_node/agentic/apply_aiter_pybind11_fix.sh create mode 100644 benchmarks/single_node/agentic/apply_triton_mla_cudagraph_fix.sh diff --git a/benchmarks/single_node/agentic/apply_aiter_pybind11_fix.sh b/benchmarks/single_node/agentic/apply_aiter_pybind11_fix.sh new file mode 100644 index 000000000..d69021ffa --- /dev/null +++ b/benchmarks/single_node/agentic/apply_aiter_pybind11_fix.sh @@ -0,0 +1,118 @@ +#!/usr/bin/env bash +# Fix AITER JIT modules building against a different pybind11 than the prebuilt +# .so files shipped in the image. +# +# Symptom (kills the server during warmup, before it binds): +# TypeError: fmha_fwd_bf16_opus_fwd(): incompatible function arguments +# RuntimeError: Engine core initialization failed +# +# Cause: aiter/jit/utils/cpp_extension.py adds the standalone pybind11 include +# via -I, which outranks the -isystem path holding torch's bundled pybind11. +# The prebuilt modules (module_aiter_core.so) use torch's pybind11 +# (PYBIND11_INTERNALS_VERSION 11); the standalone package is version 12. +# pybind11 keeps a SEPARATE type registry per internals id, so the JIT module +# cannot see aiter_tensor_t registered by the core module. +# +# Effect: lets ROCM_AITER_FA be used for MLA prefill instead of falling back to +# FLASH_ATTN. Measured on 8x MI355X / Kimi-K3 MXFP4 TP8: +# ~24k ctx 12,953 -> 13,524 tok/s (+4.4%) +# ~93k ctx 11,174 -> 13,423 tok/s (+20.1%) +# +# Idempotent: safe to run repeatedly. No-op if already patched or not needed. +set -euo pipefail + +PY=${PYTHON:-python3} + +TARGET=$($PY - <<'EOF' +import os +try: + import aiter.jit.utils.cpp_extension as m + print(os.path.abspath(m.__file__)) +except Exception: + print("") +EOF +) + +if [ -z "$TARGET" ] || [ ! -f "$TARGET" ]; then + echo "[aiter-fix] aiter cpp_extension not found; nothing to do." + exit 0 +fi + +# Is there actually an internals-version mismatch to fix? +NEED=$($PY - <<'EOF' +import os, re +try: + import torch, pybind11 +except Exception: + print("no"); raise SystemExit +def ver(p): + f = os.path.join(p, "pybind11", "detail", "internals.h") + if not os.path.isfile(f): + return None + m = re.search(r"define\s+PYBIND11_INTERNALS_VERSION\s+(\d+)", open(f).read()) + return int(m.group(1)) if m else None +t = ver(os.path.join(os.path.dirname(torch.__file__), "include")) +s = ver(pybind11.get_include()) +print("yes" if (t is not None and s is not None and t != s) else "no") +EOF +) + +if [ "$NEED" != "yes" ]; then + echo "[aiter-fix] pybind11 internals versions already agree; no patch needed." + exit 0 +fi + +if grep -q "_use_torch_pybind11" "$TARGET"; then + echo "[aiter-fix] already patched: $TARGET" +else + cp -n "$TARGET" "$TARGET.orig" || true + $PY - "$TARGET" <<'EOF' +import sys, io +path = sys.argv[1] +src = io.open(path, encoding="utf-8").read() +old = " extra_include_paths.append(pybind11.get_include())\n" +new = ( + " # PATCHED: prefer torch's bundled pybind11 so JIT modules land in the\n" + " # same pybind11 type registry as the prebuilt .so files. Mismatched\n" + " # PYBIND11_INTERNALS_VERSION otherwise yields:\n" + " # TypeError: ...(): incompatible function arguments\n" + " _use_torch_pybind11 = False\n" + " if not torch_exclude:\n" + " _use_torch_pybind11 = os.path.isdir(\n" + " os.path.join(TORCH_INCLUDE_ROOT, \"pybind11\")\n" + " )\n" + " if not _use_torch_pybind11:\n" + " extra_include_paths.append(pybind11.get_include())\n" +) +if old not in src: + sys.stderr.write("[aiter-fix] ERROR: anchor line not found; aborting.\n") + sys.exit(1) +if src.count(old) != 1: + sys.stderr.write("[aiter-fix] ERROR: anchor line not unique; aborting.\n") + sys.exit(1) +io.open(path, "w", encoding="utf-8").write(src.replace(old, new)) +print("[aiter-fix] patched", path) +EOF +fi + +# Drop JIT artifacts built against the wrong pybind11 so they rebuild. +# Ask aiter for its jit dir: it honours AITER_JIT_DIR and falls back to ~/.aiter +# when dist-packages is not writable, so deriving it from $TARGET is wrong. +JITDIR=$($PY -c 'from aiter.jit.core import get_user_jit_dir; print(get_user_jit_dir())' 2>/dev/null || true) +if [ -z "$JITDIR" ] || [ ! -d "$JITDIR" ]; then + JITDIR=$(dirname "$(dirname "$TARGET")") +fi +echo "[aiter-fix] jit dir: $JITDIR" +# Sweep every module, not just the one we happened to hit first: any module +# JIT-built before the patch carries the wrong internals id. +shopt -s nullglob +for so in "$JITDIR"/*.so; do + if grep -qa "__pybind11_internals_v12" "$so" 2>/dev/null; then + rm -f "$so" + rm -rf "$JITDIR/build/$(basename "${so%.so}")" + echo "[aiter-fix] removed stale v12 module: $(basename "$so")" + fi +done +shopt -u nullglob + +echo "[aiter-fix] done." diff --git a/benchmarks/single_node/agentic/apply_triton_mla_cudagraph_fix.sh b/benchmarks/single_node/agentic/apply_triton_mla_cudagraph_fix.sh new file mode 100644 index 000000000..01c1634fa --- /dev/null +++ b/benchmarks/single_node/agentic/apply_triton_mla_cudagraph_fix.sh @@ -0,0 +1,65 @@ +#!/usr/bin/env bash +# Let DSpark spec-decode keep FULL cudagraphs on ROCm. +# +# vllm/v1/attention/backends/mla/triton_mla.py declares +# TritonMLAMetadataBuilder._cudagraph_support = UNIFORM_SINGLE_TOKEN_DECODE +# which caps min_cg_support below UNIFORM_BATCH, so config/compilation.py:1443 +# downgrades FULL_AND_PIECEWISE -> PIECEWISE whenever spec-decode is enabled: +# "CUDAGraphMode.FULL_AND_PIECEWISE is not supported with spec-decode for +# attention backend TritonMLABackend" +# v1/worker/gpu/spec_decode/dflash/speculator.py:110-127 then gives the DSpark +# drafter CUDAGraphMode.NONE -- fully eager, with NO warning logged. Every draft +# layer and the Markov head dispatch kernel-by-kernel from Python each step. +# +# TRITON_MLA cannot simply be swapped out: it is the only ROCm MLA backend with +# supports_non_causal_multi_token_decode = True, which the DSpark draft needs. +# ROCM_AITER_MLA fails with "non-causal attention not supported". +# +# The builder already sets supports_non_causal_multi_token_decode = True and +# calls _init_reorder_batch_threshold(1, supports_spec_as_decode=True) "so +# full-cudagraph capture admits it", so UNIFORM_BATCH is the consistent value. +# +# MEASURED on 8x MI355X, Kimi-K3 MXFP4 TP8, DSpark, single stream, 600-tok gens: +# before: 14.05 tok/s, ITL 71.16 ms (PIECEWISE, drafter eager) +# after : 77.65 tok/s, ITL 12.88 ms (FULL cudagraphs) = 5.52x +# output verified correct in both ("17*23" -> 391, finish_reason stop) +# +# Idempotent. No-op if already patched or if the anchor is absent. +set -euo pipefail +PY=${PYTHON:-python3} + +TARGET=$($PY - <<'EOF' +import os +try: + import vllm.v1.attention.backends.mla.triton_mla as m + print(os.path.abspath(m.__file__)) +except Exception: + print("") +EOF +) +if [ -z "$TARGET" ] || [ ! -f "$TARGET" ]; then + echo "[triton-mla-fix] triton_mla.py not found; nothing to do." + exit 0 +fi +if grep -q "AttentionCGSupport.UNIFORM_BATCH" "$TARGET"; then + echo "[triton-mla-fix] already patched: $TARGET" + exit 0 +fi +cp -n "$TARGET" "$TARGET.orig" || true +$PY - "$TARGET" <<'EOF' +import sys, io +path = sys.argv[1] +src = io.open(path, encoding="utf-8").read() +old = """ _cudagraph_support: ClassVar[AttentionCGSupport] = ( + AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE + )""" +new = """ # PATCHED: UNIFORM_SINGLE_TOKEN_DECODE forced a PIECEWISE downgrade under + # spec-decode, which silently made the DSpark drafter fully eager. + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH""" +if src.count(old) != 1: + sys.stderr.write("[triton-mla-fix] anchor missing or not unique; aborting.\n") + sys.exit(1) +io.open(path, "w", encoding="utf-8").write(src.replace(old, new)) +print("[triton-mla-fix] patched", path) +EOF +echo "[triton-mla-fix] done." diff --git a/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh b/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh index fe7ba1320..87295393b 100644 --- a/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh +++ b/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh @@ -81,12 +81,31 @@ amd-smi || true resolve_trace_source install_agentic_deps +# ---- AITER pybind11 fix ------------------------------------------------------ +# The image's prebuilt aiter .so files are compiled against torch's bundled +# pybind11 (PYBIND11_INTERNALS_VERSION 11), but aiter's JIT builder injects the +# standalone pybind11 3.1.0 (version 12) as a -I flag, which outranks the +# -isystem path holding torch's copy. pybind11 keeps a SEPARATE type registry +# per internals id, so a JIT-built module cannot see aiter_tensor_t registered +# by the prebuilt core, and the first call dies during model warmup with: +# TypeError: fmha_fwd_bf16_opus_fwd(): incompatible function arguments +# The script below is idempotent, verifies the mismatch actually exists before +# touching anything, and self-disables once the image ships a fixed aiter. +bash "$(dirname "$0")/apply_aiter_pybind11_fix.sh" || true + +# ---- DSpark FULL-cudagraph fix ---------------------------------------------- +# TritonMLA declares _cudagraph_support=UNIFORM_SINGLE_TOKEN_DECODE, which forces +# FULL_AND_PIECEWISE -> PIECEWISE under spec-decode and then silently gives the +# DSpark drafter CUDAGraphMode.NONE (fully eager). Measured on 8x MI355X, single +# stream 600-token generations: 14.05 -> 77.65 tok/s, ITL 71.16 -> 12.88 ms (5.52x), +# output verified correct. Idempotent; no-op if already patched. +bash "$(dirname "$0")/apply_triton_mla_cudagraph_fix.sh" || true + # ---- Reference env block ---------------------------------------------------- # Keep ALL of these. Commenting them out does not avoid the AITER FMHA crash: # that crash is gated on VLLM_ROCM_USE_AITER alone (AiterFlashAttnPrefillBackend # .is_available() consults only rocm_aiter_ops.is_enabled()), so disabling the -# others just loses the MoE kernels while keeping the failure. The crash is -# avoided by pinning the MLA prefill backend in VLLM_CMD below. +# others just loses the MoE kernels while keeping the failure. export VLLM_ROCM_AITER_MLA_ASM_PADDING=asm export VLLM_ROCM_USE_AITER=1 export SAFETENSORS_FAST_GPU=1 @@ -185,9 +204,14 @@ fi # during compile_or_warm_up_model -> _dummy_run, before the server binds. # Pinning FLASH_ATTN keeps every AITER MoE kernel (and its throughput) while # skipping only the broken FMHA prefill path. -# Set MLA_PREFILL_BACKEND=ROCM_AITER_FA to restore stock behaviour once the -# AITER packaging issue is fixed upstream. -MLA_PREFILL_BACKEND="${MLA_PREFILL_BACKEND:-FLASH_ATTN}" +# UPDATE: the AITER packaging issue is now fixed at source by +# apply_aiter_pybind11_fix.sh (run above), so ROCM_AITER_FA is usable again and +# is the default. Measured on 8x MI355X / Kimi-K3 MXFP4 TP8, cold prefill: +# ~24k ctx FLASH_ATTN 12,953 -> AITER 13,524 tok/s (+4.4%) +# ~93k ctx FLASH_ATTN 11,174 -> AITER 13,423 tok/s (+20.1%) +# This workload averages ~99k input tokens, so the ~93k figure is the relevant +# one. Set MLA_PREFILL_BACKEND=FLASH_ATTN to fall back if AITER regresses. +MLA_PREFILL_BACKEND="${MLA_PREFILL_BACKEND:-ROCM_AITER_FA}" MLA_PREFILL_ARGS=() if [ -n "$MLA_PREFILL_BACKEND" ]; then MLA_PREFILL_ARGS=( From eb81b487cae86457e5b9454c34968be172b38ae1 Mon Sep 17 00:00:00 2001 From: Sirra Date: Fri, 14 Aug 2026 08:08:24 +0530 Subject: [PATCH 18/31] [AMD] [AGENTX] KIMI-K3 Perf Tuning Signed-off-by: Sirra --- perf-changelog.yaml | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/perf-changelog.yaml b/perf-changelog.yaml index 6082f2b86..e75cdd87f 100644 --- a/perf-changelog.yaml +++ b/perf-changelog.yaml @@ -5918,3 +5918,10 @@ - "Add a TEP2 arm (tp 2, ep 2) to the qwen3.5-fp4-b200-sglang-mtp 8k/1k sweep at concurrency 16, 32, and 64" - "Rides on the NVFP4-V2 checkpoint switch from #2205" pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2550 + +- config-keys: + - kimik3-fp4-mi355x-vllm-agentic-mtp + description: + - "Kimi-K3 Perf Tuning with AITER Backend" + - "Docker IMage Pinned : vllm/vllm-openai-rocm:nightly-3ee2df30337a301164c46ae444b76ee67e71c106" + pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2602 From 836681531c7ea7deb9fddf524b1e2791d9c3e9a6 Mon Sep 17 00:00:00 2001 From: Sirra Date: Fri, 14 Aug 2026 13:44:30 +0530 Subject: [PATCH 19/31] [AMD] [WIP] [AGENTX] KIMI-K3 Perf with fixes applied. Signed-off-by: Sirra --- .../agentic/apply_aiter_pybind11_fix.sh | 118 -- .../agentic/apply_k3_container_patches.sh | 1302 ----------------- .../agentic/apply_triton_mla_cudagraph_fix.sh | 65 - .../agentic/kimik3_fp4_mi355x_mtp.sh | 54 +- 4 files changed, 34 insertions(+), 1505 deletions(-) delete mode 100644 benchmarks/single_node/agentic/apply_aiter_pybind11_fix.sh delete mode 100755 benchmarks/single_node/agentic/apply_k3_container_patches.sh delete mode 100644 benchmarks/single_node/agentic/apply_triton_mla_cudagraph_fix.sh diff --git a/benchmarks/single_node/agentic/apply_aiter_pybind11_fix.sh b/benchmarks/single_node/agentic/apply_aiter_pybind11_fix.sh deleted file mode 100644 index d69021ffa..000000000 --- a/benchmarks/single_node/agentic/apply_aiter_pybind11_fix.sh +++ /dev/null @@ -1,118 +0,0 @@ -#!/usr/bin/env bash -# Fix AITER JIT modules building against a different pybind11 than the prebuilt -# .so files shipped in the image. -# -# Symptom (kills the server during warmup, before it binds): -# TypeError: fmha_fwd_bf16_opus_fwd(): incompatible function arguments -# RuntimeError: Engine core initialization failed -# -# Cause: aiter/jit/utils/cpp_extension.py adds the standalone pybind11 include -# via -I, which outranks the -isystem path holding torch's bundled pybind11. -# The prebuilt modules (module_aiter_core.so) use torch's pybind11 -# (PYBIND11_INTERNALS_VERSION 11); the standalone package is version 12. -# pybind11 keeps a SEPARATE type registry per internals id, so the JIT module -# cannot see aiter_tensor_t registered by the core module. -# -# Effect: lets ROCM_AITER_FA be used for MLA prefill instead of falling back to -# FLASH_ATTN. Measured on 8x MI355X / Kimi-K3 MXFP4 TP8: -# ~24k ctx 12,953 -> 13,524 tok/s (+4.4%) -# ~93k ctx 11,174 -> 13,423 tok/s (+20.1%) -# -# Idempotent: safe to run repeatedly. No-op if already patched or not needed. -set -euo pipefail - -PY=${PYTHON:-python3} - -TARGET=$($PY - <<'EOF' -import os -try: - import aiter.jit.utils.cpp_extension as m - print(os.path.abspath(m.__file__)) -except Exception: - print("") -EOF -) - -if [ -z "$TARGET" ] || [ ! -f "$TARGET" ]; then - echo "[aiter-fix] aiter cpp_extension not found; nothing to do." - exit 0 -fi - -# Is there actually an internals-version mismatch to fix? -NEED=$($PY - <<'EOF' -import os, re -try: - import torch, pybind11 -except Exception: - print("no"); raise SystemExit -def ver(p): - f = os.path.join(p, "pybind11", "detail", "internals.h") - if not os.path.isfile(f): - return None - m = re.search(r"define\s+PYBIND11_INTERNALS_VERSION\s+(\d+)", open(f).read()) - return int(m.group(1)) if m else None -t = ver(os.path.join(os.path.dirname(torch.__file__), "include")) -s = ver(pybind11.get_include()) -print("yes" if (t is not None and s is not None and t != s) else "no") -EOF -) - -if [ "$NEED" != "yes" ]; then - echo "[aiter-fix] pybind11 internals versions already agree; no patch needed." - exit 0 -fi - -if grep -q "_use_torch_pybind11" "$TARGET"; then - echo "[aiter-fix] already patched: $TARGET" -else - cp -n "$TARGET" "$TARGET.orig" || true - $PY - "$TARGET" <<'EOF' -import sys, io -path = sys.argv[1] -src = io.open(path, encoding="utf-8").read() -old = " extra_include_paths.append(pybind11.get_include())\n" -new = ( - " # PATCHED: prefer torch's bundled pybind11 so JIT modules land in the\n" - " # same pybind11 type registry as the prebuilt .so files. Mismatched\n" - " # PYBIND11_INTERNALS_VERSION otherwise yields:\n" - " # TypeError: ...(): incompatible function arguments\n" - " _use_torch_pybind11 = False\n" - " if not torch_exclude:\n" - " _use_torch_pybind11 = os.path.isdir(\n" - " os.path.join(TORCH_INCLUDE_ROOT, \"pybind11\")\n" - " )\n" - " if not _use_torch_pybind11:\n" - " extra_include_paths.append(pybind11.get_include())\n" -) -if old not in src: - sys.stderr.write("[aiter-fix] ERROR: anchor line not found; aborting.\n") - sys.exit(1) -if src.count(old) != 1: - sys.stderr.write("[aiter-fix] ERROR: anchor line not unique; aborting.\n") - sys.exit(1) -io.open(path, "w", encoding="utf-8").write(src.replace(old, new)) -print("[aiter-fix] patched", path) -EOF -fi - -# Drop JIT artifacts built against the wrong pybind11 so they rebuild. -# Ask aiter for its jit dir: it honours AITER_JIT_DIR and falls back to ~/.aiter -# when dist-packages is not writable, so deriving it from $TARGET is wrong. -JITDIR=$($PY -c 'from aiter.jit.core import get_user_jit_dir; print(get_user_jit_dir())' 2>/dev/null || true) -if [ -z "$JITDIR" ] || [ ! -d "$JITDIR" ]; then - JITDIR=$(dirname "$(dirname "$TARGET")") -fi -echo "[aiter-fix] jit dir: $JITDIR" -# Sweep every module, not just the one we happened to hit first: any module -# JIT-built before the patch carries the wrong internals id. -shopt -s nullglob -for so in "$JITDIR"/*.so; do - if grep -qa "__pybind11_internals_v12" "$so" 2>/dev/null; then - rm -f "$so" - rm -rf "$JITDIR/build/$(basename "${so%.so}")" - echo "[aiter-fix] removed stale v12 module: $(basename "$so")" - fi -done -shopt -u nullglob - -echo "[aiter-fix] done." diff --git a/benchmarks/single_node/agentic/apply_k3_container_patches.sh b/benchmarks/single_node/agentic/apply_k3_container_patches.sh deleted file mode 100755 index 1a9938e57..000000000 --- a/benchmarks/single_node/agentic/apply_k3_container_patches.sh +++ /dev/null @@ -1,1302 +0,0 @@ -#!/usr/bin/env bash -# ============================================================================= -# apply_k3_cb8104839c_fp8_embedded.sh (PINNED / offline) -# -# Reproduces, BYTE-FOR-BYTE, the patched Python source of the working Kimi-K3 -# fp8-KV FULL_AND_PIECEWISE cudagraph container `k3_srok_cb810_0810_replay` on a -# FRESH container of: -# vllm/vllm-openai-rocm:nightly-cb8104839c141609d99f1254459ef3a4f1bd4263 -# -# Code changes are EMBEDDED as pristine->container diffs (no GitHub / no PR -# drift). Net effect of, in the container: -# aiter #4474 int64 KV stride (mla_gluon >2GB global_load) -# aiter #4494 a16w16 GEMM fresh split-K semaphore under cudagraph capture -# vllm #51171 FULL cudagraphs for AITER MLA speculative decoding -# vllm #50578 asm decode for non-divisor small head counts (12->16 @ TP8) -# vllm #51011 fix fp8 KV cache decode on the AITER MLA backend -# vllm #51040 extend FP8 asm MLA prefill to non-divisor small head counts -# vllm #50619 (PARTIAL) cudagraph-exclude draft-attn layers + nvidia MLA -# fallback gate: gpu/attn_utils.py, gpu/model_runner.py, -# kimi_k3/nvidia/mla.py (rocm_aiter_mla.py hunks NOT taken -- -# they conflict with the #50578/#51011 asm strategy) -# vllm #51682 KDA packed decode: pass the state-index stride to the kernel so -# a non-contiguous 1-D state_indices is handled natively (only -# requires ndim==1). Replaces the earlier reshape/coerce -# workaround. NOTE: not strictly needed by this stack (it boots -# under FULL_AND_PIECEWISE without it) -- kept for robustness. -# aiter #4521 fp8 cp round-robin asm MLA verify kernels: adds the qh16/qh32 -# qseqlen4 gqaratio16/32 cprr .co + mla_asm.csv + asm_mla.cu + -# v1_2_device.cuh + aiter/mla.py + aiter/ops/attention.py, then -# rebuilds module_mla_asm. [NEEDS NETWORK + hipcc + a GPU: -# unlike the offline Python diffs, this fetches the binary .co -# and recompiles the asm module.] -# DSpark PS verify: route the small-head fp8 DSpark TARGET VERIFY to the ASM -# persistent (PS) decode instead of the Gluon flatten. Two edits -# on rocm_aiter_mla.py: (a) use_gluon_verify returns False for -# fp8 KV so the verify is NOT swallowed by the flatten, (b) -# _mtp_decode_qlen is sized for DSpark (1 + num_spec) so the PS -# gate opens. Needs aiter #4521 for the fp8 qseqlen4 verify -# kernels. SUPERSEDES the earlier HYBRID (gluon-flatten) verify. -# mla_gluon bh16bn128 batch<=256 relax + fp8-query dequant are -# kept (used by the bf16 verify path). -# triton 3.7.0 (AMD ROCm 7.2.0) + tabulate + lm_eval[api]==0.4.12 -# -# Run INSIDE a fresh container of that image: -# docker exec -i bash < apply_k3_cb8104839c_fp8_embedded.sh -# -# RUNTIME NOTE (NOT a code change -- set in your server script): -# * MODEL_PATH must point at the model INSIDE the container (e.g. /model/Kimi-K3 -# when launched with `-v /data:/model`). -# * fp8 PIECEWISE capture memory-faults at capture size 45 -> cap below it: -# "max_cudagraph_capture_size": 44, "cudagraph_mode": "FULL_AND_PIECEWISE", -# "custom_ops": ["+fused_rms_norm_gated"]. -# * env: VLLM_ROCM_USE_AITER=1, VLLM_ROCM_AITER_MLA_ASM_PADDING=asm, -# VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4=1, --kv-cache-dtype fp8, -# --enable-prefix-caching, DSpark spec-decode attention_backend=TRITON_MLA. -# ============================================================================= -set -uo pipefail - -# Resolve install root WITHOUT importing (importing aiter runs rocminfo and -# aborts on a GPU-less container). vllm and aiter share one dist-packages dir. -ROOT="$(python -c 'import importlib.util as u, os; print(os.path.dirname(os.path.dirname(u.find_spec("vllm").origin)))')" -if [ -z "$ROOT" ] || [ ! -d "$ROOT/vllm" ] || [ ! -d "$ROOT/aiter" ]; then - echo "ERROR: could not resolve dist-packages (ROOT='$ROOT')"; exit 1 -fi -echo "[embed] ROOT=$ROOT" -WS="${WS:-/tmp/k3_embed}"; mkdir -p "$WS" -say(){ echo; echo "=================== $* ==================="; } - -say "1/4 triton 3.7.0 + tabulate + lm_eval" -python -m pip install --extra-index-url https://pypi.amd.com/triton/release/rocm-7.2.0/simple/ triton==3.7.0 2>&1 | tail -2 -python -m pip install tabulate 2>&1 | tail -1 -if [ "${WITH_LM_EVAL:-1}" = "1" ]; then - python -m pip install "lm_eval[api]==0.4.12" 2>&1 | tail -2 -fi - -# Marker-gated apply: skip if the post-state marker is already present. -apply_one(){ # $1=relpath $2=marker $3=difffile - local f="$ROOT/$1" - if grep -qF "$2" "$f" 2>/dev/null; then echo " $1: already present (skip)"; return; fi - if ( cd "$ROOT" && git apply -p1 "$3" ) 2>/dev/null; then - echo " $1: APPLIED (git apply)" - else - patch -p1 -d "$ROOT" --fuzz=3 --forward --no-backup-if-mismatch < "$3" \ - && echo " $1: APPLIED (patch)" || echo " $1: FAILED" - fi -} - -say "2/4 apply embedded code changes" -cat > "$WS/MLA_GLUON.diff" <<'DIFF_MLA_GLUON' -diff --git a/aiter/ops/triton/gluon/mla_gluon.py b/aiter/ops/triton/gluon/mla_gluon.py ---- a/aiter/ops/triton/gluon/mla_gluon.py -+++ b/aiter/ops/triton/gluon/mla_gluon.py -@@ -156,6 +156,11 @@ - num_iter = gl.cdiv(split_kv_end - split_kv_start, BLOCK_N) - start_n = split_kv_start - -+ # >2GB KV cache (global_load path): widen strides to int64 so kv offsets don't overflow int32. -+ if not WITHIN_2GB: -+ stride_kv_c_bs = stride_kv_c_bs.to(gl.int64) -+ stride_k_pe_bs = stride_k_pe_bs.to(gl.int64) -+ - # early return with empty kv slice to save compute - if split_kv_start >= split_kv_end: - return -@@ -861,6 +866,11 @@ - kv_pe_offset = 0 - use_2d_view = False - -+ if q_nope.dtype == torch.float8_e4m3fn: -+ q_nope = q_nope.to(torch.bfloat16) -+ if q_pe is not None and q_pe.dtype == torch.float8_e4m3fn: -+ q_pe = q_pe.to(torch.bfloat16) -+ - assert ( - arch_info.get_arch() == "gfx950" - ), f"mla_gluon requires gfx950 (CDNA4), got {arch_info.get_arch()}" -@@ -931,9 +941,11 @@ - # NUM_KV_SPLITS >= 1). Each clamp below keeps NUM_KV_SPLITS <= min_kv_seq_len, - if REGIME == "bh16bn128": - assert ( -- batch_size == 1 -- ), f"mla_gluon[bh16bn128] requires batch_size=1, got {batch_size}" -- NUM_KV_SPLITS = max(1, min(256 // (batch_size * qlen), min_kv_seq_len)) -+ 1 <= batch_size <= 256 -+ ), f"mla_gluon[bh16bn128] requires 1 <= batch_size <= 256, got {batch_size}" -+ NUM_KV_SPLITS = max( -+ 1, min(256 // (batch_size * qlen), triton.cdiv(min_kv_seq_len, BLOCK_N)) -+ ) - else: # bh16bn64 - # Fill ~256 WGs (total WGs = B * NUM_KV_SPLITS <= 256, one MI350 wave), - # but never split a sequence into more blocks than it has: bound by the -DIFF_MLA_GLUON -apply_one "aiter/ops/triton/gluon/mla_gluon.py" "1 <= batch_size <= 256" "$WS/MLA_GLUON.diff" - -cat > "$WS/GEMM_A16W16.diff" <<'DIFF_GEMM_A16W16' -diff --git a/aiter/ops/gemm_op_a16w16.py b/aiter/ops/gemm_op_a16w16.py ---- a/aiter/ops/gemm_op_a16w16.py -+++ b/aiter/ops/gemm_op_a16w16.py -@@ -37,6 +37,9 @@ - return torch.zeros(_SEMA_SHAPE, dtype=torch.uint32, device=device) - - -+_captured_semaphore_keepalive: list[Tensor] = [] -+ -+ - def get_semaphore_workspace(device: torch.device) -> Tensor: - """Return a per-(device, stream) zero-initialized semaphore workspace. - -@@ -52,7 +55,19 @@ - Workspace size is small (~4 KB) and stream count per process is typically - < 8, so the LRU cap of 64 leaves plenty of headroom before any in-flight - workspace risks being evicted. -+ -+ Under CUDA graph capture this returns a fresh workspace per launch instead -+ of the cached one: a captured graph bakes in the pointer and replays on a -+ stream other than the capture stream, so the cached counter can be left -+ non-zero and the reduction never fires. Allocating under capture also -+ records the zero-fill as a graph node, re-establishing the counter==0 entry -+ invariant on every replay. It is retained for the process lifetime because -+ aiter cannot observe when a graph dies. - """ -+ if torch.cuda.is_current_stream_capturing(): -+ w = torch.zeros(_SEMA_SHAPE, dtype=torch.uint32, device=device) -+ _captured_semaphore_keepalive.append(w) -+ return w - stream = torch.cuda.current_stream(device) - return _get_semaphore_workspace_keyed(device, stream.cuda_stream) - -DIFF_GEMM_A16W16 -apply_one "aiter/ops/gemm_op_a16w16.py" "is_current_stream_capturing" "$WS/GEMM_A16W16.diff" - -cat > "$WS/ROCM_AITER_MLA.diff" <<'DIFF_ROCM_AITER_MLA' -diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py ---- a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py -+++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py -@@ -26,7 +26,7 @@ - CommonAttentionMetadata, - MultipleOf, - ) --from vllm.v1.kv_cache_interface import AttentionSpec -+from vllm.v1.kv_cache_interface import AttentionSpec, is_quantized_kv_cache - - logger = init_logger(__name__) - -@@ -75,6 +75,50 @@ - except Exception: # noqa: BLE001 - return False - return True -+ -+ -+@functools.lru_cache(maxsize=1) -+def _gluon_mla_decode_supported() -> bool: -+ """The small-head Gluon MLA decode kernel only has a gfx950 (CDNA4) build. -+ -+ Its tiling needs ~160 KiB of LDS, which exceeds CDNA3's 64 KiB, so on -+ gfx942 there is no kernel to fall through to and selecting it asserts -+ (``mla_gluon requires gfx950``). Restrict Gluon decode to gfx950; other -+ archs use the asm persistent decode, which ``get_mla_padded_q`` makes -+ correct for any 1..15 heads. -+ """ -+ try: -+ from vllm.platforms.rocm import on_gfx950 -+ except Exception: # noqa: BLE001 -+ return False -+ return on_gfx950() -+ -+ -+def _aiter_mla_small_head_mode() -> str: -+ """Small-head (<16) MLA decode kernel selection. -+ -+ Controlled by ``VLLM_ROCM_AITER_MLA_ASM_PADDING``: -+ -+ - ``"auto"`` (default): let the arch decide -- divisor head counts keep the -+ Gluon decode where a build exists (gfx950), everything else (non-divisor -+ counts and all counts on gfx942) uses the padded persistent-scheduling -+ ASM decode. -+ - ``"gluon"``: prefer the Gluon path wherever a build exists. -+ - ``"asm"``: force the padded persistent-scheduling ASM decode. -+ -+ On gfx942 (no Gluon build) the ASM path is always used regardless of this -+ setting; ``"gluon"`` there falls back to ASM with a one-time warning. -+ """ -+ import vllm.envs as envs -+ -+ mode = (envs.VLLM_ROCM_AITER_MLA_ASM_PADDING or "auto").lower() -+ if mode == "gluon" and not _gluon_mla_decode_supported(): -+ logger.warning_once( -+ "VLLM_ROCM_AITER_MLA_ASM_PADDING=gluon requested, but this device " -+ "has no Gluon MLA decode build (Gluon requires gfx950); using the " -+ "padded persistent-scheduling ASM decode instead." -+ ) -+ return mode - - - class AiterMLABackend(MLACommonBackend): -@@ -134,6 +178,13 @@ - use_gluon_decode: bool = False - # Whether persistent MLA metadata was computed - has_persistent_metadata: bool = False -+ # Small-head multi-token verify: paged-KV metadata with one row per verify -+ # token holding that token's causal KV window, built in _build_decode so -+ # forward_mqa stays free of device->host syncs. -+ # flat_kv_indptr is [num_reqs * max_qo_len + 1]; flat_kv_indices is the -+ # whole persistent buffer, indexed through flat_kv_indptr. -+ flat_kv_indptr: torch.Tensor | None = None -+ flat_kv_indices: torch.Tensor | None = None - - - @dataclass -@@ -225,17 +276,17 @@ - self.compilation_config = vllm_config.compilation_config - self.decode_attn_out_dtype = vllm_config.model_config.dtype - -- # MTP/deepseek_mtp verification runs decode with qlen = num_spec + 1; -- # any other config (including no spec) stays at single-token decode. -- speculative_config = vllm_config.speculative_config -- if ( -- speculative_config is not None -- and speculative_config.method in ("mtp", "deepseek_mtp") -- and speculative_config.num_speculative_tokens is not None -- ): -- self._mtp_decode_qlen = int(speculative_config.num_speculative_tokens) + 1 -- else: -- self._mtp_decode_qlen = 1 -+ # Size the metadata from reorder_batch_threshold, the largest query -+ # length decode can be handed (MLACommonMetadataBuilder asserts -+ # max_query_len <= reorder_batch_threshold); it already accounts for the -+ # drafting scheme. A method-name whitelist instead leaves drafters not on -+ # it -- DSpark, the eagle family -- sized for qlen=1 while the router -+ # still admits up to 1 + 2 * num_spec. The persistent gate below then -+ # never opens and aiter indexes get_block_n_fp8[num_heads * qlen], a -+ # table holding only {8, 16, 24, 32, 48, 64, 128, 256, 384, 512}: at 16 -+ # heads every qlen in 5..7 and 9..15 is a KeyError, raised mid-run rather -+ # than at startup. -+ self._mtp_decode_qlen = self.reorder_batch_threshold or 1 - - # Store the kernel block size from the spec. When kernel_block_size=1 - # (no spec-dec), behavior is identical to the original. When > 1 -@@ -267,6 +318,74 @@ - self.paged_kv_indices = torch.zeros( - max_num_pages, dtype=torch.int32, device=device - ) -+ -+ # Small-head (< 16) multi-token verify expands each request's paged-KV -+ # range into one row per verify token, each holding that token's causal -+ # window. reorder_batch_threshold is the longest query block the decode -+ # path admits, so it bounds the row count per request. Sizing the -+ # buffers here keeps the expansion at fixed addresses, which is what -+ # lets the mla_gluon call in forward_mqa be captured in a full CUDA -+ # graph. -+ # -+ # The flatten is selected by the *impl's* per-layer query head count, so -+ # the buffers are reserved for any multi-token decode block this group -+ # can admit rather than from this builder's own num_heads, which is not -+ # required to agree with it. That reserves them for >= 16-head -+ # deployments too, where mla_decode_fwd serves the block and never reads -+ # them. -+ self._flat_max_qo_len = max(1, int(self.reorder_batch_threshold or 1)) -+ self._flat_kv_enabled = self._flat_max_qo_len > 1 -+ if self._flat_kv_enabled: -+ # The rows write at most max_qo_len times the sum of the batch's -+ # sequence lengths. max_num_pages bounds that sum by assuming every -+ # request is max_model_len long at the same time, which needs many -+ # times more entries than the KV cache can hold. Without prefix -+ # caching no two requests share a slot, so the pool's own token -+ # capacity is the real bound. -+ # -+ # cache_config.kv_cache_size_tokens is that capacity, -+ # max_concurrency * max_model_len, and it is a genuine upper bound -+ # on the sum even though every group draws block ids from one -+ # shared pool: a group's block count for a request of L tokens is -+ # either constant in L or concave in L through the origin, so it is -+ # never below L / max_model_len of what a full-length request -+ # takes. Summing that over the pool gives exactly this figure. -+ # num_gpu_blocks * block_size counts only this group's slots and so -+ # overstates the bound on hybrid layouts, where the other groups' -+ # blocks come out of the same pool; it is kept as the fallback for -+ # engines that have not published the group-aware capacity. -+ cache_config = vllm_config.cache_config -+ flat_pages = max_num_pages -+ if not cache_config.enable_prefix_caching: -+ kv_capacity = cache_config.kv_cache_size_tokens -+ if not kv_capacity and cache_config.num_gpu_blocks: -+ kv_capacity = ( -+ int(cache_config.num_gpu_blocks) * self.kernel_block_size -+ ) -+ if kv_capacity: -+ flat_pages = min(flat_pages, int(kv_capacity)) -+ self.flat_kv_indptr = torch.zeros( -+ max_num_reqs * self._flat_max_qo_len + 1, -+ dtype=torch.int32, -+ device=device, -+ ) -+ self.flat_kv_indices = torch.zeros( -+ flat_pages * self._flat_max_qo_len, dtype=torch.int32, device=device -+ ) -+ # [0, 1, ..., max_qo_len - 1]. Added to a request's context length -+ # this gives each verify row its own causal KV bound; materialised -+ # once so the per-step build allocates nothing. -+ self._flat_causal_offsets = torch.arange( -+ self._flat_max_qo_len, dtype=torch.int32, device=device -+ ) -+ logger.info( -+ "AITER MLA small-head verify buffers allocated " -+ "(max_qo_len=%d, pages=%d of %d, %.1f MiB)", -+ self._flat_max_qo_len, -+ flat_pages, -+ max_num_pages, -+ self.flat_kv_indices.numel() * 4 / (1024 * 1024), -+ ) - - from aiter import dtypes, get_mla_metadata_info_v1 - -@@ -283,6 +402,9 @@ - torch.float16: dtypes.fp16, - torch.bfloat16: dtypes.bf16, - }[kv_cache_spec.dtype] -+ # _build_decode needs the cache dtype to pick the decode kernel; keep -+ # the normalized string instead of dropping it at the end of __init__. -+ self._kv_cache_dtype_str = kv_cache_dtype_str - # MLAAttention quantizes decode Q to FP8 before calling this backend - # whenever the KV cache is FP8 and supports_quant_query_input is true. - q_dtype = ( -@@ -329,9 +451,12 @@ - device=device, - ) - -- # FP8 MLA prefill (kn_mla_reduce_v1) only supports 16-aligned heads. -- self._fp8_prefill_enabled = ( -- _fp8_mla_prefill_supported() and self.num_heads % 16 == 0 -+ # FP8 MLA prefill (kn_mla_reduce_v1) only supports 16-aligned heads, and -+ # only runs when the KV cache is FP8 (otherwise the bf16 path is used and -+ # the PS workspace must not be reserved). -+ self._fp8_prefill_enabled = _fp8_mla_prefill_supported() and ( -+ kv_cache_dtype_str == "fp8" -+ and (self.num_heads % 16 == 0 or 0 < self.num_heads < 16) - ) - if self._fp8_prefill_enabled: - max_prefill_qlen = min( -@@ -387,7 +512,11 @@ - - # After kv_b_proj decompression, K has num_heads heads (same as Q). - # So gqa_ratio=1 and num_head_k=num_heads for the PS kernel. -- num_head_k = self.num_heads -+ # Non-divisor head counts (e.g. K3's 12/rank at TP8) are padded to 16 in -+ # _mla_fp8_prefill_attn; build the PS metadata for the padded head count so -+ # the work/reduce maps match. This also lowers the partial-tile count: -+ # gcd(16, cu_num=256)=16 (~960 tiles) vs gcd(12,256)=4 (~4032), saving ~6 GiB. -+ num_head_k = max(16, self.num_heads) - v_head_dim = self.mla_dims.v_head_dim - # gqa_ratio = 1 - # qlen_granularity = _FP8_PREFILL_TILE_Q // max(gqa_ratio, 1) -@@ -481,7 +610,11 @@ - kv_indptr_cpu = qo_indptr_cpu.clone() - seq_lens_cpu = (qo_indptr_cpu[1:] - qo_indptr_cpu[:-1]).to(torch.int32) - -- num_head_k = self.num_heads -+ # Non-divisor head counts (e.g. K3's 12/rank at TP8) are padded to 16 in -+ # _mla_fp8_prefill_attn; build the PS metadata for the padded head count so -+ # the work/reduce maps match. This also lowers the partial-tile count: -+ # gcd(16, cu_num=256)=16 (~960 tiles) vs gcd(12,256)=4 (~4032), saving ~6 GiB. -+ num_head_k = max(16, self.num_heads) - # gqa_ratio = 1 - # qhead_granularity = max(gqa_ratio, 1) - # qlen_granularity = _FP8_PREFILL_TILE_Q // qhead_granularity -@@ -580,7 +713,7 @@ - ] - ) - use_gluon_decode = AiterMLAHelper.use_gluon_decode( -- self.num_heads, int(max_qo_len) -+ self.num_heads, int(max_qo_len), self._kv_cache_dtype_str - ) - - if self.compilation_config.cudagraph_mode.has_full_cudagraphs(): -@@ -596,9 +729,9 @@ - block_table_tensor, - block_table_tensor.stride(0), - paged_kv_indptr, -- seq_lens_for_kernel, - KERNEL_BLOCK_SIZE=self.kernel_block_size, - BLOCK_SIZE=1024, -+ QLEN=1, - ) - paged_kv_indices = self.paged_kv_indices - -@@ -650,12 +783,24 @@ - qo_indptr = query_start_loc_device[: 1 + num_kernel_reqs] - - # Pass persistent metadata for every uniform decode we sized buffers for -- # (normal qlen==1 through MTP verification qlen==K): the fp8 nhead=32 fold -- # path breaks without it. qlen>K falls back to kernel-internal metadata. -- # Small-head (<16) decode takes the Gluon paths and never consumes it. -+ # (qlen==1 through verification qlen==K); qlen>K falls back to -+ # kernel-internal metadata. Only the asm decode consumes the schedule, so -+ # gate on the routing and not on the raw head count: a non-divisor rank is -+ # padded to 16 and runs the same asm kernels as a native 16-head rank, yet -+ # `num_heads >= 16` reads as False for it and denies it the schedule. The -+ # kernel then falls back on its internal metadata, which bf16 tolerates -+ # and fp8 does not, and which the fp8 fold path rejects once qlen > 4: -+ # -+ # asm_mla.cu:903 mla_decode_stage1_asm_fwd: only support gqa_ratio=16 -+ # fp8 mla decoding with qo_len <= 4 and qo_len > 4 in persistent mode - has_persistent_metadata = False - use_persistent_metadata = ( -- self.num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS -+ not AiterMLAHelper.use_gluon_decode( -+ self.num_heads, max_qo_len, self._kv_cache_dtype_str -+ ) -+ and not AiterMLAHelper.use_gluon_verify( -+ self.num_heads, max_qo_len, self._kv_cache_dtype_str -+ ) - and max_qo_len >= 1 - and max_qo_len <= self._mtp_decode_qlen - ) -@@ -688,6 +833,79 @@ - ) - has_persistent_metadata = True - -+ # Small-head multi-token verify: build the per-verify-token causal -+ # paged-KV view here, once per step, instead of once per MLA layer in -+ # forward_mqa. That removes four device->host syncs per layer (an -+ # .item(), two tensor-driven repeat_interleave calls and a .min()) plus -+ # a data-dependent allocation, all of which abort HIP graph capture. -+ flat_kv_indptr = None -+ flat_kv_indices = None -+ min_kv_seq_len = 1 -+ if self._flat_kv_enabled and max_qo_len > 1: -+ qlen = int(max_qo_len) -+ assert qlen <= self._flat_max_qo_len, ( -+ f"verify block {qlen} exceeds the reserved maximum " -+ f"{self._flat_max_qo_len}" -+ ) -+ num_rows = num_kernel_reqs * qlen -+ # Row r * qlen + t is request r's verify token t. seq_lens counts -+ # the tokens scheduled in this step, so a request's KV range already -+ # spans its whole verify block and context_r = seq_len_r - qlen. -+ # Causal masking lets token t attend to KV positions -+ # [0, context_r + t], i.e. seq_len_r - (qlen - 1) + t entries, so -+ # only the last row of a block may see the full range. Rows clamp to -+ # zero for cudagraph padding requests, whose seq_len is 0. -+ per_req_len = paged_kv_indptr[1:] - paged_kv_indptr[:-1] -+ row_len = ( -+ ( -+ per_req_len.unsqueeze(1) -+ - (qlen - 1) -+ + self._flat_causal_offsets[:qlen] -+ ) -+ .clamp_(min=0) -+ .flatten() -+ ) -+ # Element 0 stays zero from the initial torch.zeros; assigning a -+ # Python scalar to it would be a blocking host->device copy. -+ self.flat_kv_indptr[1 : num_rows + 1].copy_( -+ row_len.cumsum(dim=0, dtype=torch.int32), non_blocking=True -+ ) -+ # A replayed cudagraph reads seq_info out to its captured row count, -+ # which can exceed num_rows. Repeating the final offset rather than -+ # zeroing makes every such row report length 0 instead of a large -+ # negative one, the same reason paged_kv_indptr's tail above is -+ # filled with its last entry. -+ self.flat_kv_indptr[num_rows + 1 :].fill_(self.flat_kv_indptr[num_rows]) -+ flat_kv_indptr = self.flat_kv_indptr[: num_rows + 1] -+ # One device->host read serves both uses below; a sync is legal here -+ # because the builder runs outside the captured region. Gluon turns -+ # min_kv_seq_len into its split count, so it has to be the shortest -+ # row actually submitted, not the shortest per-request length those -+ # rows were cut from. -+ min_kv_seq_len, total_entries = torch.stack( -+ (row_len.min(), self.flat_kv_indptr[num_rows]) -+ ).tolist() -+ # flat_kv_indices is reserved from the KV pool's token capacity, -+ # which bounds this sum. Check it rather than let a bound that is -+ # wrong for some future layout corrupt memory silently. -+ assert total_entries <= self.flat_kv_indices.numel(), ( -+ f"verify KV view needs {total_entries} entries but only " -+ f"{self.flat_kv_indices.numel()} are reserved" -+ ) -+ # No need to clear flat_kv_indices: the kernel writes exactly the -+ # [flat_kv_indptr[row], flat_kv_indptr[row + 1]) range that -+ # mla_gluon reads back for that row. -+ _expand_page_indices_kernel[(num_rows,)]( -+ self.flat_kv_indices, -+ block_table_tensor, -+ block_table_tensor.stride(0), -+ flat_kv_indptr, -+ KERNEL_BLOCK_SIZE=self.kernel_block_size, -+ BLOCK_SIZE=1024, -+ QLEN=qlen, -+ ) -+ flat_kv_indices = self.flat_kv_indices -+ - attn_metadata = AiterMLADecodeMetadata( - block_table=block_table_tensor, - seq_lens=seq_lens_for_kernel, -@@ -697,9 +915,12 @@ - qo_indptr=qo_indptr, - dcp_tot_seq_lens=dcp_tot_seq_lens_device, - max_qo_len=max_qo_len, -+ min_kv_seq_len=min_kv_seq_len, - use_gluon_decode=use_gluon_decode, - attn_out_dtype=self.decode_attn_out_dtype, - has_persistent_metadata=has_persistent_metadata, -+ flat_kv_indptr=flat_kv_indptr, -+ flat_kv_indices=flat_kv_indices, - ) - - return attn_metadata -@@ -734,9 +955,9 @@ - block_table, - block_table_stride, - cu_num_tokens, -- seq_lens, - KERNEL_BLOCK_SIZE: tl.constexpr, - BLOCK_SIZE: tl.constexpr, -+ QLEN: tl.constexpr, - ): - """Expand block table entries into per-token flat page indices. - -@@ -750,11 +971,19 @@ - - When KERNEL_BLOCK_SIZE=K: block table entry b (covering K tokens) - is expanded to flat indices b*K, b*K+1, ..., b*K+(K-1). -+ -+ QLEN is the number of output rows per request: 1 for ordinary decode, and -+ the verify block length for the small-head multi-token verify expansion, -+ where output row ``r * QLEN + t`` is request ``r``'s verify token ``t`` and -+ takes the first ``cu_num_tokens[row + 1] - cu_num_tokens[row]`` tokens of -+ that request -- its causal window, since block_table lists a request's -+ blocks in ascending position order. - """ -- req_idx = tl.program_id(0) -+ row_idx = tl.program_id(0) -+ req_idx = row_idx // QLEN - row_ptr = block_table + req_idx * block_table_stride -- start_idx = tl.load(cu_num_tokens + req_idx) -- num_tokens = tl.load(seq_lens + req_idx) -+ start_idx = tl.load(cu_num_tokens + row_idx) -+ num_tokens = tl.load(cu_num_tokens + row_idx + 1) - start_idx - - offset = tl.arange(0, BLOCK_SIZE) - for i in tl.range(0, num_tokens, BLOCK_SIZE): -@@ -781,8 +1010,11 @@ - - class AiterMLAHelper: - """ -- AITER MLA implementation requires num_heads >= 16. If num_heads < 16 and -- 16 % num_heads == 0, we can pad q to 16 heads; otherwise AITER has to fail. -+ AITER MLA persistent (asm) decode requires num_heads >= 16. Head counts -+ < 16 are padded up to exactly 16: divisors of 16 by repeat_interleave, -+ other counts (e.g. 12 heads/rank at TP8, 6 at TP16) by tiling the query -+ heads and slicing to 16. Non-divisor padded decodes take the asm path; -+ divisors and max_qo_len > 1 small-head verify still use Gluon. - """ - - _AITER_MIN_MLA_HEADS: Final = 16 -@@ -791,8 +1023,9 @@ - @staticmethod - def check_num_heads_validity(num_heads: int): - assert AiterMLAHelper.is_valid_num_heads(num_heads), ( -- "ROCM AITER MLA requires 1-15 heads for Gluon decode or a multiple " -- f"of 16 heads for persistent decode, but got {num_heads}.\n" -+ "ROCM AITER MLA requires 1-15 heads (padded to 16 for asm " -+ "persistent decode; exact divisors of 16 may keep Gluon) or a " -+ f"multiple of 16 heads, but got {num_heads}.\n" - f"Try adjusting tensor_parallel_size value." - ) - -@@ -809,25 +1042,87 @@ - - @staticmethod - def get_mla_padded_q(num_heads: int, q: torch.Tensor) -> torch.Tensor: -- return ( -- q -- if num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS -- else q.repeat_interleave( -- AiterMLAHelper._AITER_MIN_MLA_HEADS // num_heads, dim=1 -- ) -- ) -+ m = AiterMLAHelper._AITER_MIN_MLA_HEADS -+ if num_heads >= m: -+ return q -+ if m % num_heads == 0: -+ return q.repeat_interleave(m // num_heads, dim=1) -+ # Non-divisor head counts (e.g. 12 heads/rank at TP8, 6 at TP16) cannot -+ # be padded by repeat_interleave. Tile the query heads and slice to -+ # exactly m; this reaches m for any 0 < num_heads < m (unlike a single -+ # append, which under-pads when num_heads < m - num_heads). MLA -+ # attention is independent per query head over the shared KV, so the -+ # padding heads cannot affect heads [0:num_heads]; they are sliced back -+ # off in get_mla_unpadded_o. -+ reps = -(-m // num_heads) # ceil(m / num_heads) -+ # Slicing a tiled tensor down to m yields a non-contiguous view whenever -+ # reps * num_heads > m (the common case: TP8 12->24->16, TP16 6->18->16). -+ # The asm persistent decode reads q as a packed [tokens, m, head_dim] -+ # buffer, so materialize a contiguous copy. No-op when already contiguous. -+ return q.repeat(1, reps, 1)[:, :m, :].contiguous() - - @staticmethod - def get_mla_unpadded_o(num_heads: int, o: torch.Tensor) -> torch.Tensor: -- return ( -- o -- if num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS -- else o[:, :: AiterMLAHelper._AITER_MIN_MLA_HEADS // num_heads, :] -- ) -+ m = AiterMLAHelper._AITER_MIN_MLA_HEADS -+ if num_heads >= m: -+ return o -+ if m % num_heads == 0: -+ return o[:, :: m // num_heads, :] -+ # Undo the tile-padding from get_mla_padded_q: the real heads are the -+ # first num_heads. -+ return o[:, :num_heads, :] - - @staticmethod -- def use_gluon_decode(num_heads: int, max_qo_len: int) -> bool: -- return num_heads < AiterMLAHelper._AITER_MIN_MLA_HEADS and max_qo_len == 1 -+ def use_gluon_decode(num_heads: int, max_qo_len: int, kv_cache_dtype: str) -> bool: -+ # Small-head (<16) single-token decode can use either the Gluon kernel -+ # or the padded asm persistent decode, selected by -+ # VLLM_ROCM_AITER_MLA_ASM_PADDING (see _aiter_mla_small_head_mode) and -+ # the arch: Gluon only has a gfx950 build. In "auto" (default) mode -+ # divisor counts keep Gluon on gfx950 and everything else -- non-divisor -+ # counts (e.g. 12 heads/rank at TP8) and all counts on gfx942 -- takes -+ # the asm path, which get_mla_padded_q pads to exactly 16. -+ m = AiterMLAHelper._AITER_MIN_MLA_HEADS -+ if num_heads >= m or max_qo_len != 1: -+ return False -+ # Gluon has exactly one fp8-KV regime, bh16bn128. It is a bf16-query -+ # kernel that upcasts the cache in registers with a hardcoded scale of -+ # 1.0, and it asserts batch_size == 1, so it cannot serve a real decode -+ # batch at any head count. A quantized cache always goes to the asm -+ # decode, which ships true fp8 kernels for gqa=16 -+ # (mla_a8w8_qh16_qseqlen*_gqaratio16*.co). This precedes the mode knob: -+ # an explicit "gluon" request under fp8 would assert immediately. -+ if is_quantized_kv_cache(kv_cache_dtype): -+ return False -+ mode = _aiter_mla_small_head_mode() -+ if mode == "asm": -+ return False -+ gluon_supported = _gluon_mla_decode_supported() -+ if mode == "gluon": -+ return gluon_supported -+ return m % num_heads == 0 and gluon_supported -+ -+ @staticmethod -+ def use_gluon_verify(num_heads: int, max_qo_len: int, kv_cache_dtype: str) -> bool: -+ """Whether a small-head multi-token verify is flattened onto Gluon. -+ -+ The bf16 asm kernels have no gqa < 16, qseqlen > 1 entry, so a small-head -+ verify is flattened into per-token qseqlen=1 Gluon decodes. fp8 does have -+ one, reached by the q-row fold (16 heads x qlen 8 folds onto the -+ nhead=32, qseqlen=4 kernel, which ships in the package), and must not -+ come here: the flatten hands Gluon a batch of exactly the size that its -+ fp8 regime asserts against. -+ -+ This lives next to use_gluon_decode rather than inline in forward_mqa so -+ that the builder, which has to know whether the asm decode will run, sees -+ the same answer the impl acts on. -+ """ -+ if num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS or max_qo_len <= 1: -+ return False -+ # HYBRID: small-head multi-token verify always uses the Gluon flatten, -+ # independent of kv dtype and VLLM_ROCM_AITER_MLA_ASM_PADDING. fp8 KV is -+ # served by the batch<=256 + fp8-query-dequant mla_gluon relaxation; the -+ # asm fp8 q-row-fold verify faults on gfx950 (HSA 0x1016 in DSpark). -+ return _gluon_mla_decode_supported() - - - class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]): -@@ -873,10 +1168,13 @@ - self.flash_attn_varlen_func = flash_attn_varlen_func - - # FP8 MLA prefill kernel imports (lazy, only when enabled). -- # Auto-enabled on gfx950 when AITER ships the kernels. -- # FP8 MLA prefill (kn_mla_reduce_v1) only supports 16-aligned heads. -- self._fp8_prefill_enabled = ( -- _fp8_mla_prefill_supported() and self.num_heads % 16 == 0 -+ # Auto-enabled on gfx950 when AITER ships the kernels. Only runs when the -+ # KV cache is FP8, and supports non-divisor small head counts via pad-to-16. -+ from vllm.utils.torch_utils import is_quantized_kv_cache -+ -+ self._fp8_prefill_enabled = _fp8_mla_prefill_supported() and ( -+ is_quantized_kv_cache(kv_cache_dtype) -+ and (self.num_heads % 16 == 0 or 0 < self.num_heads < 16) - ) - if self._fp8_prefill_enabled: - from aiter import mla_prefill_ps_asm_fwd, mla_reduce_v1 -@@ -919,7 +1217,19 @@ - - fp8_dtype = current_platform.fp8_dtype() - total_q = q.shape[0] -- nhead = self.num_heads -+ # PS asm prefill + mla_reduce_v1 require 16-aligned heads and the PS -+ # metadata is built for max(16, num_heads). For non-divisor small head -+ # counts (K3 = 12/rank at TP8) replicate-pad q/k/v to 16 — MLA attention -+ # is independent per query head over the shared KV, so the padding heads -+ # cannot affect the real ones (exact, same as the decode path) — then -+ # slice the output back to the real head count. -+ _real_nhead = self.num_heads -+ _pad16 = _real_nhead < 16 -+ if _pad16: -+ q = AiterMLAHelper.get_mla_padded_q(_real_nhead, q) -+ k = AiterMLAHelper.get_mla_padded_q(_real_nhead, k) -+ v = AiterMLAHelper.get_mla_padded_q(_real_nhead, v) -+ nhead = 16 if _pad16 else self.num_heads - v_head_dim = self.v_head_dim - tile_q = _FP8_PREFILL_TILE_Q - -@@ -946,7 +1256,13 @@ - # Reuse the caller's output buffer to skip the per-call alloc + copy. - # The ASM and reduce kernels both write to a [total_q, nhead, v_head_dim] - # view, which aliases the [total_q, nhead * v_head_dim] storage of out. -- out_3d = out.view(total_q, nhead, v_head_dim) -+ if _pad16: -+ # Padded heads can't alias the real-head `out` storage; use scratch. -+ out_3d = torch.empty( -+ total_q, nhead, v_head_dim, dtype=out.dtype, device=out.device -+ ) -+ else: -+ out_3d = out.view(total_q, nhead, v_head_dim) - - # Per-call scratch (logits, attn_lse, final_lse) is served from the - # workspace manager so allocator churn in the prefill hot path is -@@ -993,6 +1309,11 @@ - final_lse, - ) - -+ if _pad16: -+ out.view(total_q, _real_nhead, v_head_dim).copy_( -+ out_3d[:, :_real_nhead, :] -+ ) -+ - def forward_mha( - self, - q: torch.Tensor, -@@ -1113,11 +1434,12 @@ - # target is checking draft tokens, so position t must not see t+1 -- - # and attention rows are independent, so giving row t the KV range - # [0, context + t] is exactly causal multi-token attention. -- if ( -- self.num_heads < AiterMLAHelper._AITER_MIN_MLA_HEADS -- and int(decode.max_qo_len) > 1 -+ # Arch, mode and dtype gating all live in use_gluon_verify, so that the -+ # builder -- which has to know whether the asm decode will run -- sees -+ # the same answer as this branch. -+ if AiterMLAHelper.use_gluon_verify( -+ self.num_heads, int(decode.max_qo_len), self.kv_cache_dtype - ): -- qlen = int(decode.max_qo_len) - if type(q) is tuple: - q_nope, q_pe = q - else: -@@ -1133,56 +1455,35 @@ - device=q_nope.device, - ) - kv_buffer = kv_c_and_k_pe_cache.reshape(-1, kv_c_and_k_pe_cache.shape[-1]) -- # Expand per-request paged-KV to per-verify-token. Row r*qlen+t is -- # request r's verify token t, and seq_lens counts the tokens -- # scheduled in this step, so a request's KV range already spans its -- # whole verify block and context_r = seq_len_r - qlen. Token t may -- # attend to [0, context_r + t], i.e. seq_len_r - (qlen - 1) + t -- # entries. paged_kv_indices lists a request's pages in ascending -- # position order, so each row's causal window is a prefix of that -- # request's slice and only the row length changes. Rows clamp to -- # zero for cudagraph padding requests, whose seq_len is 0. Fully -- # vectorized (no host loop). -- old_indptr = decode.paged_kv_indptr -- per_req_len = old_indptr[1:] - old_indptr[:-1] -- dev = q_nope.device -- row_req = torch.arange(per_req_len.shape[0], device=dev).repeat_interleave( -- qlen -- ) -- row_len = ( -- ( -- per_req_len.unsqueeze(1) -- - (qlen - 1) -- + torch.arange(qlen, device=dev, dtype=per_req_len.dtype) -- ) -- .clamp_(min=0) -- .flatten() -- ) -- new_indptr = torch.cat([old_indptr.new_zeros(1), row_len.cumsum(0)]).to( -- torch.int32 -- ) -- total = int(new_indptr[-1].item()) -- within = torch.arange(total, device=dev, dtype=torch.int64) - new_indptr[ -- :-1 -- ].to(torch.int64).repeat_interleave(row_len) -- src = ( -- old_indptr[row_req].to(torch.int64).repeat_interleave(row_len) + within -- ) -- new_indices = decode.paged_kv_indices[src] -+ # The per-verify-token view -- row r*qlen+t reads request r's -+ # committed prefix plus verify tokens 0..t, i.e. its causal window -- -+ # is built once per step in _build_decode, where device->host syncs -+ # are legal. Reading it back here keeps this path free of the syncs -+ # that previously aborted HIP graph capture. -+ assert decode.flat_kv_indptr is not None -+ assert decode.flat_kv_indices is not None -+ # A non-causal block would need the untruncated range instead, and -+ # cannot arrive here: this builder leaves -+ # supports_non_causal_multi_token_decode False, so -+ # MLACommonMetadataBuilder.build rejects causal=False before -+ # _build_decode ever runs. -+ assert attn_metadata.causal, ( -+ "AITER MLA small-head verify flatten is causal-only" -+ ) - mla_gluon = _get_mla_gluon() - mla_gluon( - q_nope=q_nope, - q_pe=q_pe, - kv_c=kv_buffer, - o=o, -- page_table=new_indices, -- seq_info=new_indptr, -+ page_table=decode.flat_kv_indices, -+ seq_info=decode.flat_kv_indptr, - sm_scale=self.scale, - k_pe=None, - kv_pe_offset=self.kv_lora_rank, - use_2d_view=False, - kv_scale=1.0, -- min_kv_seq_len=int(row_len.min()), -+ min_kv_seq_len=decode.min_kv_seq_len, - ) - return o, None - -DIFF_ROCM_AITER_MLA -apply_one "vllm/v1/attention/backends/mla/rocm_aiter_mla.py" "flat_kv_indices" "$WS/ROCM_AITER_MLA.diff" - -# --- DSpark PS verify: supersede the HYBRID gluon-flatten verify ------------- -# Two edits on the file the diff above just produced (HYBRID). Done as exact -# string replacements (not a context diff) so whitespace/line-drift can't break -# it, and idempotent via the "Local DSpark PS extension" guard. The base marker -# above was changed to "flat_kv_indices" (untouched here) so re-runs still skip. -# (a) use_gluon_verify -> False for fp8 KV: the small-head multi-token verify -# is no longer swallowed by the Gluon flatten and falls through to the ASM -# persistent (PS) decode (aiter #4521 qseqlen4 cprr kernels). -# (b) size _mtp_decode_qlen for DSpark (1 + num_spec) so the PS gate opens. -python - "$ROOT/vllm/v1/attention/backends/mla/rocm_aiter_mla.py" <<'PYDSPARK' -import ast, sys -F = sys.argv[1] -src = open(F).read() -if "Local DSpark PS extension" in src: - print(" rocm_aiter_mla.py (DSpark PS): already present (skip)"); sys.exit(0) -OLD1 = " self._mtp_decode_qlen = self.reorder_batch_threshold or 1\n" -NEW1 = ( - OLD1 - + " # Local DSpark PS extension: reorder_batch_threshold's method\n" - + " # whitelist does not size DSpark, leaving its verify (qlen =\n" - + " # 1 + num_spec) at 1 so the persistent gate below never opens. Size\n" - + " # it explicitly so the ASM PS fp8 verify (qh16/qh32 qseqlen4 cprr\n" - + " # kernels, aiter #4521) is reachable.\n" - + " _spec = vllm_config.speculative_config\n" - + " if _spec is not None and (\n" - + " getattr(_spec, \"use_dspark\", False)\n" - + " or getattr(_spec, \"method\", None) == \"dspark\"\n" - + " ):\n" - + " self._mtp_decode_qlen = max(\n" - + " self._mtp_decode_qlen, 1 + int(_spec.num_speculative_tokens or 0)\n" - + " )\n" -) -OLD2 = ( - " # HYBRID: small-head multi-token verify always uses the Gluon flatten,\n" - " # independent of kv dtype and VLLM_ROCM_AITER_MLA_ASM_PADDING. fp8 KV is\n" - " # served by the batch<=256 + fp8-query-dequant mla_gluon relaxation; the\n" - " # asm fp8 q-row-fold verify faults on gfx950 (HSA 0x1016 in DSpark).\n" - " return _gluon_mla_decode_supported()\n" -) -NEW2 = ( - " # Local DSpark PS extension: with aiter #4521 the asm fp8 q-row-fold\n" - " # verify (qh16/qh32 qseqlen4 cprr kernels) works on gfx950, so an fp8\n" - " # KV small-head multi-token verify must NOT be swallowed by the Gluon\n" - " # flatten -- let it fall through to the ASM persistent (PS) path.\n" - " if is_quantized_kv_cache(kv_cache_dtype):\n" - " return False\n" - " return _gluon_mla_decode_supported()\n" -) -for tag, OLD in (("mtp_qlen sizing", OLD1), ("use_gluon_verify", OLD2)): - if src.count(OLD) != 1: - print(f" rocm_aiter_mla.py (DSpark PS): ABORT {tag} (found {src.count(OLD)})") - sys.exit(2) -src = src.replace(OLD1, NEW1, 1).replace(OLD2, NEW2, 1) -ast.parse(src) -open(F, "w").write(src) -print(" rocm_aiter_mla.py (DSpark PS): APPLIED") -PYDSPARK - -cat > "$WS/TRITON_MLA.diff" <<'DIFF_TRITON_MLA' -diff --git a/vllm/v1/attention/backends/mla/triton_mla.py b/vllm/v1/attention/backends/mla/triton_mla.py ---- a/vllm/v1/attention/backends/mla/triton_mla.py -+++ b/vllm/v1/attention/backends/mla/triton_mla.py -@@ -6,6 +6,7 @@ - import torch - - import vllm.envs as envs -+from vllm.config import VllmConfig - from vllm.config.cache import CacheDType - from vllm.logger import init_logger - from vllm.model_executor.layers.attention.mla_attention import ( -@@ -25,6 +26,7 @@ - MultipleOf, - ) - from vllm.v1.attention.ops.triton_decode_attention import decode_attention_fwd -+from vllm.v1.kv_cache_interface import KVCacheSpec - from vllm.v1.worker.workspace import ( - current_workspace_manager, - is_workspace_manager_initialized, -@@ -54,6 +56,34 @@ - # Non-causal DSpark block is flattened to one decode row per query token in - # forward_mqa, so no intra-block causal masking is required. - supports_non_causal_multi_token_decode: ClassVar[bool] = True -+ -+ @classmethod -+ def get_cudagraph_support( -+ cls, -+ vllm_config: VllmConfig, -+ kv_cache_spec: KVCacheSpec, -+ ) -> AttentionCGSupport: -+ """Report UNIFORM_BATCH where a non-causal multi-token block is served. -+ -+ ``_cudagraph_support`` is a class constant, so serving the DSpark -+ draft's (1 + num_spec) block through the decode path reports -+ UNIFORM_SINGLE_TOKEN_DECODE and, because the engine takes the minimum -+ over all attention groups, downgrades the *whole* engine off full -+ cudagraphs. ``forward_mqa`` flattens that block with -+ ``repeat_interleave`` on a Python int and performs no device->host -+ sync, so it does satisfy the UNIFORM_BATCH contract. -+ -+ ``non_causal_multi_token_decode`` is a KV-cache-group property, not a -+ per-layer one: ``MLAAttentionSpec.merge`` ORs it over every layer in -+ the group, so a group holding both a draft and its target reports it -+ for both. That is the same predicate ``__init__`` below already uses to -+ raise ``reorder_batch_threshold``, so the two stay consistent, but it -+ does mean this lifts a causal target sharing the draft's KV cache group -+ as well. -+ """ -+ if getattr(kv_cache_spec, "non_causal_multi_token_decode", False): -+ return AttentionCGSupport.UNIFORM_BATCH -+ return cls._cudagraph_support - - def __init__(self, kv_cache_spec, layer_names, vllm_config, device): - super().__init__(kv_cache_spec, layer_names, vllm_config, device) -DIFF_TRITON_MLA -apply_one "vllm/v1/attention/backends/mla/triton_mla.py" "get_cudagraph_support" "$WS/TRITON_MLA.diff" - -cat > "$WS/GPU_WORKER.diff" <<'DIFF_GPU_WORKER' -diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py ---- a/vllm/v1/worker/gpu_worker.py -+++ b/vllm/v1/worker/gpu_worker.py -@@ -64,6 +64,7 @@ - from vllm.utils.mem_constants import GiB_bytes - from vllm.utils.mem_utils import MemorySnapshot, format_gib, memory_profiling - from vllm.utils.torch_utils import set_random_seed -+from vllm.v1.core.kv_cache_utils import get_kv_cache_capacity - from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput - from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec - from vllm.v1.outputs import ( -@@ -652,6 +653,19 @@ - - # Update local config with adjusted num blocks after profiling, - # so that it's available to the warmup stage. -+ # num_gpu_blocks * block_size is not the pool's token capacity when a -+ # request occupies more than one KV cache group, which is why -+ # kv_cache_size_tokens exists. It is only ever filled in by the engine -+ # core and the front end, so the worker's copy stays None and anything -+ # sizing a buffer off the KV pool during warmup -- the AITER MLA verify -+ # view, for one -- silently falls back to a far looser bound. Fill it in -+ # here too; get_kv_cache_capacity is documented to give the same answer -+ # for the worker's config as for the scheduler's. -+ if kv_cache_config.kv_cache_groups: -+ ( -+ self.cache_config.kv_cache_size_tokens, -+ self.cache_config.kv_cache_max_concurrency, -+ ) = get_kv_cache_capacity(self.vllm_config, kv_cache_config) - self.cache_config.num_gpu_blocks = kv_cache_config.num_blocks - - # Init kv cache connector here, because it requires -DIFF_GPU_WORKER -apply_one "vllm/v1/worker/gpu_worker.py" "import get_kv_cache_capacity" "$WS/GPU_WORKER.diff" - -cat > "$WS/VLLM_ENVS.diff" <<'DIFF_VLLM_ENVS' -diff --git a/vllm/envs.py b/vllm/envs.py ---- a/vllm/envs.py -+++ b/vllm/envs.py -@@ -133,6 +133,7 @@ - VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4: bool = False - VLLM_ROCM_USE_AITER_RMSNORM: bool = True - VLLM_ROCM_USE_AITER_MLA: bool = True -+ VLLM_ROCM_AITER_MLA_ASM_PADDING: Literal["auto", "gluon", "asm"] = "auto" - VLLM_ROCM_USE_AITER_MHA: bool = True - VLLM_ROCM_USE_AITER_FP4_ASM_GEMM: bool = False - VLLM_ROCM_USE_AITER_TRITON_ROPE: bool = False -@@ -1236,6 +1237,20 @@ - "VLLM_ROCM_USE_AITER_MLA": lambda: ( - os.getenv("VLLM_ROCM_USE_AITER_MLA", "True").lower() in ("true", "1") - ), -+ # Small-head (<16) AITER MLA decode kernel selection. Small head counts -+ # (e.g. Kimi-K3: 12 heads/rank at TP8, 6 at TP16) can decode either through -+ # the Gluon small-head kernel or through the padded persistent-scheduling -+ # (PS) ASM kernel. "auto" (default) keeps Gluon for head counts that divide -+ # 16 where a Gluon build exists (gfx950/CDNA4) and otherwise uses the padded -+ # PS ASM decode; "gluon" forces the Gluon path wherever a build exists; -+ # "asm" forces the padded PS ASM decode. On gfx942/CDNA3 there is no Gluon -+ # build, so the ASM path is always used regardless of this setting. -+ "VLLM_ROCM_AITER_MLA_ASM_PADDING": env_with_choices( -+ "VLLM_ROCM_AITER_MLA_ASM_PADDING", -+ "auto", -+ ["auto", "gluon", "asm"], -+ case_sensitive=False, -+ ), - # Whether to use aiter mha ops. - # By default is enabled. - "VLLM_ROCM_USE_AITER_MHA": lambda: ( -DIFF_VLLM_ENVS -apply_one "vllm/envs.py" "VLLM_ROCM_AITER_MLA_ASM_PADDING" "$WS/VLLM_ENVS.diff" - -cat > "$WS/KIMI_NVIDIA_MLA.diff" <<'DIFF_KIMI_NVIDIA_MLA' -diff --git a/vllm/models/kimi_k3/nvidia/mla.py b/vllm/models/kimi_k3/nvidia/mla.py ---- a/vllm/models/kimi_k3/nvidia/mla.py -+++ b/vllm/models/kimi_k3/nvidia/mla.py -@@ -594,8 +594,7 @@ - cos_sin_cache: torch.Tensor | None, - slot_mapping: torch.Tensor, - ) -> torch.Tensor: -- """Fused decode query-concat + latent cache insert, dispatched by cache -- dtype (same policy as prefill: fp8 cache -> fp8 query).""" -+ """Build the decode query and update the cache for its dtype/backend.""" - if self.kv_cache_dtype == "fp8_ds_mla": - cache = self.kv_cache - if cache.dtype != torch.uint8: -@@ -612,10 +611,21 @@ - cos_sin_cache=cos_sin_cache, - ) - if is_quantized_kv_cache(self.kv_cache_dtype): -- assert self.impl.supports_quant_query_input, ( # type: ignore[attr-defined] -- "Kimi-K3 fp8 KV cache decode requires a backend that accepts an " -- "fp8 (quantized) query input." -- ) -+ if not self.impl.supports_quant_query_input: # type: ignore[attr-defined] -+ if positions is not None: -+ assert self.rotary_emb is not None -+ q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe) -+ q_pe = q_pe.to(ql_nope.dtype) -+ k_pe = k_pe.to(kv_c_normed.dtype) -+ self.impl.do_kv_cache_update( # type: ignore[attr-defined] -+ kv_c_normed, -+ k_pe, -+ self.kv_cache, -+ slot_mapping, -+ self.kv_cache_dtype, -+ self._k_scale, -+ ) -+ return torch.cat((ql_nope, q_pe), dim=-1) - cache = self.kv_cache - if cache.dtype != torch.float8_e4m3fn: - cache = cache.view(torch.float8_e4m3fn) -DIFF_KIMI_NVIDIA_MLA -apply_one "vllm/models/kimi_k3/nvidia/mla.py" "if not self.impl.supports_quant_query_input" "$WS/KIMI_NVIDIA_MLA.diff" - -cat > "$WS/ATTN_UTILS.diff" <<'DIFF_ATTN_UTILS' -diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py ---- a/vllm/v1/worker/gpu/attn_utils.py -+++ b/vllm/v1/worker/gpu/attn_utils.py -@@ -92,6 +92,7 @@ - kv_cache_config: KVCacheConfig, - vllm_config: VllmConfig, - device: torch.device, -+ cg_support_exclude_layers: set[str] | None = None, - active_layer_names: set[str] | None = None, - ) -> tuple[list[list[AttentionGroup]], AttentionCGSupportInfo, list[int]]: - # Phase 1: discover attention groups for each kv cache group. -@@ -165,6 +166,15 @@ - else: - if hasattr(builder, "set_workspace_buffer"): - builder.set_workspace_buffer(attn_backend_workspace) -+ # A group owned entirely by a separately-managed model part must -+ # not constrain this runner: a spec-decode draft gets its own -+ # CudaGraphManager and has a first-class eager fallback, so letting -+ # it in here downgrades the target for a decision it does not share. -+ if ( -+ cg_support_exclude_layers is not None -+ and set(group.layer_names) <= cg_support_exclude_layers -+ ): -+ continue - # Check cudagraph support for the attention backend - cg_support = builder.get_cudagraph_support( - vllm_config, -DIFF_ATTN_UTILS -apply_one "vllm/v1/worker/gpu/attn_utils.py" "cg_support_exclude_layers" "$WS/ATTN_UTILS.diff" - -cat > "$WS/MODEL_RUNNER.diff" <<'DIFF_MODEL_RUNNER' -diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py ---- a/vllm/v1/worker/gpu/model_runner.py -+++ b/vllm/v1/worker/gpu/model_runner.py -@@ -488,7 +488,14 @@ - max_num_blocks_per_group.append(max_num_blocks) - - self.attn_groups, attn_cg_support, self.kernel_block_sizes = init_attn_backend( -- self.kv_cache_config, self.vllm_config, self.device -+ self.kv_cache_config, -+ self.vllm_config, -+ self.device, -+ cg_support_exclude_layers=( -+ self.speculator.draft_attn_layer_names -+ if isinstance(self.speculator, DraftModelSpeculator) -+ else None -+ ), - ) - attn_cg_support = attn_cg_support.narrow( - *self.model_state.get_additional_cg_support() -DIFF_MODEL_RUNNER -apply_one "vllm/v1/worker/gpu/model_runner.py" "cg_support_exclude_layers" "$WS/MODEL_RUNNER.diff" - -cat > "$WS/KDA_FUSED_RECURRENT.diff" <<'DIFF_KDA_FUSED_RECURRENT' -diff --git a/vllm/models/kimi_k3/amd/ops/third_party/kda/fused_recurrent.py b/vllm/models/kimi_k3/amd/ops/third_party/kda/fused_recurrent.py -index 2f512df62643..db519fb6f0db 100644 ---- a/vllm/models/kimi_k3/amd/ops/third_party/kda/fused_recurrent.py -+++ b/vllm/models/kimi_k3/amd/ops/third_party/kda/fused_recurrent.py -@@ -459,6 +459,7 @@ def fused_recurrent_kda_packed_decode_kernel( - stride_g_token: tl.constexpr, - stride_beta_token: tl.constexpr, - stride_state_token: tl.constexpr, -+ stride_state_indices, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, -@@ -476,7 +477,7 @@ def fused_recurrent_kda_packed_decode_kernel( - mask_v = o_v < V - mask_state = mask_v[:, None] & mask_k[None, :] - -- state_idx = tl.load(state_indices + i_n).to(tl.int64) -+ state_idx = tl.load(state_indices + i_n * stride_state_indices).to(tl.int64) - p_out = out + (i_n * H + i_h) * V + o_v - if state_idx <= 0: - tl.store(p_out, tl.zeros([BV], dtype=tl.float32), mask=mask_v) -@@ -560,8 +561,8 @@ def fused_recurrent_kda_packed_decode( - raise ValueError("`raw_beta` heads must be contiguous.") - if initial_state.stride()[1:] != (V * K, K, 1): - raise ValueError("`initial_state` must be contiguous within each cache slot.") -- if state_indices.ndim != 1 or state_indices.stride(0) != 1: -- raise ValueError("`state_indices` must be contiguous and one-dimensional.") -+ if state_indices.ndim != 1: -+ raise ValueError("`state_indices` must be one-dimensional.") - if A_log.ndim != 1 or not A_log.is_contiguous(): - raise ValueError("`A_log` must be contiguous and one-dimensional.") - if not dt_bias.is_contiguous(): -@@ -608,6 +609,7 @@ def fused_recurrent_kda_packed_decode( - stride_g_token=raw_g.stride(1), - stride_beta_token=raw_beta.stride(1), - stride_state_token=initial_state.stride(0), -+ stride_state_indices=state_indices.stride(0), - H=H, - K=K, - V=V, -DIFF_KDA_FUSED_RECURRENT -apply_one "vllm/models/kimi_k3/amd/ops/third_party/kda/fused_recurrent.py" "stride_state_indices" "$WS/KDA_FUSED_RECURRENT.diff" - - -say "3/4 aiter #4521 (fp8 cp round-robin asm MLA verify kernels) [needs network + hipcc + GPU]" -# Unlike the offline Python diffs above, #4521 ships BINARY .co kernels (not in -# the GitHub .diff) and a C++ asm_mla.cu change that must be recompiled, so this -# section fetches from GitHub and rebuilds module_mla_asm (which imports aiter -> -# needs a GPU). Skipped-idempotent via the csv marker. -PR4521_SHA="0cbedbb1bc5b3b254dd12ca4e8d3c7638b86830b" # merged head of ROCm/aiter#4521 -PR4521_RAW="https://raw.githubusercontent.com/ROCm/aiter/$PR4521_SHA" -META="$ROOT/aiter_meta"; MLADIR="$META/hsa/gfx950/mla" -if grep -qF "mla_a8w8_qh16_qseqlen4_gqaratio16_cprr_v3_ps.co" "$MLADIR/mla_asm.csv" 2>/dev/null; then - echo " #4521: already present (skip)" -elif [ "${WITH_PR4521:-1}" != "1" ]; then - echo " #4521: SKIPPED (WITH_PR4521!=1)" -else - # 1) binary .co verify kernels (4 new cprr + 4 updated); leave orphans, csv gates load - for f in \ - mla_a8w8_qh16_qseqlen4_gqaratio16_cprr_v3_ps.co \ - mla_a8w8_qh16_qseqlen4_gqaratio16_lse_cprr_v3_ps.co \ - mla_a8w8_qh16_qseqlen4_gqaratio16_lse_v3_ps.co \ - mla_a8w8_qh16_qseqlen4_gqaratio16_v3_ps.co \ - mla_a8w8_qh32_qseqlen4_gqaratio32_cprr_ps.co \ - mla_a8w8_qh32_qseqlen4_gqaratio32_lse_cprr_ps.co \ - mla_a8w8_qh32_qseqlen4_gqaratio32_lse_ps.co \ - mla_a8w8_qh32_qseqlen4_gqaratio32_ps.co ; do - if curl -ksSL -o "$MLADIR/$f.new" "$PR4521_RAW/hsa/gfx950/mla/$f" \ - && [ "$(stat -c %s "$MLADIR/$f.new" 2>/dev/null || echo 0)" -gt 1000 ]; then - mv "$MLADIR/$f.new" "$MLADIR/$f"; echo " co OK $f" - else - rm -f "$MLADIR/$f.new"; echo " co FAIL $f" - fi - done - # 2) text diffs. Two install roots: aiter/*.py -> $ROOT ; csrc + mla_asm.csv -> $META - curl -ksSL -o "$WS/pr4521.diff" "https://github.com/ROCm/aiter/pull/4521.diff" - awk -v A="$WS/pr4521_A.diff" -v B="$WS/pr4521_B.diff" ' - /^diff --git /{p=$0; sub(/^diff --git a\//,"",p); sub(/ .*/,"",p); a=0; b=0; - if (p ~ /^aiter\//) a=1; - else if (p ~ /^csrc\// || p=="hsa/gfx950/mla/mla_asm.csv") b=1 } - { if (a) print > A; else if (b) print > B } - ' "$WS/pr4521.diff" - git apply --directory="$ROOT" -p1 --unsafe-paths --whitespace=nowarn "$WS/pr4521_A.diff" 2>/dev/null \ - || patch -p1 -d "$ROOT" --fuzz=3 --forward --no-backup-if-mismatch < "$WS/pr4521_A.diff" - git apply --directory="$META" -p1 --unsafe-paths --whitespace=nowarn "$WS/pr4521_B.diff" 2>/dev/null \ - || patch -p1 -d "$META" --fuzz=3 --forward --no-backup-if-mismatch < "$WS/pr4521_B.diff" - # 3) force module_mla_asm rebuild (aiter JIT only rebuilds when the .so is gone) - rm -f "$ROOT/aiter/jit/module_mla_asm.so" - UJ="$(python -c 'from aiter.jit.core import get_user_jit_dir as g; print(g())' 2>/dev/null)" - [ -n "$UJ" ] && rm -f "$UJ/module_mla_asm.so" - python - <<'PYBUILD' -from aiter.jit.core import get_args_of_build, build_module -d = get_args_of_build("module_mla_asm") -build_module("module_mla_asm", d["srcs"], d["flags_extra_cc"], d["flags_extra_hip"], - d["blob_gen_cmd"], d["extra_include"], d["extra_ldflags"], d["verbose"], - d["is_python_module"], d["is_standalone"], d["torch_exclude"], - d.get("third_party", []), d.get("hipify", False), - d.get("flags_extra_hip_per_source", {})) -print(" module_mla_asm rebuilt") -PYBUILD - echo " #4521: APPLIED" -fi - -say "4/4 verify markers + py_compile + import" -echo "chk mla_gluon.py = $(grep -c '1 <= batch_size <= 256' "$ROOT/aiter/ops/triton/gluon/mla_gluon.py")" -echo "chk gemm_op_a16w16.py = $(grep -c 'is_current_stream_capturing' "$ROOT/aiter/ops/gemm_op_a16w16.py")" -echo "chk rocm_aiter_mla.py (base) = $(grep -c 'flat_kv_indices' "$ROOT/vllm/v1/attention/backends/mla/rocm_aiter_mla.py")" -echo "chk rocm_aiter_mla.py (DSpark) = $(grep -c 'Local DSpark PS extension' "$ROOT/vllm/v1/attention/backends/mla/rocm_aiter_mla.py") (expect 2)" -echo "chk #4521 mla_asm.csv = $(grep -c 'qh16_qseqlen4_gqaratio16_cprr' "$ROOT/aiter_meta/hsa/gfx950/mla/mla_asm.csv" 2>/dev/null) (expect 2; 0 if WITH_PR4521=0)" -echo "chk #4521 module_mla_asm.so = $([ -f "$ROOT/aiter/jit/module_mla_asm.so" ] && echo present || echo MISSING)" -echo "chk triton_mla.py = $(grep -c 'get_cudagraph_support' "$ROOT/vllm/v1/attention/backends/mla/triton_mla.py")" -echo "chk gpu_worker.py = $(grep -c 'import get_kv_cache_capacity' "$ROOT/vllm/v1/worker/gpu_worker.py")" -echo "chk envs.py = $(grep -c 'VLLM_ROCM_AITER_MLA_ASM_PADDING' "$ROOT/vllm/envs.py")" -echo "chk mla.py = $(grep -c 'if not self.impl.supports_quant_query_input' "$ROOT/vllm/models/kimi_k3/nvidia/mla.py")" -echo "chk attn_utils.py = $(grep -c 'cg_support_exclude_layers' "$ROOT/vllm/v1/worker/gpu/attn_utils.py")" -echo "chk model_runner.py = $(grep -c 'cg_support_exclude_layers' "$ROOT/vllm/v1/worker/gpu/model_runner.py")" -echo "chk fused_recurrent.py = $(grep -c 'reshape(-1).contiguous()' "$ROOT/vllm/models/kimi_k3/amd/ops/third_party/kda/fused_recurrent.py")" -echo "triton = $(python -c 'import triton; print(triton.__version__)') (expect 3.7.0*)" -python -m py_compile "$ROOT/aiter/ops/triton/gluon/mla_gluon.py" \ - "$ROOT/aiter/ops/gemm_op_a16w16.py" \ - "$ROOT/vllm/v1/attention/backends/mla/rocm_aiter_mla.py" \ - "$ROOT/vllm/v1/attention/backends/mla/triton_mla.py" \ - "$ROOT/vllm/v1/worker/gpu_worker.py" \ - "$ROOT/vllm/envs.py" \ - "$ROOT/vllm/models/kimi_k3/nvidia/mla.py" \ - "$ROOT/vllm/v1/worker/gpu/attn_utils.py" \ - "$ROOT/vllm/v1/worker/gpu/model_runner.py" \ - "$ROOT/vllm/models/kimi_k3/amd/ops/third_party/kda/fused_recurrent.py" && echo "PY_COMPILE_OK" || { echo "PY_COMPILE_FAIL"; exit 1; } -# Runtime import needs a GPU (aiter probes rocminfo); best-effort. -python - <<'PYEOF' -import importlib, traceback -mods = ("vllm.envs", - "vllm.v1.attention.backends.mla.rocm_aiter_mla", - "vllm.v1.attention.backends.mla.triton_mla", - "vllm.v1.worker.gpu_worker", - "vllm.v1.worker.gpu.attn_utils", - "vllm.v1.worker.gpu.model_runner") -try: - for m in mods: importlib.import_module(m) - import aiter.ops.gemm_op_a16w16 # noqa: F401 - import aiter.ops.triton.gluon.mla_gluon # noqa: F401 - print("IMPORT_OK") -except Exception as e: - print("IMPORT_SKIPPED (needs GPU?):", type(e).__name__, str(e).splitlines()[-1] if str(e) else "") -PYEOF -echo -echo "[embed] DONE. Launch server_final_CI.sh (MODEL_PATH + max_cudagraph_capture_size=44)." diff --git a/benchmarks/single_node/agentic/apply_triton_mla_cudagraph_fix.sh b/benchmarks/single_node/agentic/apply_triton_mla_cudagraph_fix.sh deleted file mode 100644 index 01c1634fa..000000000 --- a/benchmarks/single_node/agentic/apply_triton_mla_cudagraph_fix.sh +++ /dev/null @@ -1,65 +0,0 @@ -#!/usr/bin/env bash -# Let DSpark spec-decode keep FULL cudagraphs on ROCm. -# -# vllm/v1/attention/backends/mla/triton_mla.py declares -# TritonMLAMetadataBuilder._cudagraph_support = UNIFORM_SINGLE_TOKEN_DECODE -# which caps min_cg_support below UNIFORM_BATCH, so config/compilation.py:1443 -# downgrades FULL_AND_PIECEWISE -> PIECEWISE whenever spec-decode is enabled: -# "CUDAGraphMode.FULL_AND_PIECEWISE is not supported with spec-decode for -# attention backend TritonMLABackend" -# v1/worker/gpu/spec_decode/dflash/speculator.py:110-127 then gives the DSpark -# drafter CUDAGraphMode.NONE -- fully eager, with NO warning logged. Every draft -# layer and the Markov head dispatch kernel-by-kernel from Python each step. -# -# TRITON_MLA cannot simply be swapped out: it is the only ROCm MLA backend with -# supports_non_causal_multi_token_decode = True, which the DSpark draft needs. -# ROCM_AITER_MLA fails with "non-causal attention not supported". -# -# The builder already sets supports_non_causal_multi_token_decode = True and -# calls _init_reorder_batch_threshold(1, supports_spec_as_decode=True) "so -# full-cudagraph capture admits it", so UNIFORM_BATCH is the consistent value. -# -# MEASURED on 8x MI355X, Kimi-K3 MXFP4 TP8, DSpark, single stream, 600-tok gens: -# before: 14.05 tok/s, ITL 71.16 ms (PIECEWISE, drafter eager) -# after : 77.65 tok/s, ITL 12.88 ms (FULL cudagraphs) = 5.52x -# output verified correct in both ("17*23" -> 391, finish_reason stop) -# -# Idempotent. No-op if already patched or if the anchor is absent. -set -euo pipefail -PY=${PYTHON:-python3} - -TARGET=$($PY - <<'EOF' -import os -try: - import vllm.v1.attention.backends.mla.triton_mla as m - print(os.path.abspath(m.__file__)) -except Exception: - print("") -EOF -) -if [ -z "$TARGET" ] || [ ! -f "$TARGET" ]; then - echo "[triton-mla-fix] triton_mla.py not found; nothing to do." - exit 0 -fi -if grep -q "AttentionCGSupport.UNIFORM_BATCH" "$TARGET"; then - echo "[triton-mla-fix] already patched: $TARGET" - exit 0 -fi -cp -n "$TARGET" "$TARGET.orig" || true -$PY - "$TARGET" <<'EOF' -import sys, io -path = sys.argv[1] -src = io.open(path, encoding="utf-8").read() -old = """ _cudagraph_support: ClassVar[AttentionCGSupport] = ( - AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE - )""" -new = """ # PATCHED: UNIFORM_SINGLE_TOKEN_DECODE forced a PIECEWISE downgrade under - # spec-decode, which silently made the DSpark drafter fully eager. - _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH""" -if src.count(old) != 1: - sys.stderr.write("[triton-mla-fix] anchor missing or not unique; aborting.\n") - sys.exit(1) -io.open(path, "w", encoding="utf-8").write(src.replace(old, new)) -print("[triton-mla-fix] patched", path) -EOF -echo "[triton-mla-fix] done." diff --git a/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh b/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh index 87295393b..64d609daa 100644 --- a/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh +++ b/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh @@ -81,25 +81,14 @@ amd-smi || true resolve_trace_source install_agentic_deps -# ---- AITER pybind11 fix ------------------------------------------------------ -# The image's prebuilt aiter .so files are compiled against torch's bundled -# pybind11 (PYBIND11_INTERNALS_VERSION 11), but aiter's JIT builder injects the -# standalone pybind11 3.1.0 (version 12) as a -I flag, which outranks the -# -isystem path holding torch's copy. pybind11 keeps a SEPARATE type registry -# per internals id, so a JIT-built module cannot see aiter_tensor_t registered -# by the prebuilt core, and the first call dies during model warmup with: -# TypeError: fmha_fwd_bf16_opus_fwd(): incompatible function arguments -# The script below is idempotent, verifies the mismatch actually exists before -# touching anything, and self-disables once the image ships a fixed aiter. -bash "$(dirname "$0")/apply_aiter_pybind11_fix.sh" || true - -# ---- DSpark FULL-cudagraph fix ---------------------------------------------- -# TritonMLA declares _cudagraph_support=UNIFORM_SINGLE_TOKEN_DECODE, which forces -# FULL_AND_PIECEWISE -> PIECEWISE under spec-decode and then silently gives the -# DSpark drafter CUDAGraphMode.NONE (fully eager). Measured on 8x MI355X, single -# stream 600-token generations: 14.05 -> 77.65 tok/s, ITL 71.16 -> 12.88 ms (5.52x), -# output verified correct. Idempotent; no-op if already patched. -bash "$(dirname "$0")/apply_triton_mla_cudagraph_fix.sh" || true +# ---- In-container patches ---------------------------------------------------- +# Three fixes, all confined to this container's site-packages, all idempotent +# and all self-disabling once the image ships them: +# [1] aiter pybind11 internals mismatch -> unblocks ROCM_AITER_FA prefill +# [2] TritonMLA cudagraph support -> FULL cudagraphs for DSpark (5.52x TPOT) +# [3] KV block-pool negative-count clamp -> stops the mid-run engine crash +# Set SKIP_KIMI_PATCHES=1 to run stock. +bash "$(dirname "$0")/apply_kimi_k3_patches.sh" || true # ---- Reference env block ---------------------------------------------------- # Keep ALL of these. Commenting them out does not avoid the AITER FMHA crash: @@ -195,6 +184,30 @@ else ) fi +# ---- Async scheduling / KV block-pool stability ------------------------------ +# DSpark is the ONLY spec method exempted from vLLM's async-scheduling disable +# list (config/vllm.py:1181), so async_scheduling resolves True here. That gives +# max_concurrent_batches = pp_size + 1 = 2 (vllm.py:563-569), and with +# kv_role=kv_both (is_kv_consumer=True) the scheduler sets defer_block_free=True +# (sched/scheduler.py:155-157). Its own comment: "a step may still be writing a +# freed request's KV blocks. A consumer KV Connector can reallocate and fill +# those blocks via a load that isn't ordered against that write." +# +# That limbo state matches our crash signature exactly -- the engine dies with +# block_pool.py:667 assert block.ref_cnt == 0 +# i.e. a block sitting on the FREE list that is still referenced. Crash time +# scales inversely with concurrency: c10 survived 3612 s, c12 died at 487 s, +# c16 at 354 s. Note vLLM already disables async scheduling for ROCm DeepEP DBO +# because "that combination can corrupt" state. +# +# Setting max_concurrent_batches back to 1 makes defer_block_free unreachable. +# Cost: async scheduling exists to fill GPU-utilisation gaps, so expect to give +# some throughput back. Set ASYNC_SCHEDULING=1 to restore the default. +ASYNC_SCHED_ARGS=() +if [ "${ASYNC_SCHEDULING:-0}" != "1" ]; then + ASYNC_SCHED_ARGS=(--no-async-scheduling) +fi + # ---- MLA prefill backend ----------------------------------------------------- # On ROCm the prefill priority is [ROCM_AITER_FA, FLASH_ATTN]. ROCM_AITER_FA # JIT-builds module_fmha_fwd_bf16_opus at runtime; that module registers its own @@ -205,7 +218,7 @@ fi # Pinning FLASH_ATTN keeps every AITER MoE kernel (and its throughput) while # skipping only the broken FMHA prefill path. # UPDATE: the AITER packaging issue is now fixed at source by -# apply_aiter_pybind11_fix.sh (run above), so ROCM_AITER_FA is usable again and +# apply_kimi_k3_patches.sh (run above), so ROCM_AITER_FA is usable again and # is the default. Measured on 8x MI355X / Kimi-K3 MXFP4 TP8, cold prefill: # ~24k ctx FLASH_ATTN 12,953 -> AITER 13,524 tok/s (+4.4%) # ~93k ctx FLASH_ATTN 11,174 -> AITER 13,423 tok/s (+20.1%) @@ -251,6 +264,7 @@ VLLM_CMD=( --max-model-len 1048576 --enable-prefix-caching --kv-cache-dtype "fp8" + "${ASYNC_SCHED_ARGS[@]}" "${MLA_PREFILL_ARGS[@]}" "${COMPILATION_CONFIG_ARGS[@]}" "${SPEC_ARGS[@]}" From 9bc673b0edc34ed8ed47c4d27d74a709fb21e02a Mon Sep 17 00:00:00 2001 From: Sirra Date: Fri, 14 Aug 2026 13:45:02 +0530 Subject: [PATCH 20/31] [AMD] [WIP] [AGENTX] KIMI-K3 Perf with fixes applied. Signed-off-by: Sirra --- .../agentic/apply_kimi_k3_patches.sh | 245 ++++++++++++++++++ 1 file changed, 245 insertions(+) create mode 100644 benchmarks/single_node/agentic/apply_kimi_k3_patches.sh diff --git a/benchmarks/single_node/agentic/apply_kimi_k3_patches.sh b/benchmarks/single_node/agentic/apply_kimi_k3_patches.sh new file mode 100644 index 000000000..3e2993875 --- /dev/null +++ b/benchmarks/single_node/agentic/apply_kimi_k3_patches.sh @@ -0,0 +1,245 @@ +#!/usr/bin/env bash +# ============================================================================= +# Kimi-K3 / MI355X (gfx950) in-container patches — all three in one place. +# +# Everything here patches files inside the running container only +# (site-packages). Nothing outside the container is touched. Each patch is +# idempotent, verifies its own anchor, backs up to .orig, and no-ops if +# the image already ships the fix. A failed anchor aborts that patch cleanly +# rather than corrupting the file, so a future image with different sources +# degrades to "unpatched", never to "broken". +# +# [1] aiter pybind11 internals mismatch -> unblocks ROCM_AITER_FA prefill +# [2] TritonMLA cudagraph support -> FULL cudagraphs for DSpark (5.52x TPOT) +# [3] KV block-pool negative-count clamp -> stops the mid-run engine crash +# +# Env: +# SKIP_KIMI_PATCHES=1 skip everything +# PYTHON=... interpreter to use (default python3) +# ============================================================================= +set -euo pipefail +PY=${PYTHON:-python3} + +if [ "${SKIP_KIMI_PATCHES:-0}" = "1" ]; then + echo "[kimi-patches] SKIP_KIMI_PATCHES=1, doing nothing." + exit 0 +fi + +# Locate an installed module's file, or empty string if unimportable. +_modfile() { + $PY - "$1" <<'EOF' +import importlib, os, sys +try: + print(os.path.abspath(importlib.import_module(sys.argv[1]).__file__)) +except Exception: + print("") +EOF +} + +# _patch <<'PYEOF' ... old/new python ... PYEOF +# The heredoc body must define OLD and NEW strings. +_patch() { + local target="$1" marker="$2" label="$3" + if [ -z "$target" ] || [ ! -f "$target" ]; then + echo "[$label] target not found; skipping." + return 0 + fi + if grep -q "$marker" "$target"; then + echo "[$label] already patched." + return 0 + fi + cp -n "$target" "$target.orig" 2>/dev/null || true + if $PY - "$target" "$label"; then + return 0 + else + echo "[$label] patch failed; left unchanged." >&2 + return 0 + fi +} + +# ----------------------------------------------------------------------------- +# [1] aiter: JIT modules must use torch's bundled pybind11 +# ----------------------------------------------------------------------------- +# aiter/jit/utils/cpp_extension.py appends the STANDALONE pybind11 include via +# -I, which outranks the -isystem path carrying torch's bundled copy. The 117 +# prebuilt aiter .so are built against torch's (PYBIND11_INTERNALS_VERSION 11); +# the standalone package here is version 12. pybind11 keeps a SEPARATE type +# registry per internals id, so a JIT-built module cannot see aiter_tensor_t +# registered by the prebuilt core and the first call dies during warmup with +# TypeError: fmha_fwd_bf16_opus_fwd(): incompatible function arguments +# even though arity and types match exactly. +patch_aiter_pybind11() { + local label="aiter-pybind11" + local target; target=$(_modfile aiter.jit.utils.cpp_extension) + if [ -z "$target" ] || [ ! -f "$target" ]; then + echo "[$label] aiter not present; skipping."; return 0 + fi + + # Only act if the two pybind11s actually disagree. + local need; need=$($PY - <<'EOF' +import os, re +try: + import torch, pybind11 +except Exception: + print("no"); raise SystemExit +def ver(p): + f = os.path.join(p, "pybind11", "detail", "internals.h") + if not os.path.isfile(f): return None + m = re.search(r"define\s+PYBIND11_INTERNALS_VERSION\s+(\d+)", open(f).read()) + return int(m.group(1)) if m else None +t = ver(os.path.join(os.path.dirname(torch.__file__), "include")) +s = ver(pybind11.get_include()) +print("yes" if (t is not None and s is not None and t != s) else "no") +EOF +) + if [ "$need" != "yes" ]; then + echo "[$label] pybind11 internals already agree; nothing to do."; return 0 + fi + + if grep -q "_use_torch_pybind11" "$target"; then + echo "[$label] already patched." + else + cp -n "$target" "$target.orig" 2>/dev/null || true + $PY - "$target" <<'EOF' || echo "[aiter-pybind11] patch failed; unchanged." >&2 +import sys, io +p = sys.argv[1] +src = io.open(p, encoding="utf-8").read() +old = " extra_include_paths.append(pybind11.get_include())\n" +new = ( + " # PATCHED: prefer torch's bundled pybind11 so JIT modules land in the\n" + " # same pybind11 type registry as the prebuilt .so files.\n" + " _use_torch_pybind11 = False\n" + " if not torch_exclude:\n" + " _use_torch_pybind11 = os.path.isdir(\n" + " os.path.join(TORCH_INCLUDE_ROOT, \"pybind11\")\n" + " )\n" + " if not _use_torch_pybind11:\n" + " extra_include_paths.append(pybind11.get_include())\n" +) +if src.count(old) != 1: + sys.stderr.write("[aiter-pybind11] anchor missing or not unique; aborting.\n") + sys.exit(1) +io.open(p, "w", encoding="utf-8").write(src.replace(old, new)) +print("[aiter-pybind11] patched", p) +EOF + fi + + # Drop JIT artifacts built against the wrong pybind11 so they rebuild. + # aiter honours AITER_JIT_DIR and falls back to ~/.aiter when dist-packages + # is read-only, so ask aiter rather than deriving the path from $target. + local jitdir + jitdir=$($PY -c 'from aiter.jit.core import get_user_jit_dir; print(get_user_jit_dir())' 2>/dev/null || true) + [ -n "$jitdir" ] && [ -d "$jitdir" ] || jitdir=$(dirname "$(dirname "$target")") + shopt -s nullglob + for so in "$jitdir"/*.so; do + if grep -qa "__pybind11_internals_v12" "$so" 2>/dev/null; then + rm -f "$so"; rm -rf "$jitdir/build/$(basename "${so%.so}")" + echo "[$label] removed stale v12 module: $(basename "$so")" + fi + done + shopt -u nullglob +} + +# ----------------------------------------------------------------------------- +# [2] vLLM: let DSpark spec-decode keep FULL cudagraphs +# ----------------------------------------------------------------------------- +# TritonMLAMetadataBuilder._cudagraph_support = UNIFORM_SINGLE_TOKEN_DECODE caps +# min_cg_support below UNIFORM_BATCH, so config/compilation.py downgrades +# FULL_AND_PIECEWISE -> PIECEWISE under spec-decode. dflash/speculator.py then +# gives the DSpark drafter CUDAGraphMode.NONE -- fully eager -- and logs nothing. +# TRITON_MLA cannot be swapped out: it is the only ROCm MLA backend with +# supports_non_causal_multi_token_decode=True, which DSpark requires +# (ROCM_AITER_MLA fails with "non-causal attention not supported"). +# The builder already calls _init_reorder_batch_threshold(1, +# supports_spec_as_decode=True) "so full-cudagraph capture admits it", so +# UNIFORM_BATCH is the self-consistent value. +# MEASURED 8x MI355X single stream, 600-token gens: +# before 14.05 tok/s ITL 71.16 ms -> after 77.65 tok/s ITL 12.88 ms (5.52x) +patch_triton_mla_cudagraph() { + local label="triton-mla-cudagraph" + local target; target=$(_modfile vllm.v1.attention.backends.mla.triton_mla) + if [ -z "$target" ] || [ ! -f "$target" ]; then + echo "[$label] target not found; skipping."; return 0 + fi + if grep -q "AttentionCGSupport.UNIFORM_BATCH" "$target"; then + echo "[$label] already patched."; return 0 + fi + cp -n "$target" "$target.orig" 2>/dev/null || true + $PY - "$target" <<'EOF' || echo "[triton-mla-cudagraph] patch failed; unchanged." >&2 +import sys, io +p = sys.argv[1] +src = io.open(p, encoding="utf-8").read() +old = """ _cudagraph_support: ClassVar[AttentionCGSupport] = ( + AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE + )""" +new = """ # PATCHED: UNIFORM_SINGLE_TOKEN_DECODE forced a PIECEWISE downgrade under + # spec-decode, which silently made the DSpark drafter fully eager. + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH""" +if src.count(old) != 1: + sys.stderr.write("[triton-mla-cudagraph] anchor missing or not unique; aborting.\n") + sys.exit(1) +io.open(p, "w", encoding="utf-8").write(src.replace(old, new)) +print("[triton-mla-cudagraph] patched", p) +EOF +} + +# ----------------------------------------------------------------------------- +# [3] vLLM: clamp the negative block count that corrupts the KV free list +# ----------------------------------------------------------------------------- +# single_type_kv_cache_manager.py, in allocate_external_computed_blocks(), is the +# ONLY unguarded get_new_blocks() call site in that file (siblings clamp or +# early-return). When len(req_blocks) exceeds the block count implied by +# num_total_computed_tokens the argument goes NEGATIVE, and a negative count is +# silently destructive rather than rejected: +# * block_pool.get_new_blocks only rejects num_blocks > free +# * popleft_n passes its own assert num_free_blocks >= n +# * it runs num_free_blocks -= n -> an INCREASE +# * range(n) iterates zero times, so the linked list is untouched +# num_free_blocks is then inflated relative to the real free list; a later +# legitimate pop walks past the tail and the engine dies mid-run on +# kv_cache_utils.py assert curr_block is not None +# block_pool.py assert block.ref_cnt == 0 +# Load-dependent: c10 died at 3612 s, c12 at 487 s, c16 at 354 s. On the EXTERNAL +# block path, so it needs --kv-transfer-config to appear. NOTE +# --no-async-scheduling was tested and does NOT help (c12 died at 490 s). +patch_kv_blockpool() { + local label="kv-blockpool" + local target; target=$(_modfile vllm.v1.core.single_type_kv_cache_manager) + if [ -z "$target" ] || [ ! -f "$target" ]; then + echo "[$label] target not found; skipping."; return 0 + fi + # NB: the marker must be unique to OUR patch. "num_new_blocks = max(" is NOT + # -- stock already has it at three other call sites (lines ~208/1511/1601), + # so using it silently skipped the patch on a clean image. + if grep -q "KIMI-PATCH-KV-BLOCKPOOL" "$target"; then + echo "[$label] already patched."; return 0 + fi + cp -n "$target" "$target.orig" 2>/dev/null || true + $PY - "$target" <<'EOF' || echo "[kv-blockpool] patch failed; unchanged." >&2 +import sys, io +p = sys.argv[1] +src = io.open(p, encoding="utf-8").read() +old = """ req_blocks = self.req_to_blocks[request_id] + allocated_blocks = self.block_pool.get_new_blocks( + cdiv(num_total_computed_tokens, self.block_size) - len(req_blocks) + )""" +new = """ req_blocks = self.req_to_blocks[request_id] + # KIMI-PATCH-KV-BLOCKPOOL: clamp to >= 0; a negative count silently + # inflates FreeKVCacheBlockQueue.num_free_blocks and corrupts the free list. + num_new_blocks = max( + 0, cdiv(num_total_computed_tokens, self.block_size) - len(req_blocks) + ) + allocated_blocks = self.block_pool.get_new_blocks(num_new_blocks)""" +if src.count(old) != 1: + sys.stderr.write("[kv-blockpool] anchor missing or not unique; aborting.\n") + sys.exit(1) +io.open(p, "w", encoding="utf-8").write(src.replace(old, new)) +print("[kv-blockpool] patched", p) +EOF +} + +echo "[kimi-patches] applying in-container patches..." +patch_aiter_pybind11 || true +patch_triton_mla_cudagraph || true +patch_kv_blockpool || true +echo "[kimi-patches] done." From f08086516392927827b4465c2325a8a7ed91e2b1 Mon Sep 17 00:00:00 2001 From: Sirra Date: Fri, 14 Aug 2026 13:47:56 +0530 Subject: [PATCH 21/31] [AMD] [WIP] [AGENTX] KIMI-K3 Perf with fixes applied. Signed-off-by: Sirra --- configs/amd-master.yaml | 31 +------------------------------ 1 file changed, 1 insertion(+), 30 deletions(-) diff --git a/configs/amd-master.yaml b/configs/amd-master.yaml index eed4c22f6..003b154db 100644 --- a/configs/amd-master.yaml +++ b/configs/amd-master.yaml @@ -649,7 +649,7 @@ kimik3-fp4-mi355x-vllm-agentic-mtp: - dram-utilization: 0.50 search-space: # - { tp: 8, kv-offloading: none, conc-list: [1, 4, 8] , spec-decoding: mtp} - - { tp: 8, ep: 1, kv-offloading: dram, kv-offload-backend: { name: vllm-simple }, conc-list: [10], spec-decoding: mtp } + - { tp: 8, ep: 1, kv-offloading: dram, kv-offload-backend: { name: vllm-simple }, conc-list: [1, 4, 8, 10], spec-decoding: mtp } dsr1-fp4-mi355x-sglang-disagg: image: lmsysorg/sglang-rocm:v0.5.12-rocm720-mi35x-20260519 @@ -1522,21 +1522,6 @@ minimaxm3-fp8-mi300x-vllm-agentic: - { tp: 8, kv-offloading: dram, kv-offload-backend: { name: mooncake, version: "0.3.11.post1" }, conc-list: [5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 16, 18, 20] } - { tp: 8, ep: 8, kv-offloading: dram, kv-offload-backend: { name: mooncake, version: "0.3.11.post1" }, conc-list: [5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 16, 18, 20] } -minimaxm3-fp8-mi300x-vllm-agentic-mtp: - image: vllm/vllm-openai-rocm:v0.27.1 - model: MiniMaxAI/MiniMax-M3-MXFP8 - model-prefix: minimaxm3 - runner: cluster:mi300x-amds - precision: fp8 - framework: vllm - multinode: false - scenarios: - agentic-coding: - - dram-utilization: 0.80 - search-space: - - { tp: 8, spec-decoding: mtp, kv-offloading: none, conc-list: [2, 4, 6, 8, 10] } - - { tp: 8, spec-decoding: mtp, kv-offloading: dram, kv-offload-backend: { name: lmcache, version: "0.5.3" }, conc-list: [16] } - # GLM-5.2 FP8 full-context AgentX refresh on MI325X. This preserves the TP8 # GPU-resident-KV c1/c2/c3/c4/c5/c6/c8 curve from Actions run 29657732517 # and enables EAGLE MTP with the committed thinking-on golden AL. @@ -1580,20 +1565,6 @@ minimaxm3-fp8-mi325x-vllm-agentic: - { tp: 8, ep: 8, kv-offloading: dram, kv-offload-backend: { name: mooncake, version: "0.3.11.post1" }, conc-list: [10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 32] } - { tp: 8, ep: 8, dp-attn: true, kv-offloading: dram, kv-offload-backend: { name: mooncake, version: "0.3.11.post1" }, conc-list: [24, 32, 36, 40, 44, 48, 52, 56, 60, 64, 72, 80, 96], router: { name: vllm-router, version: "0.1.14" } } -minimaxm3-fp8-mi325x-vllm-agentic-mtp: - image: vllm/vllm-openai-rocm:v0.27.1 - model: MiniMaxAI/MiniMax-M3-MXFP8 - model-prefix: minimaxm3 - runner: cluster:mi325x-amds - precision: fp8 - framework: vllm - multinode: false - scenarios: - agentic-coding: - - dram-utilization: 0.20 - search-space: - - { tp: 8, spec-decoding: mtp, kv-offloading: none, conc-list: [1, 2, 4, 8, 10, 12, 14, 16, 18] } - minimaxm3-fp4-mi355x-vllm-agentic: image: vllm/vllm-openai-rocm:nightly-dcfebf93f4eccf30f71872283331eee757915daf model: amd/MiniMax-M3-MXFP4 From 4901088df4c6e98f2a5d171e6e841284b3d456ae Mon Sep 17 00:00:00 2001 From: Sirra Date: Fri, 14 Aug 2026 15:37:19 +0530 Subject: [PATCH 22/31] [AMD] [WIP] [AGENTX] KIMI-K3 with fixes applied & Perf Boost Signed-off-by: Sirra --- .../agentic/apply_k3_container_patches.sh | 245 ++++++++++++++++++ configs/amd-master.yaml | 38 ++- 2 files changed, 278 insertions(+), 5 deletions(-) create mode 100644 benchmarks/single_node/agentic/apply_k3_container_patches.sh diff --git a/benchmarks/single_node/agentic/apply_k3_container_patches.sh b/benchmarks/single_node/agentic/apply_k3_container_patches.sh new file mode 100644 index 000000000..3e2993875 --- /dev/null +++ b/benchmarks/single_node/agentic/apply_k3_container_patches.sh @@ -0,0 +1,245 @@ +#!/usr/bin/env bash +# ============================================================================= +# Kimi-K3 / MI355X (gfx950) in-container patches — all three in one place. +# +# Everything here patches files inside the running container only +# (site-packages). Nothing outside the container is touched. Each patch is +# idempotent, verifies its own anchor, backs up to .orig, and no-ops if +# the image already ships the fix. A failed anchor aborts that patch cleanly +# rather than corrupting the file, so a future image with different sources +# degrades to "unpatched", never to "broken". +# +# [1] aiter pybind11 internals mismatch -> unblocks ROCM_AITER_FA prefill +# [2] TritonMLA cudagraph support -> FULL cudagraphs for DSpark (5.52x TPOT) +# [3] KV block-pool negative-count clamp -> stops the mid-run engine crash +# +# Env: +# SKIP_KIMI_PATCHES=1 skip everything +# PYTHON=... interpreter to use (default python3) +# ============================================================================= +set -euo pipefail +PY=${PYTHON:-python3} + +if [ "${SKIP_KIMI_PATCHES:-0}" = "1" ]; then + echo "[kimi-patches] SKIP_KIMI_PATCHES=1, doing nothing." + exit 0 +fi + +# Locate an installed module's file, or empty string if unimportable. +_modfile() { + $PY - "$1" <<'EOF' +import importlib, os, sys +try: + print(os.path.abspath(importlib.import_module(sys.argv[1]).__file__)) +except Exception: + print("") +EOF +} + +# _patch <<'PYEOF' ... old/new python ... PYEOF +# The heredoc body must define OLD and NEW strings. +_patch() { + local target="$1" marker="$2" label="$3" + if [ -z "$target" ] || [ ! -f "$target" ]; then + echo "[$label] target not found; skipping." + return 0 + fi + if grep -q "$marker" "$target"; then + echo "[$label] already patched." + return 0 + fi + cp -n "$target" "$target.orig" 2>/dev/null || true + if $PY - "$target" "$label"; then + return 0 + else + echo "[$label] patch failed; left unchanged." >&2 + return 0 + fi +} + +# ----------------------------------------------------------------------------- +# [1] aiter: JIT modules must use torch's bundled pybind11 +# ----------------------------------------------------------------------------- +# aiter/jit/utils/cpp_extension.py appends the STANDALONE pybind11 include via +# -I, which outranks the -isystem path carrying torch's bundled copy. The 117 +# prebuilt aiter .so are built against torch's (PYBIND11_INTERNALS_VERSION 11); +# the standalone package here is version 12. pybind11 keeps a SEPARATE type +# registry per internals id, so a JIT-built module cannot see aiter_tensor_t +# registered by the prebuilt core and the first call dies during warmup with +# TypeError: fmha_fwd_bf16_opus_fwd(): incompatible function arguments +# even though arity and types match exactly. +patch_aiter_pybind11() { + local label="aiter-pybind11" + local target; target=$(_modfile aiter.jit.utils.cpp_extension) + if [ -z "$target" ] || [ ! -f "$target" ]; then + echo "[$label] aiter not present; skipping."; return 0 + fi + + # Only act if the two pybind11s actually disagree. + local need; need=$($PY - <<'EOF' +import os, re +try: + import torch, pybind11 +except Exception: + print("no"); raise SystemExit +def ver(p): + f = os.path.join(p, "pybind11", "detail", "internals.h") + if not os.path.isfile(f): return None + m = re.search(r"define\s+PYBIND11_INTERNALS_VERSION\s+(\d+)", open(f).read()) + return int(m.group(1)) if m else None +t = ver(os.path.join(os.path.dirname(torch.__file__), "include")) +s = ver(pybind11.get_include()) +print("yes" if (t is not None and s is not None and t != s) else "no") +EOF +) + if [ "$need" != "yes" ]; then + echo "[$label] pybind11 internals already agree; nothing to do."; return 0 + fi + + if grep -q "_use_torch_pybind11" "$target"; then + echo "[$label] already patched." + else + cp -n "$target" "$target.orig" 2>/dev/null || true + $PY - "$target" <<'EOF' || echo "[aiter-pybind11] patch failed; unchanged." >&2 +import sys, io +p = sys.argv[1] +src = io.open(p, encoding="utf-8").read() +old = " extra_include_paths.append(pybind11.get_include())\n" +new = ( + " # PATCHED: prefer torch's bundled pybind11 so JIT modules land in the\n" + " # same pybind11 type registry as the prebuilt .so files.\n" + " _use_torch_pybind11 = False\n" + " if not torch_exclude:\n" + " _use_torch_pybind11 = os.path.isdir(\n" + " os.path.join(TORCH_INCLUDE_ROOT, \"pybind11\")\n" + " )\n" + " if not _use_torch_pybind11:\n" + " extra_include_paths.append(pybind11.get_include())\n" +) +if src.count(old) != 1: + sys.stderr.write("[aiter-pybind11] anchor missing or not unique; aborting.\n") + sys.exit(1) +io.open(p, "w", encoding="utf-8").write(src.replace(old, new)) +print("[aiter-pybind11] patched", p) +EOF + fi + + # Drop JIT artifacts built against the wrong pybind11 so they rebuild. + # aiter honours AITER_JIT_DIR and falls back to ~/.aiter when dist-packages + # is read-only, so ask aiter rather than deriving the path from $target. + local jitdir + jitdir=$($PY -c 'from aiter.jit.core import get_user_jit_dir; print(get_user_jit_dir())' 2>/dev/null || true) + [ -n "$jitdir" ] && [ -d "$jitdir" ] || jitdir=$(dirname "$(dirname "$target")") + shopt -s nullglob + for so in "$jitdir"/*.so; do + if grep -qa "__pybind11_internals_v12" "$so" 2>/dev/null; then + rm -f "$so"; rm -rf "$jitdir/build/$(basename "${so%.so}")" + echo "[$label] removed stale v12 module: $(basename "$so")" + fi + done + shopt -u nullglob +} + +# ----------------------------------------------------------------------------- +# [2] vLLM: let DSpark spec-decode keep FULL cudagraphs +# ----------------------------------------------------------------------------- +# TritonMLAMetadataBuilder._cudagraph_support = UNIFORM_SINGLE_TOKEN_DECODE caps +# min_cg_support below UNIFORM_BATCH, so config/compilation.py downgrades +# FULL_AND_PIECEWISE -> PIECEWISE under spec-decode. dflash/speculator.py then +# gives the DSpark drafter CUDAGraphMode.NONE -- fully eager -- and logs nothing. +# TRITON_MLA cannot be swapped out: it is the only ROCm MLA backend with +# supports_non_causal_multi_token_decode=True, which DSpark requires +# (ROCM_AITER_MLA fails with "non-causal attention not supported"). +# The builder already calls _init_reorder_batch_threshold(1, +# supports_spec_as_decode=True) "so full-cudagraph capture admits it", so +# UNIFORM_BATCH is the self-consistent value. +# MEASURED 8x MI355X single stream, 600-token gens: +# before 14.05 tok/s ITL 71.16 ms -> after 77.65 tok/s ITL 12.88 ms (5.52x) +patch_triton_mla_cudagraph() { + local label="triton-mla-cudagraph" + local target; target=$(_modfile vllm.v1.attention.backends.mla.triton_mla) + if [ -z "$target" ] || [ ! -f "$target" ]; then + echo "[$label] target not found; skipping."; return 0 + fi + if grep -q "AttentionCGSupport.UNIFORM_BATCH" "$target"; then + echo "[$label] already patched."; return 0 + fi + cp -n "$target" "$target.orig" 2>/dev/null || true + $PY - "$target" <<'EOF' || echo "[triton-mla-cudagraph] patch failed; unchanged." >&2 +import sys, io +p = sys.argv[1] +src = io.open(p, encoding="utf-8").read() +old = """ _cudagraph_support: ClassVar[AttentionCGSupport] = ( + AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE + )""" +new = """ # PATCHED: UNIFORM_SINGLE_TOKEN_DECODE forced a PIECEWISE downgrade under + # spec-decode, which silently made the DSpark drafter fully eager. + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH""" +if src.count(old) != 1: + sys.stderr.write("[triton-mla-cudagraph] anchor missing or not unique; aborting.\n") + sys.exit(1) +io.open(p, "w", encoding="utf-8").write(src.replace(old, new)) +print("[triton-mla-cudagraph] patched", p) +EOF +} + +# ----------------------------------------------------------------------------- +# [3] vLLM: clamp the negative block count that corrupts the KV free list +# ----------------------------------------------------------------------------- +# single_type_kv_cache_manager.py, in allocate_external_computed_blocks(), is the +# ONLY unguarded get_new_blocks() call site in that file (siblings clamp or +# early-return). When len(req_blocks) exceeds the block count implied by +# num_total_computed_tokens the argument goes NEGATIVE, and a negative count is +# silently destructive rather than rejected: +# * block_pool.get_new_blocks only rejects num_blocks > free +# * popleft_n passes its own assert num_free_blocks >= n +# * it runs num_free_blocks -= n -> an INCREASE +# * range(n) iterates zero times, so the linked list is untouched +# num_free_blocks is then inflated relative to the real free list; a later +# legitimate pop walks past the tail and the engine dies mid-run on +# kv_cache_utils.py assert curr_block is not None +# block_pool.py assert block.ref_cnt == 0 +# Load-dependent: c10 died at 3612 s, c12 at 487 s, c16 at 354 s. On the EXTERNAL +# block path, so it needs --kv-transfer-config to appear. NOTE +# --no-async-scheduling was tested and does NOT help (c12 died at 490 s). +patch_kv_blockpool() { + local label="kv-blockpool" + local target; target=$(_modfile vllm.v1.core.single_type_kv_cache_manager) + if [ -z "$target" ] || [ ! -f "$target" ]; then + echo "[$label] target not found; skipping."; return 0 + fi + # NB: the marker must be unique to OUR patch. "num_new_blocks = max(" is NOT + # -- stock already has it at three other call sites (lines ~208/1511/1601), + # so using it silently skipped the patch on a clean image. + if grep -q "KIMI-PATCH-KV-BLOCKPOOL" "$target"; then + echo "[$label] already patched."; return 0 + fi + cp -n "$target" "$target.orig" 2>/dev/null || true + $PY - "$target" <<'EOF' || echo "[kv-blockpool] patch failed; unchanged." >&2 +import sys, io +p = sys.argv[1] +src = io.open(p, encoding="utf-8").read() +old = """ req_blocks = self.req_to_blocks[request_id] + allocated_blocks = self.block_pool.get_new_blocks( + cdiv(num_total_computed_tokens, self.block_size) - len(req_blocks) + )""" +new = """ req_blocks = self.req_to_blocks[request_id] + # KIMI-PATCH-KV-BLOCKPOOL: clamp to >= 0; a negative count silently + # inflates FreeKVCacheBlockQueue.num_free_blocks and corrupts the free list. + num_new_blocks = max( + 0, cdiv(num_total_computed_tokens, self.block_size) - len(req_blocks) + ) + allocated_blocks = self.block_pool.get_new_blocks(num_new_blocks)""" +if src.count(old) != 1: + sys.stderr.write("[kv-blockpool] anchor missing or not unique; aborting.\n") + sys.exit(1) +io.open(p, "w", encoding="utf-8").write(src.replace(old, new)) +print("[kv-blockpool] patched", p) +EOF +} + +echo "[kimi-patches] applying in-container patches..." +patch_aiter_pybind11 || true +patch_triton_mla_cudagraph || true +patch_kv_blockpool || true +echo "[kimi-patches] done." diff --git a/configs/amd-master.yaml b/configs/amd-master.yaml index 003b154db..39cff4be5 100644 --- a/configs/amd-master.yaml +++ b/configs/amd-master.yaml @@ -637,7 +637,7 @@ dsr1-fp8-mi355x-sglang-disagg-mtp: - "DECODE_MTP_SIZE=2" kimik3-fp4-mi355x-vllm-agentic-mtp: - image: vllm/vllm-openai-rocm:nightly-3ee2df30337a301164c46ae444b76ee67e71c106 + image: vllm/vllm-openai-rocm:nightly-ac7509e2b1db40fec2f03dde1ed4e9dfdc2338c9 model: moonshotai/Kimi-K3 model-prefix: kimik3 runner: cluster:mi355x-amds @@ -646,10 +646,9 @@ kimik3-fp4-mi355x-vllm-agentic-mtp: multinode: false scenarios: agentic-coding: - - dram-utilization: 0.50 + - dram-utilization: 0.65 search-space: - # - { tp: 8, kv-offloading: none, conc-list: [1, 4, 8] , spec-decoding: mtp} - - { tp: 8, ep: 1, kv-offloading: dram, kv-offload-backend: { name: vllm-simple }, conc-list: [1, 4, 8, 10], spec-decoding: mtp } + - { tp: 8, ep: 1, kv-offloading: dram, kv-offload-backend: { name: vllm-simple }, conc-list: [1, 4, 8, 10, 12, 14, 16], spec-decoding: mtp } dsr1-fp4-mi355x-sglang-disagg: image: lmsysorg/sglang-rocm:v0.5.12-rocm720-mi35x-20260519 @@ -1522,6 +1521,21 @@ minimaxm3-fp8-mi300x-vllm-agentic: - { tp: 8, kv-offloading: dram, kv-offload-backend: { name: mooncake, version: "0.3.11.post1" }, conc-list: [5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 16, 18, 20] } - { tp: 8, ep: 8, kv-offloading: dram, kv-offload-backend: { name: mooncake, version: "0.3.11.post1" }, conc-list: [5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 16, 18, 20] } +minimaxm3-fp8-mi300x-vllm-agentic-mtp: + image: vllm/vllm-openai-rocm:v0.27.1 + model: MiniMaxAI/MiniMax-M3-MXFP8 + model-prefix: minimaxm3 + runner: cluster:mi300x-amds + precision: fp8 + framework: vllm + multinode: false + scenarios: + agentic-coding: + - dram-utilization: 0.80 + search-space: + - { tp: 8, spec-decoding: mtp, kv-offloading: none, conc-list: [2, 4, 6, 8, 10] } + - { tp: 8, spec-decoding: mtp, kv-offloading: dram, kv-offload-backend: { name: lmcache, version: "0.5.3" }, conc-list: [16] } + # GLM-5.2 FP8 full-context AgentX refresh on MI325X. This preserves the TP8 # GPU-resident-KV c1/c2/c3/c4/c5/c6/c8 curve from Actions run 29657732517 # and enables EAGLE MTP with the committed thinking-on golden AL. @@ -1565,6 +1579,20 @@ minimaxm3-fp8-mi325x-vllm-agentic: - { tp: 8, ep: 8, kv-offloading: dram, kv-offload-backend: { name: mooncake, version: "0.3.11.post1" }, conc-list: [10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 32] } - { tp: 8, ep: 8, dp-attn: true, kv-offloading: dram, kv-offload-backend: { name: mooncake, version: "0.3.11.post1" }, conc-list: [24, 32, 36, 40, 44, 48, 52, 56, 60, 64, 72, 80, 96], router: { name: vllm-router, version: "0.1.14" } } +minimaxm3-fp8-mi325x-vllm-agentic-mtp: + image: vllm/vllm-openai-rocm:v0.27.1 + model: MiniMaxAI/MiniMax-M3-MXFP8 + model-prefix: minimaxm3 + runner: cluster:mi325x-amds + precision: fp8 + framework: vllm + multinode: false + scenarios: + agentic-coding: + - dram-utilization: 0.20 + search-space: + - { tp: 8, spec-decoding: mtp, kv-offloading: none, conc-list: [1, 2, 4, 8, 10, 12, 14, 16, 18] } + minimaxm3-fp4-mi355x-vllm-agentic: image: vllm/vllm-openai-rocm:nightly-dcfebf93f4eccf30f71872283331eee757915daf model: amd/MiniMax-M3-MXFP4 @@ -1717,4 +1745,4 @@ glm5.2-fp4-mi355x-sglang-agentic-mtp: - dram-utilization: 0.8 search-space: - { tp: 4, ep: 4, kv-offloading: dram, kv-offload-backend: { name: hicache }, conc-list: [1, 2, 4, 8, 10, 12, 16], spec-decoding: mtp } - - { tp: 8, ep: 8, kv-offloading: none, conc-list: [1, 2, 4], spec-decoding: mtp } + - { tp: 8, ep: 8, kv-offloading: none, conc-list: [1, 2, 4], spec-decoding: mtp } \ No newline at end of file From 84bc10350c074e7817cdfd7ac64313e680224d1d Mon Sep 17 00:00:00 2001 From: Sirra Date: Fri, 14 Aug 2026 15:40:12 +0530 Subject: [PATCH 23/31] [AMD] [WIP] [AGENTX] KIMI-K3 with fixes applied & Perf Boost Signed-off-by: Sirra --- .../agentic/apply_kimi_k3_patches.sh | 245 ------------------ .../agentic/kimik3_fp4_mi355x_mtp.sh | 2 +- 2 files changed, 1 insertion(+), 246 deletions(-) delete mode 100644 benchmarks/single_node/agentic/apply_kimi_k3_patches.sh diff --git a/benchmarks/single_node/agentic/apply_kimi_k3_patches.sh b/benchmarks/single_node/agentic/apply_kimi_k3_patches.sh deleted file mode 100644 index 3e2993875..000000000 --- a/benchmarks/single_node/agentic/apply_kimi_k3_patches.sh +++ /dev/null @@ -1,245 +0,0 @@ -#!/usr/bin/env bash -# ============================================================================= -# Kimi-K3 / MI355X (gfx950) in-container patches — all three in one place. -# -# Everything here patches files inside the running container only -# (site-packages). Nothing outside the container is touched. Each patch is -# idempotent, verifies its own anchor, backs up to .orig, and no-ops if -# the image already ships the fix. A failed anchor aborts that patch cleanly -# rather than corrupting the file, so a future image with different sources -# degrades to "unpatched", never to "broken". -# -# [1] aiter pybind11 internals mismatch -> unblocks ROCM_AITER_FA prefill -# [2] TritonMLA cudagraph support -> FULL cudagraphs for DSpark (5.52x TPOT) -# [3] KV block-pool negative-count clamp -> stops the mid-run engine crash -# -# Env: -# SKIP_KIMI_PATCHES=1 skip everything -# PYTHON=... interpreter to use (default python3) -# ============================================================================= -set -euo pipefail -PY=${PYTHON:-python3} - -if [ "${SKIP_KIMI_PATCHES:-0}" = "1" ]; then - echo "[kimi-patches] SKIP_KIMI_PATCHES=1, doing nothing." - exit 0 -fi - -# Locate an installed module's file, or empty string if unimportable. -_modfile() { - $PY - "$1" <<'EOF' -import importlib, os, sys -try: - print(os.path.abspath(importlib.import_module(sys.argv[1]).__file__)) -except Exception: - print("") -EOF -} - -# _patch <<'PYEOF' ... old/new python ... PYEOF -# The heredoc body must define OLD and NEW strings. -_patch() { - local target="$1" marker="$2" label="$3" - if [ -z "$target" ] || [ ! -f "$target" ]; then - echo "[$label] target not found; skipping." - return 0 - fi - if grep -q "$marker" "$target"; then - echo "[$label] already patched." - return 0 - fi - cp -n "$target" "$target.orig" 2>/dev/null || true - if $PY - "$target" "$label"; then - return 0 - else - echo "[$label] patch failed; left unchanged." >&2 - return 0 - fi -} - -# ----------------------------------------------------------------------------- -# [1] aiter: JIT modules must use torch's bundled pybind11 -# ----------------------------------------------------------------------------- -# aiter/jit/utils/cpp_extension.py appends the STANDALONE pybind11 include via -# -I, which outranks the -isystem path carrying torch's bundled copy. The 117 -# prebuilt aiter .so are built against torch's (PYBIND11_INTERNALS_VERSION 11); -# the standalone package here is version 12. pybind11 keeps a SEPARATE type -# registry per internals id, so a JIT-built module cannot see aiter_tensor_t -# registered by the prebuilt core and the first call dies during warmup with -# TypeError: fmha_fwd_bf16_opus_fwd(): incompatible function arguments -# even though arity and types match exactly. -patch_aiter_pybind11() { - local label="aiter-pybind11" - local target; target=$(_modfile aiter.jit.utils.cpp_extension) - if [ -z "$target" ] || [ ! -f "$target" ]; then - echo "[$label] aiter not present; skipping."; return 0 - fi - - # Only act if the two pybind11s actually disagree. - local need; need=$($PY - <<'EOF' -import os, re -try: - import torch, pybind11 -except Exception: - print("no"); raise SystemExit -def ver(p): - f = os.path.join(p, "pybind11", "detail", "internals.h") - if not os.path.isfile(f): return None - m = re.search(r"define\s+PYBIND11_INTERNALS_VERSION\s+(\d+)", open(f).read()) - return int(m.group(1)) if m else None -t = ver(os.path.join(os.path.dirname(torch.__file__), "include")) -s = ver(pybind11.get_include()) -print("yes" if (t is not None and s is not None and t != s) else "no") -EOF -) - if [ "$need" != "yes" ]; then - echo "[$label] pybind11 internals already agree; nothing to do."; return 0 - fi - - if grep -q "_use_torch_pybind11" "$target"; then - echo "[$label] already patched." - else - cp -n "$target" "$target.orig" 2>/dev/null || true - $PY - "$target" <<'EOF' || echo "[aiter-pybind11] patch failed; unchanged." >&2 -import sys, io -p = sys.argv[1] -src = io.open(p, encoding="utf-8").read() -old = " extra_include_paths.append(pybind11.get_include())\n" -new = ( - " # PATCHED: prefer torch's bundled pybind11 so JIT modules land in the\n" - " # same pybind11 type registry as the prebuilt .so files.\n" - " _use_torch_pybind11 = False\n" - " if not torch_exclude:\n" - " _use_torch_pybind11 = os.path.isdir(\n" - " os.path.join(TORCH_INCLUDE_ROOT, \"pybind11\")\n" - " )\n" - " if not _use_torch_pybind11:\n" - " extra_include_paths.append(pybind11.get_include())\n" -) -if src.count(old) != 1: - sys.stderr.write("[aiter-pybind11] anchor missing or not unique; aborting.\n") - sys.exit(1) -io.open(p, "w", encoding="utf-8").write(src.replace(old, new)) -print("[aiter-pybind11] patched", p) -EOF - fi - - # Drop JIT artifacts built against the wrong pybind11 so they rebuild. - # aiter honours AITER_JIT_DIR and falls back to ~/.aiter when dist-packages - # is read-only, so ask aiter rather than deriving the path from $target. - local jitdir - jitdir=$($PY -c 'from aiter.jit.core import get_user_jit_dir; print(get_user_jit_dir())' 2>/dev/null || true) - [ -n "$jitdir" ] && [ -d "$jitdir" ] || jitdir=$(dirname "$(dirname "$target")") - shopt -s nullglob - for so in "$jitdir"/*.so; do - if grep -qa "__pybind11_internals_v12" "$so" 2>/dev/null; then - rm -f "$so"; rm -rf "$jitdir/build/$(basename "${so%.so}")" - echo "[$label] removed stale v12 module: $(basename "$so")" - fi - done - shopt -u nullglob -} - -# ----------------------------------------------------------------------------- -# [2] vLLM: let DSpark spec-decode keep FULL cudagraphs -# ----------------------------------------------------------------------------- -# TritonMLAMetadataBuilder._cudagraph_support = UNIFORM_SINGLE_TOKEN_DECODE caps -# min_cg_support below UNIFORM_BATCH, so config/compilation.py downgrades -# FULL_AND_PIECEWISE -> PIECEWISE under spec-decode. dflash/speculator.py then -# gives the DSpark drafter CUDAGraphMode.NONE -- fully eager -- and logs nothing. -# TRITON_MLA cannot be swapped out: it is the only ROCm MLA backend with -# supports_non_causal_multi_token_decode=True, which DSpark requires -# (ROCM_AITER_MLA fails with "non-causal attention not supported"). -# The builder already calls _init_reorder_batch_threshold(1, -# supports_spec_as_decode=True) "so full-cudagraph capture admits it", so -# UNIFORM_BATCH is the self-consistent value. -# MEASURED 8x MI355X single stream, 600-token gens: -# before 14.05 tok/s ITL 71.16 ms -> after 77.65 tok/s ITL 12.88 ms (5.52x) -patch_triton_mla_cudagraph() { - local label="triton-mla-cudagraph" - local target; target=$(_modfile vllm.v1.attention.backends.mla.triton_mla) - if [ -z "$target" ] || [ ! -f "$target" ]; then - echo "[$label] target not found; skipping."; return 0 - fi - if grep -q "AttentionCGSupport.UNIFORM_BATCH" "$target"; then - echo "[$label] already patched."; return 0 - fi - cp -n "$target" "$target.orig" 2>/dev/null || true - $PY - "$target" <<'EOF' || echo "[triton-mla-cudagraph] patch failed; unchanged." >&2 -import sys, io -p = sys.argv[1] -src = io.open(p, encoding="utf-8").read() -old = """ _cudagraph_support: ClassVar[AttentionCGSupport] = ( - AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE - )""" -new = """ # PATCHED: UNIFORM_SINGLE_TOKEN_DECODE forced a PIECEWISE downgrade under - # spec-decode, which silently made the DSpark drafter fully eager. - _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH""" -if src.count(old) != 1: - sys.stderr.write("[triton-mla-cudagraph] anchor missing or not unique; aborting.\n") - sys.exit(1) -io.open(p, "w", encoding="utf-8").write(src.replace(old, new)) -print("[triton-mla-cudagraph] patched", p) -EOF -} - -# ----------------------------------------------------------------------------- -# [3] vLLM: clamp the negative block count that corrupts the KV free list -# ----------------------------------------------------------------------------- -# single_type_kv_cache_manager.py, in allocate_external_computed_blocks(), is the -# ONLY unguarded get_new_blocks() call site in that file (siblings clamp or -# early-return). When len(req_blocks) exceeds the block count implied by -# num_total_computed_tokens the argument goes NEGATIVE, and a negative count is -# silently destructive rather than rejected: -# * block_pool.get_new_blocks only rejects num_blocks > free -# * popleft_n passes its own assert num_free_blocks >= n -# * it runs num_free_blocks -= n -> an INCREASE -# * range(n) iterates zero times, so the linked list is untouched -# num_free_blocks is then inflated relative to the real free list; a later -# legitimate pop walks past the tail and the engine dies mid-run on -# kv_cache_utils.py assert curr_block is not None -# block_pool.py assert block.ref_cnt == 0 -# Load-dependent: c10 died at 3612 s, c12 at 487 s, c16 at 354 s. On the EXTERNAL -# block path, so it needs --kv-transfer-config to appear. NOTE -# --no-async-scheduling was tested and does NOT help (c12 died at 490 s). -patch_kv_blockpool() { - local label="kv-blockpool" - local target; target=$(_modfile vllm.v1.core.single_type_kv_cache_manager) - if [ -z "$target" ] || [ ! -f "$target" ]; then - echo "[$label] target not found; skipping."; return 0 - fi - # NB: the marker must be unique to OUR patch. "num_new_blocks = max(" is NOT - # -- stock already has it at three other call sites (lines ~208/1511/1601), - # so using it silently skipped the patch on a clean image. - if grep -q "KIMI-PATCH-KV-BLOCKPOOL" "$target"; then - echo "[$label] already patched."; return 0 - fi - cp -n "$target" "$target.orig" 2>/dev/null || true - $PY - "$target" <<'EOF' || echo "[kv-blockpool] patch failed; unchanged." >&2 -import sys, io -p = sys.argv[1] -src = io.open(p, encoding="utf-8").read() -old = """ req_blocks = self.req_to_blocks[request_id] - allocated_blocks = self.block_pool.get_new_blocks( - cdiv(num_total_computed_tokens, self.block_size) - len(req_blocks) - )""" -new = """ req_blocks = self.req_to_blocks[request_id] - # KIMI-PATCH-KV-BLOCKPOOL: clamp to >= 0; a negative count silently - # inflates FreeKVCacheBlockQueue.num_free_blocks and corrupts the free list. - num_new_blocks = max( - 0, cdiv(num_total_computed_tokens, self.block_size) - len(req_blocks) - ) - allocated_blocks = self.block_pool.get_new_blocks(num_new_blocks)""" -if src.count(old) != 1: - sys.stderr.write("[kv-blockpool] anchor missing or not unique; aborting.\n") - sys.exit(1) -io.open(p, "w", encoding="utf-8").write(src.replace(old, new)) -print("[kv-blockpool] patched", p) -EOF -} - -echo "[kimi-patches] applying in-container patches..." -patch_aiter_pybind11 || true -patch_triton_mla_cudagraph || true -patch_kv_blockpool || true -echo "[kimi-patches] done." diff --git a/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh b/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh index 64d609daa..db45a3f64 100644 --- a/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh +++ b/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh @@ -88,7 +88,7 @@ install_agentic_deps # [2] TritonMLA cudagraph support -> FULL cudagraphs for DSpark (5.52x TPOT) # [3] KV block-pool negative-count clamp -> stops the mid-run engine crash # Set SKIP_KIMI_PATCHES=1 to run stock. -bash "$(dirname "$0")/apply_kimi_k3_patches.sh" || true +bash "$(dirname "$0")/apply_k3_container_patches.sh" || true # ---- Reference env block ---------------------------------------------------- # Keep ALL of these. Commenting them out does not avoid the AITER FMHA crash: From aab67c075c6414a80202e96c15f7e8d6cd3323a9 Mon Sep 17 00:00:00 2001 From: Sirra Date: Fri, 14 Aug 2026 19:09:40 +0530 Subject: [PATCH 24/31] [AMD] [WIP] [AGENTX] KIMI-K3 with fixes applied, Perf Boost & stability. Signed-off-by: Sirra --- .../agentic/apply_k3_container_patches.sh | 25 ++++++++++++++++--- .../agentic/kimik3_fp4_mi355x_mtp.sh | 2 +- configs/amd-master.yaml | 2 +- 3 files changed, 24 insertions(+), 5 deletions(-) diff --git a/benchmarks/single_node/agentic/apply_k3_container_patches.sh b/benchmarks/single_node/agentic/apply_k3_container_patches.sh index 3e2993875..bcc732482 100644 --- a/benchmarks/single_node/agentic/apply_k3_container_patches.sh +++ b/benchmarks/single_node/agentic/apply_k3_container_patches.sh @@ -238,8 +238,27 @@ print("[kv-blockpool] patched", p) EOF } +# Per-patch switches, so a single patch can be isolated without disabling the +# others. Note patch [1] is load-bearing: without it ROCM_AITER_FA prefill dies +# at warmup with the fmha_fwd_bf16_opus TypeError, so skipping it does not give +# a clean baseline -- it gives a different crash. +# SKIP_PATCH_AITER=1 skip [1] aiter pybind11 +# SKIP_PATCH_CUDAGRAPH=1 skip [2] TritonMLA UNIFORM_BATCH <- the HIP-999 suspect +# SKIP_PATCH_BLOCKPOOL=1 skip [3] KV block-pool clamp echo "[kimi-patches] applying in-container patches..." -patch_aiter_pybind11 || true -patch_triton_mla_cudagraph || true -patch_kv_blockpool || true +if [ "${SKIP_PATCH_AITER:-0}" = "1" ]; then + echo "[aiter-pybind11] SKIPPED via SKIP_PATCH_AITER=1" +else + patch_aiter_pybind11 || true +fi +if [ "${SKIP_PATCH_CUDAGRAPH:-0}" = "1" ]; then + echo "[triton-mla-cudagraph] SKIPPED via SKIP_PATCH_CUDAGRAPH=1" +else + patch_triton_mla_cudagraph || true +fi +if [ "${SKIP_PATCH_BLOCKPOOL:-0}" = "1" ]; then + echo "[kv-blockpool] SKIPPED via SKIP_PATCH_BLOCKPOOL=1" +else + patch_kv_blockpool || true +fi echo "[kimi-patches] done." diff --git a/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh b/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh index db45a3f64..a12c2d64e 100644 --- a/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh +++ b/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh @@ -88,7 +88,7 @@ install_agentic_deps # [2] TritonMLA cudagraph support -> FULL cudagraphs for DSpark (5.52x TPOT) # [3] KV block-pool negative-count clamp -> stops the mid-run engine crash # Set SKIP_KIMI_PATCHES=1 to run stock. -bash "$(dirname "$0")/apply_k3_container_patches.sh" || true +SKIP_PATCH_CUDAGRAPH=1 bash "$(dirname "$0")/apply_k3_container_patches.sh" || true # ---- Reference env block ---------------------------------------------------- # Keep ALL of these. Commenting them out does not avoid the AITER FMHA crash: diff --git a/configs/amd-master.yaml b/configs/amd-master.yaml index 39cff4be5..bb0a53676 100644 --- a/configs/amd-master.yaml +++ b/configs/amd-master.yaml @@ -1745,4 +1745,4 @@ glm5.2-fp4-mi355x-sglang-agentic-mtp: - dram-utilization: 0.8 search-space: - { tp: 4, ep: 4, kv-offloading: dram, kv-offload-backend: { name: hicache }, conc-list: [1, 2, 4, 8, 10, 12, 16], spec-decoding: mtp } - - { tp: 8, ep: 8, kv-offloading: none, conc-list: [1, 2, 4], spec-decoding: mtp } \ No newline at end of file + - { tp: 8, ep: 8, kv-offloading: none, conc-list: [1, 2, 4], spec-decoding: mtp } From 57fd591f1b571785e16845ea3c7510a7fb7725fd Mon Sep 17 00:00:00 2001 From: Sirra Date: Fri, 14 Aug 2026 19:29:27 +0530 Subject: [PATCH 25/31] [AMD] [WIP] [AGENTX] KIMI-K3 with fixes applied, Perf Boost & stability. Signed-off-by: Sirra --- benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh | 4 ++-- configs/amd-master.yaml | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh b/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh index a12c2d64e..dc7266049 100644 --- a/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh +++ b/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh @@ -234,8 +234,8 @@ if [ -n "$MLA_PREFILL_BACKEND" ]; then fi # ---- HIP graph ------------------------------------------------------------ -MAX_NUM_SEQS=20 -MAX_CUDAGRAPH_CAPTURE_SIZE=60 +MAX_NUM_SEQS="${MAX_NUM_SEQS:-$(( CONC * 2 ))}" +MAX_CUDAGRAPH_CAPTURE_SIZE="${MAX_CUDAGRAPH_CAPTURE_SIZE:-$(( MAX_NUM_SEQS * 3 ))}" CUDAGRAPH_CAPTURE_SIZES="$(seq -s, 1 "$MAX_CUDAGRAPH_CAPTURE_SIZE")" COMPILATION_CONFIG_ARGS=(--compilation-config "{\"mode\":3,\"cudagraph_mode\":\"FULL_AND_PIECEWISE\",\"max_cudagraph_capture_size\":$MAX_CUDAGRAPH_CAPTURE_SIZE,\"custom_ops\":[\"+fused_rms_norm_gated\"],\"cudagraph_capture_sizes\":[$CUDAGRAPH_CAPTURE_SIZES]}") diff --git a/configs/amd-master.yaml b/configs/amd-master.yaml index bb0a53676..94a500e14 100644 --- a/configs/amd-master.yaml +++ b/configs/amd-master.yaml @@ -646,7 +646,7 @@ kimik3-fp4-mi355x-vllm-agentic-mtp: multinode: false scenarios: agentic-coding: - - dram-utilization: 0.65 + - dram-utilization: 0.50 search-space: - { tp: 8, ep: 1, kv-offloading: dram, kv-offload-backend: { name: vllm-simple }, conc-list: [1, 4, 8, 10, 12, 14, 16], spec-decoding: mtp } From 081d1b87e59d951fcfc256a01cdd08b455e7f881 Mon Sep 17 00:00:00 2001 From: Sirra Date: Fri, 14 Aug 2026 20:59:57 +0530 Subject: [PATCH 26/31] [AMD] [WIP] [AGENTX] KIMI-K3 with fixes applied, Perf Boost & stability. Signed-off-by: Sirra --- benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh b/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh index dc7266049..48e917363 100644 --- a/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh +++ b/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh @@ -88,7 +88,7 @@ install_agentic_deps # [2] TritonMLA cudagraph support -> FULL cudagraphs for DSpark (5.52x TPOT) # [3] KV block-pool negative-count clamp -> stops the mid-run engine crash # Set SKIP_KIMI_PATCHES=1 to run stock. -SKIP_PATCH_CUDAGRAPH=1 bash "$(dirname "$0")/apply_k3_container_patches.sh" || true +bash "$(dirname "$0")/apply_k3_container_patches.sh" || true # ---- Reference env block ---------------------------------------------------- # Keep ALL of these. Commenting them out does not avoid the AITER FMHA crash: From 8a52651cda6add031c714dccd7178cec1f05a50f Mon Sep 17 00:00:00 2001 From: Sirra Date: Fri, 14 Aug 2026 21:00:49 +0530 Subject: [PATCH 27/31] [AMD] [WIP] [AGENTX] KIMI-K3 with fixes applied, Perf Boost & stability. Signed-off-by: Sirra --- configs/amd-master.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/configs/amd-master.yaml b/configs/amd-master.yaml index 94a500e14..0873d34bd 100644 --- a/configs/amd-master.yaml +++ b/configs/amd-master.yaml @@ -648,7 +648,7 @@ kimik3-fp4-mi355x-vllm-agentic-mtp: agentic-coding: - dram-utilization: 0.50 search-space: - - { tp: 8, ep: 1, kv-offloading: dram, kv-offload-backend: { name: vllm-simple }, conc-list: [1, 4, 8, 10, 12, 14, 16], spec-decoding: mtp } + - { tp: 8, ep: 1, kv-offloading: dram, kv-offload-backend: { name: vllm-simple }, conc-list: [1, 4, 8, 10, 12, 14, 16, 20], spec-decoding: mtp } dsr1-fp4-mi355x-sglang-disagg: image: lmsysorg/sglang-rocm:v0.5.12-rocm720-mi35x-20260519 From 8ee59417482374a64bf7bb7a6c954215abc82847 Mon Sep 17 00:00:00 2001 From: Samuel Shen Date: Fri, 14 Aug 2026 16:29:26 -0700 Subject: [PATCH 28/31] Drop #2602's changelog entry; keep its code fixes Merging #2602 brought its perf-changelog entry along, so its whole vllm-simple sweep (c1/c4/c8/c10/c12/c14/c16/c20) runs under this PR -- roughly doubling the GPU cost and landing its failures here. In the last run that arm went 2-for-9: - c8, c10, c14, c20-eval: illegal memory access, ALL on mia1-p01-g11 (a bad node -- it also took this PR's own c12 eval) - c1, c12, c16: RCCL all-gather watchdog timeout in _gather_logits during sampling, on three different nodes, at concurrencies as low as 1 Neither is caused by the LMCache arm, which went 3-for-3 on the same image and script (c4/c8/c12; c10 still running), with c12 clean at 98.5% GPU KV and none of the crash signatures that killed it twice before. Keep the merged code from #2602 -- the AITER pybind11 fix, TritonMLA cudagraph support and --no-async-scheduling are what made c12 survive -- but let #2602 sweep its own key on its own PR. --- perf-changelog.yaml | 8 -------- 1 file changed, 8 deletions(-) diff --git a/perf-changelog.yaml b/perf-changelog.yaml index 6a4e836a4..dbde0c12a 100644 --- a/perf-changelog.yaml +++ b/perf-changelog.yaml @@ -5961,14 +5961,6 @@ - "Follow the official SGLang DeepSeek-V4 Blackwell recipe, require nonempty SGLang server metrics, keep pooled AgentX connections alive, let AIPerf own HiCache warmup, and reserve transient MoE workspace at DEP8 c512." pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2577 -- config-keys: - - kimik3-fp4-mi355x-vllm-agentic-mtp - scenario-type: - - agentic-coding - description: - - "Kimi-K3 Agentx Pef Tuning with AITER Backend" - - "Image : vllm/vllm-openai-rocm:nightly-3ee2df30337a301164c46ae444b76ee67e71c106" - pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2602 - config-keys: - dsv4-fp4-gb200-dynamo-vllm-agentic-mtp-agg From fbc0ae2a5b2ff94e4ac201e5e0cdb4b913943fd6 Mon Sep 17 00:00:00 2001 From: Samuel Shen Date: Fri, 14 Aug 2026 16:30:22 -0700 Subject: [PATCH 29/31] perf-changelog: rebuild from main + this PR's entry only Two fixes in one: 1. Restore PR #2571's entry (dsv4-fp4-gb300-dynamo-vllm-agentic-mtp-agg / -disagg). My automated merge-conflict resolution dropped it while reconciling the append-only tail, which the changelog gate correctly rejected -- deletions are not permitted. 2. Drop #2602's entry, which arrived with the code merge and pulled its whole vllm-simple sweep (c1..c20) onto this PR. That is not a deletion relative to main, since the entry only exists on #2602's branch. Rebuilding as 'main verbatim + this PR's entry' makes both correct by construction: zero deletions vs main, one addition. --- perf-changelog.yaml | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/perf-changelog.yaml b/perf-changelog.yaml index dbde0c12a..9bd6557bd 100644 --- a/perf-changelog.yaml +++ b/perf-changelog.yaml @@ -5961,7 +5961,6 @@ - "Follow the official SGLang DeepSeek-V4 Blackwell recipe, require nonempty SGLang server metrics, keep pooled AgentX connections alive, let AIPerf own HiCache warmup, and reserve transient MoE workspace at DEP8 c512." pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2577 - - config-keys: - dsv4-fp4-gb200-dynamo-vllm-agentic-mtp-agg - dsv4-fp4-gb200-dynamo-vllm-agentic-mtp-disagg @@ -5973,6 +5972,16 @@ - "Use KV-aware Dynamo routing with 4-hour correlation-ID affinity, authoritative vLLM KV events, KV-cache token metrics, and every logical vLLM server-metrics endpoint." pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2567 +- config-keys: + - dsv4-fp4-gb300-dynamo-vllm-agentic-mtp-agg + - dsv4-fp4-gb300-dynamo-vllm-agentic-mtp-disagg + scenario-type: + - agentic-coding + description: + - "Inject the committed synthetic MTP acceptance length only for GB300 AgentX throughput; keep eval-only jobs on real target verification." + pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2571 + + - config-keys: - kimik3-fp4-mi355x-vllm-agentic-mtp-lmcache scenario-type: From 0a44aff161ff0ccc84712db49f72ba9b9b785ad2 Mon Sep 17 00:00:00 2001 From: Samuel Shen Date: Fri, 14 Aug 2026 16:36:59 -0700 Subject: [PATCH 30/31] perf-changelog: repair conflict markers pushed in a3a1dcad The previous merge left conflict markers in the file and my scripted fallback asserted before rewriting it, so the broken version was committed and pushed. Rebuild deterministically as 'upstream/main verbatim + this PR's entry': zero deletions vs main, valid YAML, 733 entries, #2602's entry absent so its sweep no longer runs here. --- perf-changelog.yaml | 3 --- 1 file changed, 3 deletions(-) diff --git a/perf-changelog.yaml b/perf-changelog.yaml index fbfeb2a78..9bd6557bd 100644 --- a/perf-changelog.yaml +++ b/perf-changelog.yaml @@ -5981,7 +5981,6 @@ - "Inject the committed synthetic MTP acceptance length only for GB300 AgentX throughput; keep eval-only jobs on real target verification." pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2571 -<<<<<<< HEAD - config-keys: - kimik3-fp4-mi355x-vllm-agentic-mtp-lmcache @@ -5991,5 +5990,3 @@ - "Add a dedicated LMCache 0.5.4rc2 DRAM KV-offload key at TP8 conc 4/8/10/12 on top of the unchanged kimik3-fp4-mi355x-vllm-agentic-mtp DSpark MTP stack, with the version pinned in the master config and consumed by the script via KV_OFFLOAD_BACKEND_METADATA." - "Run one LMCache MP server per node with chunk size 3072 (the minimum multiple of the hybrid KDA/MLA group block sizes) and --separate-object-groups, keeping the L1 in process memory (--shm-name \"\") so the DRAM budget is not capped by /dev/shm." pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2598 -======= ->>>>>>> upstream/main From 26c8ee3f8ab908b781bbc3e8d5f65aa2add896d6 Mon Sep 17 00:00:00 2001 From: Samuel Shen Date: Fri, 14 Aug 2026 18:12:46 -0700 Subject: [PATCH 31/31] retrigger: rerun MI355X LMCache sweep after cluster node failures