Add vLLM ROCm 0.23 kernel tracing support

README.md

@@ -403,13 +403,13 @@ pytest tests/test_graders.py -v     # Grader tests only
 The pipeline runs E2E benchmarks inside Docker containers. Override the vLLM image:
 
 ```bash
---docker-image vllm/vllm-openai-rocm:v0.19.0
+--docker-image vllm/vllm-openai-rocm:v0.23.0
 ```
 
 Or set the environment variable:
 
 ```bash
-export APEX_VLLM_ROCM_IMAGE=vllm/vllm-openai-rocm:v0.19.0
+export APEX_VLLM_ROCM_IMAGE=vllm/vllm-openai-rocm:v0.23.0
 ```
 
 ### Benchmark Caching

graders/score.py

@@ -320,7 +320,7 @@ def _detect_run_mode() -> str:
         return "local"
 
 
-VLLM_ROCM_IMAGE_DEFAULT = "vllm/vllm-openai-rocm:v0.19.0"
+VLLM_ROCM_IMAGE_DEFAULT = "vllm/vllm-openai-rocm:v0.23.0"
 
 
 def _ensure_docker_image(

pipeline/kernel_tracing/README.md

@@ -12,7 +12,7 @@ This feature is meant to answer one practical question:
 python3 workload_optimizer.py trace-kernel \
   -r /path/to/results_trace \
   --kernel-id aiter.triton.unified_attention_2d \
-  --docker-image vllm/vllm-openai-rocm:v0.19.1 \
+  --docker-image vllm/vllm-openai-rocm:v0.23.0 \
   --max-records 200 \
   --sample-rate 1.0 \
   -b /root/Magpie/examples/benchmarks/benchmark_vllm_gptoss_20b.yaml
@@ -22,9 +22,9 @@ List supported IDs with:
 
 ```bash
 python3 workload_optimizer.py list-trace-kernels \
-  --docker-image vllm/vllm-openai-rocm:v0.19.1
+  --docker-image vllm/vllm-openai-rocm:v0.23.0
 python3 workload_optimizer.py list-trace-kernels \
-  --docker-image vllm/vllm-openai-rocm:v0.19.1 \
+  --docker-image vllm/vllm-openai-rocm:v0.23.0 \
   --repo vllm --kernel-type hip
 python3 workload_optimizer.py list-trace-kernels --supported-images
 ```
@@ -42,13 +42,15 @@ python3 workload_optimizer.py trace-kernel \
   -b /root/Magpie/examples/benchmarks/benchmark_sglang_dsr1.yaml
 ```
 
-For the DeepSeek R1 SGLang multi-kernel shape-analysis workflow, use the
-checked-in example script and override the benchmark config as needed:
+Checked-in Kimi K2.6 and DeepSeek R1 examples pair each trace script with one
+benchmark workload under `pipeline/kernel_tracing/examples/`:
 
 ```bash
-BENCH_CONFIG=tmp/benchmark_sglang_dsr1_isl8192_osl1024.yaml \
-RESULTS_DIR=tmp/results_trace_deepseek_r1_isl8192_osl1024_example \
-bash pipeline/kernel_tracing/trace_deepseek_r1_multi_kernels_example.sh
+HF_CACHE_PATH=/path/to/huggingface \
+  bash pipeline/kernel_tracing/examples/trace_kimi_k2_6_mxfp4_multi_kernels.sh
+
+HF_CACHE_PATH=/path/to/huggingface \
+  bash pipeline/kernel_tracing/examples/trace_deepseek_r1_1k_multi_kernels.sh
 ```
 
 `trace-kernel` has two phases:
@@ -190,6 +192,12 @@ APEX_TRACE_OUTPUT_DIR=/apex_trace/trace_raw
 
 The importer records a `module_import` event whenever a patched overlay module is imported. This diagnostic event is never sampled and does not count against `--max-records`.
 
+The generated `benchmark/trace_benchmark_config.yaml` also pins
+`benchmark.docker_image` to the exact image selected for the trace registry.
+This overrides an older or mutable image tag in the input benchmark config
+without modifying the original YAML, keeping registry discovery, source
+extraction, and the executed workload on the same image.
+
 ## Interpreting Results
 
 ### Case 1: `module_import` and target events exist
@@ -232,7 +240,7 @@ Use this workflow for a new target:
 
 Kernel tracing intentionally supports only these Docker images:
 
-- `vllm/vllm-openai-rocm:v0.19.1`
+- `vllm/vllm-openai-rocm:v0.23.0`
 - `lmsysorg/sglang:v0.5.12-rocm720-mi35x`
 
 Each image has its own generated registry under
@@ -257,10 +265,11 @@ python3 workload_optimizer.py update-trace-kernel-registry --write
 The generator copies Python source from the two supported images, discovers
 traceable launch/wrapper sites, records image/package provenance, and writes:
 
-- `pipeline/kernel_tracing/registries/vllm_v0_19_1.yaml`
+- `pipeline/kernel_tracing/registries/vllm_v0_23_0.yaml`
 - `pipeline/kernel_tracing/registries/sglang_v0_5_12_rocm720_mi35x.yaml`
 
-Custom, nightly, digest-only, or other architecture images are not accepted by
+The former vLLM `v0.19.1` image, custom tags, nightly, digest-only, or other
+architecture images are not accepted by
 `trace-kernel`. During Docker tracing, Apex extracts the source file from the
 resolved supported image and fails hard if the file cannot be found, rather than
 falling back to a possibly mismatched local checkout.
@@ -284,7 +293,7 @@ This traces a known aiter Triton launch used by GPT-OSS 20B when vLLM routes att
 python3 workload_optimizer.py trace-kernel \
   -r /root/Apex/results_trace_gptoss20b_aiter_unified_attention_2d \
   --kernel-id aiter.triton.unified_attention_2d \
-  --docker-image vllm/vllm-openai-rocm:v0.19.1 \
+  --docker-image vllm/vllm-openai-rocm:v0.23.0 \
   --max-records 200 \
   --sample-rate 1.0 \
   --benchmark-timeout 2700 \
@@ -305,7 +314,7 @@ This traces the Python-visible vLLM cache wrapper before it calls the compiled c
 python3 workload_optimizer.py trace-kernel \
   -r /root/Apex/results_trace_gptoss20b_vllm_reshape_and_cache_flash \
   --kernel-id vllm.hip.reshape_and_cache_flash \
-  --docker-image vllm/vllm-openai-rocm:v0.19.1 \
+  --docker-image vllm/vllm-openai-rocm:v0.23.0 \
   --max-records 200 \
   --sample-rate 1.0 \
   --benchmark-timeout 2700 \
@@ -325,7 +334,7 @@ Use this when you know a high-level aiter path is involved, but you do not yet k
 python3 workload_optimizer.py trace-kernel \
   -r /root/Apex/results_trace_gptoss20b_aiter_compile_ops_discovery \
   --kernel-id aiter.hip.fmoe \
-  --docker-image vllm/vllm-openai-rocm:v0.19.1 \
+  --docker-image vllm/vllm-openai-rocm:v0.23.0 \
   --trace-all \
   --max-records 200 \
   --sample-rate 1.0 \
@@ -356,7 +365,7 @@ Then rerun with the actual supported low-level `--kernel-id`, for example:
 python3 workload_optimizer.py trace-kernel \
   -r /root/Apex/results_trace_gptoss20b_aiter_fmoe \
   --kernel-id aiter.hip.fmoe \
-  --docker-image vllm/vllm-openai-rocm:v0.19.1 \
+  --docker-image vllm/vllm-openai-rocm:v0.23.0 \
   --max-records 10000 \
   --sample-rate 0.01 \
   --benchmark-timeout 2700 \
@@ -371,7 +380,7 @@ Use `--dry-run` to validate that the patch can be generated and compiled without
 python3 workload_optimizer.py trace-kernel \
   -r /tmp/apex_trace_dry_unified_attention \
   --kernel-id aiter.triton.unified_attention_2d \
-  --docker-image vllm/vllm-openai-rocm:v0.19.1 \
+  --docker-image vllm/vllm-openai-rocm:v0.23.0 \
   --dry-run
 ```
 
@@ -387,7 +396,7 @@ select the fixed registry that defines the supported kernel ID.
 python3 workload_optimizer.py trace-kernel \
   -r /root/Apex/results_trace_local_pa_decode \
   --kernel-id aiter.triton.paged_attn_decode_v1_wo_dot \
-  --docker-image vllm/vllm-openai-rocm:v0.19.1 \
+  --docker-image vllm/vllm-openai-rocm:v0.23.0 \
   --max-records 200 \
   --sample-rate 1.0 \
   --run-cmd 'python3 -m pytest /root/Apex/tools/rocm/aiter/op_tests/triton_tests/attention/test_pa_decode.py -x'

pipeline/kernel_tracing/examples/benchmark_sglang_dsr1_isl1024_osl1024.yaml

@@ -0,0 +1,38 @@
+# SGLang AMD DeepSeek-R1 MXFP4 trace benchmark: 1K input, 1K output.
+#
+# Run it through the paired trace_deepseek_r1_1k_multi_kernels.sh script.
+# Export HF_TOKEN in the host environment if the model requires authentication;
+# Magpie forwards it to the benchmark container without storing it in this file.
+
+benchmark:
+  framework: sglang
+  model: amd/DeepSeek-R1-MXFP4
+  precision: fp4
+  run_mode: docker
+  docker_image: "lmsysorg/sglang:v0.5.12-rocm720-mi35x"
+
+  envs:
+    TP: 8
+    CONC: 64
+    ISL: 1024
+    OSL: 1024
+    RANDOM_RANGE_RATIO: 0.8
+    PROFILER_ENABLED: false
+
+  profiler:
+    torch_profiler:
+      enabled: false
+    system_profiler:
+      enabled: false
+      profile_args: []
+    tracelens:
+      enabled: false
+
+  gap_analysis:
+    enabled: false
+    find_kernel_sources: false
+    top_k: 100
+
+  benchmark_script: "single_node/fixed_seq_len/dsr1_fp4_mi355x.sh"
+  hf_cache_path: "${HF_CACHE_PATH}"
+  timeout_seconds: 3600

pipeline/kernel_tracing/examples/benchmark_vllm_amd_kimi_k2_6_mxfp4.yaml

@@ -0,0 +1,47 @@
+# vLLM Benchmark Configuration for amd/Kimi-K2.6-MXFP4
+# ====================================================
+#
+# Usage (Docker - default):
+#   HF_CACHE_PATH=/path/to/huggingface \
+#     python -m Magpie benchmark --benchmark-config \
+#     pipeline/kernel_tracing/examples/benchmark_vllm_amd_kimi_k2_6_mxfp4.yaml
+#
+# Usage (Local - run directly on host, e.g. inside a pod):
+#   HF_CACHE_PATH=/path/to/huggingface \
+#     python -m Magpie benchmark --benchmark-config \
+#     pipeline/kernel_tracing/examples/benchmark_vllm_amd_kimi_k2_6_mxfp4.yaml \
+#     --run-mode local
+
+benchmark:
+  framework: vllm
+  model: amd/Kimi-K2.6-MXFP4
+  precision: mxfp4
+
+  envs:
+    TP: 4
+    CONC: 32
+    ISL: 8192
+    OSL: 1024
+    RANDOM_RANGE_RATIO: 0.8
+    MAX_MODEL_LEN: 131072
+    VLLM_ROCM_USE_AITER: "1"
+    EXTRA_VLLM_ARGS: "--block-size 1 --enable-auto-tool-choice --tool-call-parser kimi_k2"
+
+  profiler:
+    torch_profiler:
+      enabled: false
+    system_profiler:
+      enabled: false
+    tracelens:
+      enabled: false
+
+  docker_image: "vllm/vllm-openai-rocm:v0.23.0"
+  # Magpie expands environment variables before creating the cache bind mount.
+  hf_cache_path: "${HF_CACHE_PATH}"
+  benchmark_script: "vllm_mi355x.sh"
+  timeout_seconds: 3600
+
+  gap_analysis:
+    enabled: false
+    top_k: 100
+    find_kernel_sources: false

pipeline/kernel_tracing/trace_deepseek_r1_multi_kernels_example.sh → pipeline/kernel_tracing/examples/trace_deepseek_r1_1k_multi_kernels.sh

@@ -1,18 +1,18 @@
 #!/usr/bin/env bash
 set -euo pipefail
 
-# Example DeepSeek R1 multi-kernel tracing run.
+# DeepSeek R1 MXFP4 1K/1K multi-kernel tracing example.
 #
-# Override the workload without editing this file:
-#   BENCH_CONFIG=tmp/benchmark_sglang_dsr1_isl8192_osl1024.yaml \
-#   RESULTS_DIR=tmp/results_trace_deepseek_r1_isl8192_osl1024_example \
-#   bash pipeline/kernel_tracing/trace_deepseek_r1_multi_kernels_example.sh
+# Basic use:
+#   HF_CACHE_PATH=/path/to/huggingface \
+#     bash pipeline/kernel_tracing/examples/trace_deepseek_r1_1k_multi_kernels.sh
 #
 # Common knobs:
-#   MAX_RECORDS=20000 SAMPLE_RATE=1.0 BENCHMARK_TIMEOUT=7200
+#   MAX_RECORDS=300000 SAMPLE_RATE=0.01 BENCHMARK_TIMEOUT=7200
 #   TRACE_ALL=1 DRY_RUN=1 DISABLE_BENCHMARK_CUDA_GRAPH=0
 
-APEX_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+APEX_ROOT="$(cd "$SCRIPT_DIR/../../.." && pwd)"
 cd "$APEX_ROOT"
 
 if [[ -f "$APEX_ROOT/.venv/bin/activate" ]]; then
@@ -31,21 +31,27 @@ if [[ -z "${MAGPIE_ROOT:-}" ]]; then
   fi
 fi
 
+export HF_CACHE_PATH="${HF_CACHE_PATH:-${HF_HOME:-$HOME/.cache/huggingface}}"
 export MAGPIE_RUN_MODE="${MAGPIE_RUN_MODE:-docker}"
 
 TS="$(date -u +%Y%m%d_%H%M%S)"
-DEFAULT_RESULTS_DIR="$APEX_ROOT/tmp/results_trace_deepseek_r1_multi_kernels_${TS}"
+DEFAULT_RESULTS_DIR="$APEX_ROOT/tmp/results_trace_deepseek_r1_1k_multi_kernels_${TS}"
 RESULTS_DIR="${RESULTS_DIR:-$DEFAULT_RESULTS_DIR}"
 mkdir -p "$RESULTS_DIR"
 
-# Default tmp config is aligned with the DeepSeek R1 SGLang FP4 fixed-seq run.
-DEFAULT_TRACE_BENCH="$APEX_ROOT/tmp/benchmark_sglang_dsr1_isl1024_osl1024.yaml"
+# Keep the script and its one supported workload paired in this directory.
+DEFAULT_TRACE_BENCH="$SCRIPT_DIR/benchmark_sglang_dsr1_isl1024_osl1024.yaml"
 TRACE_BENCH="${BENCH_CONFIG:-$DEFAULT_TRACE_BENCH}"
 if [[ ! -f "$TRACE_BENCH" ]]; then
   echo "Benchmark config not found: $TRACE_BENCH" >&2
   exit 1
 fi
 
+DOCKER_IMAGE="${DOCKER_IMAGE:-lmsysorg/sglang:v0.5.12-rocm720-mi35x}"
+MAX_RECORDS="${MAX_RECORDS:-300000}"
+SAMPLE_RATE="${SAMPLE_RATE:-0.01}"
+BENCHMARK_TIMEOUT="${BENCHMARK_TIMEOUT:-7200}"
+
 # These cover the CSV analysis targets that are present in the fixed-image
 # trace registries under pipeline/kernel_tracing/registries/.
 # Keep this list explicit: different workload/model analyses often need
@@ -120,6 +126,8 @@ echo "Apex root:       $APEX_ROOT"
 echo "Magpie root:     $MAGPIE_ROOT"
 echo "Results dir:     $RESULTS_DIR"
 echo "Benchmark config:$TRACE_BENCH"
+echo "Docker image:    $DOCKER_IMAGE"
+echo "HF cache path:   $HF_CACHE_PATH"
 echo "Run mode:        $MAGPIE_RUN_MODE"
 echo "Framework:       ${FRAMEWORK:-sglang}"
 echo "CUDA graph:      disable overlay=${DISABLE_BENCHMARK_CUDA_GRAPH:-1}"
@@ -133,9 +141,10 @@ python3 workload_optimizer.py trace-kernel \
   "${TRACE_ALL_ARGS[@]}" \
   "${DRY_RUN_ARGS[@]}" \
   "${CUDA_GRAPH_ARGS[@]}" \
-  --max-records "${MAX_RECORDS:-20000}" \
-  --sample-rate "${SAMPLE_RATE:-1.0}" \
-  --benchmark-timeout "${BENCHMARK_TIMEOUT:-7200}" \
+  --docker-image "$DOCKER_IMAGE" \
+  --max-records "$MAX_RECORDS" \
+  --sample-rate "$SAMPLE_RATE" \
+  --benchmark-timeout "$BENCHMARK_TIMEOUT" \
   --framework "${FRAMEWORK:-sglang}" \
   -b "$TRACE_BENCH"