From f2968297fcd26725a55ef1f84ed0bce902874f32 Mon Sep 17 00:00:00 2001 From: Sirra Date: Thu, 13 Aug 2026 12:39:37 +0530 Subject: [PATCH 01/18] [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 02/18] [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 03/18] [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 c0119e6be379f35102386fccd5500e90398dd3ad Mon Sep 17 00:00:00 2001 From: Sirra Date: Fri, 14 Aug 2026 07:51:03 +0530 Subject: [PATCH 04/18] [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 05/18] [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 06/18] [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 07/18] [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 08/18] [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 09/18] [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 10/18] [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 11/18] [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 12/18] [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 13/18] [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 14/18] [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 9a9654752f3217248c7c242097c28436272324d9 Mon Sep 17 00:00:00 2001 From: Sirra Date: Sat, 15 Aug 2026 07:23:12 +0530 Subject: [PATCH 15/18] [AMD] [WIP] [AGENTX] KIMI-K3 with fixes applied, Perf Boost & stability. Signed-off-by: Sirra --- configs/amd-master.yaml | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/configs/amd-master.yaml b/configs/amd-master.yaml index 2bfa094c6..8d164e32d 100644 --- a/configs/amd-master.yaml +++ b/configs/amd-master.yaml @@ -646,9 +646,10 @@ kimik3-fp4-mi355x-vllm-agentic-mtp: multinode: false scenarios: agentic-coding: - - dram-utilization: 0.50 + - dram-utilization: 0.30 search-space: - - { 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 } + - { tp: 8, ep: 1, kv-offloading: none, conc-list: [1, 4], spec-decoding: mtp } + - { tp: 8, ep: 1, kv-offloading: dram, kv-offload-backend: { name: vllm-simple }, conc-list: [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 d820ae1f2d4f150357f6a82a92bbfdfa1bf881f6 Mon Sep 17 00:00:00 2001 From: Sirra Date: Sat, 15 Aug 2026 07:31:32 +0530 Subject: [PATCH 16/18] [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, 2 insertions(+) diff --git a/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh b/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh index 48e917363..a69a9eb00 100644 --- a/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh +++ b/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh @@ -46,6 +46,8 @@ source "$(dirname "$0")/../../benchmark_lib.sh" wait_for_amd_gpu_clean +du -h /dev/shm + check_env_vars MODEL TP CONC KV_OFFLOADING TOTAL_CPU_DRAM_GB RESULT_DIR DURATION EP_SIZE if [[ -n "${SLURM_JOB_ID:-}" ]]; then From 25d1e54546f143d5e7246b0bcec3fa3524b3e4fd Mon Sep 17 00:00:00 2001 From: Sirra Date: Sat, 15 Aug 2026 07:31:59 +0530 Subject: [PATCH 17/18] [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 8d164e32d..7dc0f60a8 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.30 + - dram-utilization: 0.40 search-space: - { tp: 8, ep: 1, kv-offloading: none, conc-list: [1, 4], spec-decoding: mtp } - { tp: 8, ep: 1, kv-offloading: dram, kv-offload-backend: { name: vllm-simple }, conc-list: [8, 10, 12, 14, 16, 20], spec-decoding: mtp } From 8988497f2c4a3404b93d623b68766b9b615ce519 Mon Sep 17 00:00:00 2001 From: Sirra Date: Sat, 15 Aug 2026 07:42:51 +0530 Subject: [PATCH 18/18] [AMD] [WIP] [AGENTX] KIMI-K3 with fixes applied, Perf Boost & stability. Signed-off-by: Sirra --- runners/launch_mi355x-amds.sh | 2 ++ 1 file changed, 2 insertions(+) diff --git a/runners/launch_mi355x-amds.sh b/runners/launch_mi355x-amds.sh index 49c3af0cf..cce79befb 100644 --- a/runners/launch_mi355x-amds.sh +++ b/runners/launch_mi355x-amds.sh @@ -313,6 +313,8 @@ else BENCHMARK_SCRIPT="$SCRIPT_FALLBACK" fi + srun --jobid=$JOB_ID bash -c "du -h /dev/shm" || true + srun --jobid=$JOB_ID \ --container-image=$SQUASH_FILE \ --container-mounts=$GITHUB_WORKSPACE:/workspace/,$HF_HUB_CACHE_MOUNT:$HF_HUB_CACHE,$AIPERF_MMAP_CACHE_HOST_PATH:/aiperf_mmap_cache \