[AMD] [AGENTX] Kimi-K3 Performance Tuning

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."

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."

benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh

@@ -81,7 +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.
 export VLLM_ROCM_AITER_MLA_ASM_PADDING=asm
 export VLLM_ROCM_USE_AITER=1
 export SAFETENSORS_FAST_GPU=1
@@ -147,7 +171,7 @@ esac
 fi
 
 # ---- LLM server  ------------------------------------------------------------
-bash "$(dirname "$0")/apply_k3_container_patches.sh"
+
 
 # ---- Parallelism ------------------------------------------------------------
 EP_ARGS=()
@@ -171,6 +195,31 @@ 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.
+# 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=(
+        --attention-config
+        "{\"mla_prefill_backend\":\"$MLA_PREFILL_BACKEND\"}"
+    )
+fi
+
 # ---- HIP graph ------------------------------------------------------------
 MAX_NUM_SEQS=20
 MAX_CUDAGRAPH_CAPTURE_SIZE=60
@@ -202,6 +251,7 @@ VLLM_CMD=(
     --max-model-len 1048576
     --enable-prefix-caching
     --kv-cache-dtype "fp8"
+    "${MLA_PREFILL_ARGS[@]}"
     "${COMPILATION_CONFIG_ARGS[@]}"
     "${SPEC_ARGS[@]}"
     "${OFFLOAD_ARGS[@]}"

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: