fix(amd): align Kimi-K3 DSpark with ROCm baseline

Use the exact GPU-only c1 DSpark serving defaults requested for MI355X: block rejection, 128 sequences, 4096 batched tokens, 0.95 GPU memory utilization, multimodal data parallelism, and no eager or prefix-cache override. Add regression coverage for the wrapper and generated matrix.

中文：将 MI355X Kimi-K3 DSpark 对齐到指定的 ROCm GPU-only c1 基线：使用 block rejection、128 个序列、4096 个批处理 token、0.95 GPU 显存利用率、multimodal data 并行，并移除 eager 与显式 prefix-cache 覆盖。同时为 wrapper 和生成矩阵添加回归测试。

Author: edwingao28

benchmarks/single_node/agentic/kimik3_fp4_mi355x.sh

@@ -49,7 +49,7 @@ set -x
 #   KV_BLOCK_SIZE            unset  (unset -> vLLM sizes the page; 128 under fp8)
 #   MAX_MODEL_LEN            unset  (unset -> vLLM derives K3's 1M context)
 #   SPEC_DECODE              false  (enabled by the _mtp DSpark wrapper)
-#   ENFORCE_EAGER            false  (enabled by the _mtp DSpark wrapper)
+#   ENFORCE_EAGER            false  (reference; DSpark wrapper keeps it false)
 
 source "$(dirname "$0")/../../benchmark_lib.sh"
 
@@ -600,7 +600,7 @@ fi
 # and hash sizes only align with prefix caching on -- an omission has been
 # reported to trip "tokens_per_block not divisible by tokens_per_hash" at load.
 # Set PREFIX_CACHING=true/false to force it either way.
-# ON by default for EVERY arm. This trace is built around large shared
+# ON by default for non-DSpark arms. This trace is built around large shared
 # prefixes -- theoretical prefix-cache hit is 98.1%, and a live kvnone cell
 # measured 92.8% server-side -- so a run with reuse disabled is not measuring
 # the workload. Reuse also costs essentially no KV (1,414,660 vs 1,420,824
@@ -620,43 +620,35 @@ fi
 #
 # Note vLLM resolves the flag's default to False for this model, so ON must be
 # passed explicitly. PREFIX_CACHING=false forces it off for a deliberate A/B.
-PREFIX_CACHE_ARGS=(--enable-prefix-caching)
-if [ "${PREFIX_CACHING:-}" = "false" ]; then
-    PREFIX_CACHE_ARGS=(--no-enable-prefix-caching)
-fi
+case "${PREFIX_CACHING:-true}" in
+    true)
+        PREFIX_CACHE_ARGS=(--enable-prefix-caching)
+        ;;
+    false)
+        PREFIX_CACHE_ARGS=(--no-enable-prefix-caching)
+        ;;
+    auto)
+        # Match the upstream AMD command by letting vLLM resolve its default.
+        PREFIX_CACHE_ARGS=()
+        ;;
+    *)
+        echo "Error: PREFIX_CACHING must be true, false, or auto." >&2
+        exit 1
+        ;;
+esac
 
 # The upstream DSpark config pins "attention_backend": "FLASHINFER_MLA", which
 # is CUDA-only and cannot be used verbatim on gfx950; SPEC_ATTN_BACKEND
-# overrides it. AgentX throughput uses synthetic acceptance pinned to the
-# committed golden curve. Eval-only runs use real block verification because
-# synthetic acceptance bypasses correctness verification.
+# overrides it. Use real block rejection for both throughput and evaluation so
+# this path is the exact AMD reproducer and exercises target-model verification.
 SPEC_ARGS=()
 if [ "${SPEC_DECODE:-false}" = "true" ]; then
     SPEC_DRAFT_MODEL="${SPEC_DRAFT_MODEL:-Inferact/Kimi-K3-DSpark}"
     SPEC_NUM_TOKENS="${SPEC_NUM_TOKENS:-7}"
     SPEC_ATTN_BACKEND="${SPEC_ATTN_BACKEND:-TRITON_MLA}"
-    SPEC_GOLDEN_AL="${SPEC_GOLDEN_AL:-3.84}"
-    SPEC_GOLDEN_AL_FILE="${SPEC_GOLDEN_AL_FILE:-golden_al_distribution/kimik3_dspark_probabilistic_sample_method_block_rejection_sample_method.yaml}"
-    FILE_GOLDEN_AL=$(awk -v k="$SPEC_NUM_TOKENS" '
-        /^kimi-k3:/                { in_model = 1; next }
-        /^[^[:space:]#]/           { in_model = 0 }
-        in_model && /thinking_on:/ { in_thinking = 1; next }
-        in_thinking && $1 == k":"  { print $2; exit }
-    ' "$SPEC_GOLDEN_AL_FILE" 2>/dev/null)
-    if [ "$FILE_GOLDEN_AL" != "$SPEC_GOLDEN_AL" ]; then
-        echo "Error: DSpark golden AL mismatch: $SPEC_GOLDEN_AL_FILE gives" >&2
-        echo "       '${FILE_GOLDEN_AL:-<unreadable>}' for k=$SPEC_NUM_TOKENS," >&2
-        echo "       but the recipe pins $SPEC_GOLDEN_AL." >&2
-        exit 1
-    fi
-    if [ "${EVAL_ONLY:-false}" = "true" ]; then
-        SPEC_REJECTION_CONFIG="\"rejection_sample_method\":\"block\""
-    else
-        SPEC_REJECTION_CONFIG="\"rejection_sample_method\":\"synthetic\",\"synthetic_acceptance_length\":$SPEC_GOLDEN_AL"
-    fi
     SPEC_ARGS=(
         --speculative-config
-        "{\"model\":\"$SPEC_DRAFT_MODEL\",\"num_speculative_tokens\":$SPEC_NUM_TOKENS,\"method\":\"dspark\",\"attention_backend\":\"$SPEC_ATTN_BACKEND\",\"draft_sample_method\":\"probabilistic\",$SPEC_REJECTION_CONFIG}"
+        "{\"model\":\"$SPEC_DRAFT_MODEL\",\"num_speculative_tokens\":$SPEC_NUM_TOKENS,\"method\":\"dspark\",\"attention_backend\":\"$SPEC_ATTN_BACKEND\",\"draft_sample_method\":\"probabilistic\",\"rejection_sample_method\":\"block\"}"
     )
 fi
 

benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh

@@ -5,16 +5,17 @@ set -euo pipefail
 # spec-decoding=mtp rows to this suffix, while the shared base recipe owns the
 # model, KV-offload, AgentX replay, and eval plumbing.
 #
-# These overrides are the DSpark + SimpleCPUOffloadConnector combination
-# validated on gfx950 in PR #2367. Keep them together: widening the batch or
-# re-enabling CUDA graphs reproduced an eight-rank GPU memory access fault.
+# Keep this wrapper aligned with the upstream AMD Kimi-K3 DSpark reproducer.
+# The first AgentX validation is deliberately GPU-only at c1 so a server or
+# kernel failure cannot be attributed to a KV-offload connector.
 export SPEC_DECODE=true
 export KV_CACHE_DTYPE="${KV_CACHE_DTYPE:-auto}"
-export MAX_NUM_SEQS="${MAX_NUM_SEQS:-32}"
-export EVAL_MAX_NUM_SEQS="${EVAL_MAX_NUM_SEQS:-32}"
-export MAX_NUM_BATCHED_TOKENS="${MAX_NUM_BATCHED_TOKENS:-1024}"
-export LANGUAGE_MODEL_ONLY="${LANGUAGE_MODEL_ONLY:-true}"
-export SIMPLE_LAZY_OFFLOAD="${SIMPLE_LAZY_OFFLOAD:-true}"
-export ENFORCE_EAGER="${ENFORCE_EAGER:-true}"
+export GPU_MEM_UTIL="${GPU_MEM_UTIL:-0.95}"
+export MAX_NUM_SEQS="${MAX_NUM_SEQS:-128}"
+export EVAL_MAX_NUM_SEQS="${EVAL_MAX_NUM_SEQS:-128}"
+export MAX_NUM_BATCHED_TOKENS="${MAX_NUM_BATCHED_TOKENS:-4096}"
+export LANGUAGE_MODEL_ONLY="${LANGUAGE_MODEL_ONLY:-false}"
+export PREFIX_CACHING="${PREFIX_CACHING:-auto}"
+export ENFORCE_EAGER="${ENFORCE_EAGER:-false}"
 
 exec "$(dirname "$0")/kimik3_fp4_mi355x.sh" "$@"

configs/amd-master.yaml

@@ -664,12 +664,10 @@ kimik3-fp4-mi355x-vllm-agentic-mtp:
     agentic-coding:
     - dram-utilization: 0.80
       search-space:
-      # DSpark speculative decoding with Inferact/Kimi-K3-DSpark, seven draft
-      # tokens, probabilistic draft sampling, and the committed thinking-on
-      # golden AL. The _mtp wrapper applies the gfx950 settings validated in
-      # PR #2367 as one unit: eager execution, max-num-batched-tokens 1024,
-      # max-num-seqs 32, bf16 KV, and lazy SimpleCPUOffloadConnector allocation.
-      - { tp: 8, kv-offloading: dram, kv-offload-backend: { name: vllm-simple }, spec-decoding: mtp, conc-list: [1, 2, 4, 8] }
+      # Exact upstream AMD DSpark reproducer: GPU-only TP8, seven draft tokens,
+      # probabilistic draft sampling, TRITON_MLA, and real block rejection.
+      # Start at c1; expand only after the AgentX-fast preflight is healthy.
+      - { tp: 8, kv-offloading: none, spec-decoding: mtp, conc-list: [1] }
 
 kimik3-fp4-mi355x-vllm-agentic-fp8kvtest:
   image: vllm/vllm-openai-rocm:kimi-k3

perf-changelog.yaml

@@ -5316,6 +5316,6 @@
     - kimik3-fp4-mi355x-vllm-agentic-mtp
   description:
     - "Add the Kimi-K3 MXFP4 MI355X AgentX DSpark variant with Inferact/Kimi-K3-DSpark, seven draft tokens, probabilistic draft sampling, and TRITON_MLA on ROCm."
-    - "Pin throughput to the committed thinking-on golden acceptance length 3.84 with synthetic rejection sampling; eval-only runs retain real block verification."
-    - "Run TP8 with SimpleCPUOffloadConnector at concurrency 1/2/4/8. The _mtp wrapper keeps the gfx950 bring-up settings validated in PR #2367 together: --enforce-eager, --max-num-batched-tokens 1024, --max-num-seqs 32, bf16 KV cache, language-model-only mode, and lazy host allocation."
+    - "Align the first AgentX validation with the upstream AMD reproducer: GPU-only TP8 at concurrency 1, real block rejection, --gpu-memory-utilization 0.95, --max-num-batched-tokens 4096, --max-num-seqs 128, multimodal encoder data parallelism, and no --enforce-eager or explicit prefix-cache override."
+    - "Use VLLM_ROCM_USE_AITER=1, SAFETENSORS_FAST_GPU=1, AITER_SITUV2_A8W4=1, AITER_BF16_FP8_MOE_BOUND=0, and VLLM_USE_BREAKABLE_CUDAGRAPH=0 with the stock vllm/vllm-openai-rocm:kimi-k3 image."
   pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2403

utils/matrix_logic/test_kimik3_mi355x_dspark.py

@@ -0,0 +1,70 @@
+from pathlib import Path
+import re
+
+import yaml
+
+
+REPO_ROOT = Path(__file__).resolve().parents[2]
+MTP_WRAPPER = (
+    REPO_ROOT / "benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh"
+)
+BASE_RECIPE = REPO_ROOT / "benchmarks/single_node/agentic/kimik3_fp4_mi355x.sh"
+MASTER_CONFIG = REPO_ROOT / "configs/amd-master.yaml"
+
+
+def _shell_default(script: str, variable: str) -> str:
+    match = re.search(
+        rf'export {re.escape(variable)}="\$\{{{re.escape(variable)}:-(.*?)\}}"',
+        script,
+    )
+    assert match is not None, f"{variable} must have an overridable default"
+    return match.group(1)
+
+
+def test_dspark_wrapper_matches_oren_rocm_baseline() -> None:
+    wrapper = MTP_WRAPPER.read_text()
+
+    assert _shell_default(wrapper, "KV_CACHE_DTYPE") == "auto"
+    assert _shell_default(wrapper, "GPU_MEM_UTIL") == "0.95"
+    assert _shell_default(wrapper, "MAX_NUM_SEQS") == "128"
+    assert _shell_default(wrapper, "EVAL_MAX_NUM_SEQS") == "128"
+    assert _shell_default(wrapper, "MAX_NUM_BATCHED_TOKENS") == "4096"
+    assert _shell_default(wrapper, "LANGUAGE_MODEL_ONLY") == "false"
+    assert _shell_default(wrapper, "PREFIX_CACHING") == "auto"
+    assert _shell_default(wrapper, "ENFORCE_EAGER") == "false"
+
+
+def test_dspark_defaults_to_real_block_rejection() -> None:
+    recipe = BASE_RECIPE.read_text()
+    spec_section = recipe.split("# ---- Eval-only path", maxsplit=1)[0].split(
+        "SPEC_ARGS=()", maxsplit=1
+    )[1]
+
+    assert '\\"rejection_sample_method\\":\\"block\\"' in spec_section
+    assert "synthetic_acceptance_length" not in spec_section
+
+
+def test_reference_prefix_cache_mode_emits_no_override() -> None:
+    recipe = BASE_RECIPE.read_text()
+    prefix_section = recipe.split("# The upstream DSpark config", maxsplit=1)[0].split(
+        "PREFIX_CACHE_ARGS=", maxsplit=1
+    )[1]
+
+    assert "auto)" in prefix_section
+    assert "PREFIX_CACHE_ARGS=()" in prefix_section
+
+
+def test_dspark_diagnostic_matrix_is_gpu_only_c1() -> None:
+    master = yaml.safe_load(MASTER_CONFIG.read_text())
+    search_space = master["kimik3-fp4-mi355x-vllm-agentic-mtp"]["scenarios"][
+        "agentic-coding"
+    ][0]["search-space"]
+
+    assert search_space == [
+        {
+            "tp": 8,
+            "kv-offloading": "none",
+            "spec-decoding": "mtp",
+            "conc-list": [1],
+        }
+    ]