[AMD][AgentX] Kimi-K3 FP4 MI355X agentX vLLM

.github/workflows/e2e-tests.yml

@@ -310,14 +310,15 @@ jobs:
             dp-attn: ${{ matrix.config.dp-attn }}
             conc: ${{ matrix.config.conc }}
             kv-offloading: ${{ matrix.config.kv-offloading }}
-            kv-offload-backend: ${{ matrix.config.kv-offload-backend }}
+            kv-offload-backend: ${{ matrix.config['kv-offload-backend'].name }}
+            kv-offload-backend-metadata: ${{ matrix.config['kv-offload-backend'] && toJson(matrix.config['kv-offload-backend']) || '' }}
             total-cpu-dram-gb: ${{ matrix.config.total-cpu-dram-gb }}
             duration: ${{ inputs.agentx-fast && '1200' || (inputs.duration-override != '' && inputs.duration-override || matrix.config.duration) }}
             agentx-fast: ${{ inputs.agentx-fast }}
             isl: '0'
             osl: '0'
             max-model-len: '0'
-            spec-decoding: 'none'
+            spec-decoding: ${{ matrix.config.spec-decoding }}
             disagg: ${{ 'false' }}
             run-eval: true
             eval-only: true

benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh

@@ -0,0 +1,21 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# DSpark variant of kimik3_fp4_mi355x.sh. The MI355X launcher routes
+# spec-decoding=mtp rows to this suffix, while the shared base recipe owns the
+# model, KV-offload, AgentX replay, and eval plumbing.
+#
+# 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 GPU_MEM_UTIL="${GPU_MEM_UTIL:-0.95}"
+export MAX_NUM_SEQS="${MAX_NUM_SEQS:-16}"
+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" "$@"

perf-changelog.yaml

@@ -5340,6 +5340,14 @@
     - "Validated on the node through the real launcher, every request successful: TEP4 concurrency 64 640/640 at 6895 tok/s (11% above the non-MTP arm's 6198 tok/s, mean TTFT 27.2s to 9.6s), and TEP4 concurrency 1 10/10 at 1233 tok/s (83% above the non-MTP arm's 672 tok/s, mean TPOT 12.6ms to 6.5ms)"
   pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2312
 
+- config-keys:
+    - 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."
+    - "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 16, 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
+
 - config-keys:
     - kimik2.6-fp4-b200-dynamo-vllm
   description:

runners/launch_mi355x-amds.sh

@@ -268,7 +268,36 @@ else
     export GPU_COUNT="${GPU_COUNT:-${TP:?TP must be set}}"
 
     set -x
-    salloc --partition=$PARTITION --gres=gpu:$GPU_COUNT --exclusive --cpus-per-task=128 --time=500 --no-shell --job-name="$RUNNER_NAME"
+    # Node denylist. mia1-p01-g14 chronically carries a stale ~32 GiB/GPU
+    # allocation from previous tenants, so vLLM refuses to start on it: only
+    # ~256 of 287.98 GiB is free and any gpu-memory-utilization above ~0.888
+    # trips "Free memory on device cuda:N ... is less than desired GPU memory
+    # utilization". It took out kimik3 cells in runs 30330955808 and
+    # 30331297999, and the same node did the same thing to the kimik2.7 -tune2
+    # cells.
+    #
+    # A BROAD denylist did not work: after excluding g11/g14/g15/g18, run
+    # 30412966635 still lost all three cells -- to g19 (hipErrorIllegalAddress),
+    # g17 (warmup_failure) and g16 (srun node failure), none of them on the
+    # list. Free memory on these nodes varies hour to hour, so a wide static
+    # list mostly just shrinks an already-contended pool and lengthens queue
+    # time. Don't reintroduce one.
+    #
+    # g11 is the exception and is excluded by default. It is not a
+    # free-memory-headroom case; it fails hard and across unrelated configs:
+    #   - kimik3 kvnone c2, run 30425820793: warmup_failure
+    #   - kimik3 vllm-simple-fp8 c8, run 30428781263: forward passes stalled
+    #     (returned=5 for 13 min at zero throughput), then EngineCore died
+    #   - plus the earlier kimik2.7 -tune2 cells noted above
+    # Four configs, one node, same outcome. Set SALLOC_EXCLUDE_NODES="" to
+    # clear it, or to a comma-separated list to widen it for a specific run.
+    SALLOC_EXCLUDE_NODES="${SALLOC_EXCLUDE_NODES-mia1-p01-g11}"
+    EXCLUDE_ARG=""
+    if [ -n "$SALLOC_EXCLUDE_NODES" ]; then
+        EXCLUDE_ARG="--exclude=$SALLOC_EXCLUDE_NODES"
+        echo "salloc excluding nodes: $SALLOC_EXCLUDE_NODES"
+    fi
+    salloc --partition=$PARTITION --gres=gpu:$GPU_COUNT --exclusive --cpus-per-task=128 --time=500 --no-shell --job-name="$RUNNER_NAME" $EXCLUDE_ARG
     JOB_ID=$(squeue --name="$RUNNER_NAME" -h -o %A | head -n1)
 
     srun --jobid=$JOB_ID bash -c "docker stop \$(docker ps -a -q)"
@@ -306,6 +335,14 @@ else
         export HF_HUB_CACHE_MOUNT="/it-share/hf-hub-cache/"
     fi
 
+    # Kimi-K3's MXFP4 checkpoint is 1.561 TB decimal (1.420 TiB, 96
+    # safetensors) and does not fit the node-local /var/lib NVMe hub cache
+    # alongside anything else, so it is pre-staged once on the NFS share like
+    # MiniMax-M3.
+    if [[ "$MODEL" == moonshotai/Kimi-K3* ]]; then
+        export HF_HUB_CACHE_MOUNT="/it-share/hf-hub-cache/"
+    fi
+
     SCRIPT_BASE="${EXP_NAME%%_*}_${PRECISION}_mi355x"
     SCRIPT_FW="benchmarks/single_node/${SCENARIO_SUBDIR:-fixed_seq_len/}${SCRIPT_BASE}_${FRAMEWORK}${SPEC_SUFFIX}.sh"
     SCRIPT_FALLBACK="benchmarks/single_node/${SCENARIO_SUBDIR:-fixed_seq_len/}${SCRIPT_BASE}${FRAMEWORK_SUFFIX}${SPEC_SUFFIX}.sh"

utils/changelog_gate_tests/test_run_sweep_gating.py

@@ -409,6 +409,23 @@ def test_e2e_workflow_cannot_dispatch_database_ingest() -> None:
     assert "INFX_FRONTEND_PAT" not in workflow
 
 
+def test_e2e_agentic_eval_preserves_matrix_transport_settings() -> None:
+    workflow = yaml.load(
+        (REPO_ROOT / ".github/workflows/e2e-tests.yml").read_text(),
+        Loader=yaml.BaseLoader,
+    )
+    inputs = workflow["jobs"]["test-sweep-agentic-evals"]["with"]
+
+    assert inputs["kv-offload-backend"] == (
+        "${{ matrix.config['kv-offload-backend'].name }}"
+    )
+    assert inputs["kv-offload-backend-metadata"] == (
+        "${{ matrix.config['kv-offload-backend']"
+        " && toJson(matrix.config['kv-offload-backend']) || '' }}"
+    )
+    assert inputs["spec-decoding"] == "${{ matrix.config.spec-decoding }}"
+
+
 def test_priority_classifier_runs_for_enabled_actions() -> None:
     scenario = {
         **_PR,

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") == "16"
+    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],
+        }
+    ]