[Core] Add explicit layer parallel plans

Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
2026-07-13 19:12:30 +00:00
1866 changed files with 22505 additions and 166080 deletions
@@ -1,103 +0,0 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Audit vLLM compiled libraries for PyTorch stable ABI compliance."""
 import fnmatch
 import sys
 from pathlib import Path
 from torch_abi_audit import inspect_package
 from torch_abi_audit.report import ExtensionReport, PackageReport
 # Temporary allowlist of extensions not yet on the stable ABI.
 # Shrink and remove over time.
 ALLOWED_UNSTABLE_LIBRARIES: tuple[str, ...] = (
    "_flashkda_C.abi3.so",
    "vllm_flash_attn/_vllm_fa2_C.abi3.so",
    "vllm_flash_attn/_vllm_fa3_C.abi3.so",
    "third_party/deep_gemm/_C*.so",
 )
 def _relative_path(lib: ExtensionReport, package_root: Path) -> str:
    try:
        return lib.path.relative_to(package_root).as_posix()
    except ValueError:
        return lib.path.name
 def _is_torch_unstable(lib: ExtensionReport) -> bool:
    return lib.error is None and lib.torch.uses_torch and not lib.torch.stable
 def _matches_allowlist(rel_path: str, patterns: tuple[str, ...]) -> bool:
    return any(fnmatch.fnmatch(rel_path, pattern) for pattern in patterns)
 def _iter_libs(report: PackageReport) -> tuple[ExtensionReport, ...]:
    return (*report.extensions, *report.bundled_libs)
 def _collect_unstable(report: PackageReport) -> list[str]:
    return sorted(
        _relative_path(lib, report.root)
        for lib in _iter_libs(report)
        if _is_torch_unstable(lib)
    )
 def _find_stale_allowlist_entries(
    report: PackageReport, patterns: tuple[str, ...]
 ) -> list[str]:
    """Allowlist patterns that match a built library which is no longer unstable."""
    stale: list[str] = []
    for pattern in patterns:
        for lib in _iter_libs(report):
            if lib.error is not None:
                continue
            if not fnmatch.fnmatch(_relative_path(lib, report.root), pattern):
                continue
            if not _is_torch_unstable(lib):
                stale.append(pattern)
                break
    return stale
 def check_torch_abi(
    package: str = "vllm",
    patterns: tuple[str, ...] = ALLOWED_UNSTABLE_LIBRARIES,
 ) -> int:
    report = inspect_package(package)
    if report.error:
        print(f"error: failed to inspect {package!r}: {report.error}", file=sys.stderr)
        return 2
    unstable = _collect_unstable(report)
    unexpected = [
        rel_path for rel_path in unstable if not _matches_allowlist(rel_path, patterns)
    ]
    stale = _find_stale_allowlist_entries(report, patterns)
    if unexpected or stale:
        if unexpected:
            print(
                "Not allowed: torch-unstable libraries outside "
                f"ALLOWED_UNSTABLE_LIBRARIES: {', '.join(unexpected)}",
                file=sys.stderr,
            )
        if stale:
            print(
                "Not allowed: stale ALLOWED_UNSTABLE_LIBRARIES entries: "
                f"{', '.join(stale)}",
                file=sys.stderr,
            )
        return 1
    print("Torch stable ABI check passed.")
    return 0
 if __name__ == "__main__":
    print(">>> Auditing vLLM extension modules for PyTorch stable ABI compliance")
    sys.exit(check_torch_abi())
@@ -14,7 +14,6 @@ run_all_patterns:
  - "setup.py"
  - "csrc/"
  - "cmake/"
  - ".buildkite/check-torch-abi.py"
 run_all_exclude_patterns:
  - "docker/Dockerfile."
  - "csrc/cpu/"
@@ -18,8 +18,6 @@ steps:
      TERM: "xterm-256color"
    retry:
      automatic:
        - exit_status: 1  # Transient Docker/BuildKit failure
          limit: 1
        - exit_status: -1  # Agent was lost
          limit: 1
        - exit_status: -10  # Agent was lost
@@ -48,8 +46,6 @@ steps:
      VLLM_BRANCH: "$BUILDKITE_COMMIT"
    retry:
      automatic:
        - exit_status: 1  # Transient Docker/BuildKit failure
          limit: 1
        - exit_status: -1  # Agent was lost
          limit: 1
        - exit_status: -10  # Agent was lost
@@ -76,8 +72,6 @@ steps:
      VLLM_BRANCH: "$BUILDKITE_COMMIT"
    retry:
      automatic:
        - exit_status: 1  # Transient Docker/BuildKit failure
          limit: 1
        - exit_status: -1  # Agent was lost
          limit: 1
        - exit_status: -10  # Agent was lost
@@ -18,8 +18,6 @@ steps:
  - tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
  - tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
  - tests/kernels/mamba/test_cpu_short_conv.py
  - tests/kernels/mamba/test_causal_conv1d.py
  - tests/kernels/mamba/test_mamba_ssm.py
  commands:
    - |
      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
@@ -30,9 +28,7 @@ steps:
      pytest -x -v -s tests/kernels/test_onednn.py
      pytest -x -v -s tests/kernels/test_awq_int4_to_int8.py
      pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
-      pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
+      pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py"
      pytest -x -v -s tests/kernels/mamba/test_causal_conv1d.py
      pytest -x -v -s tests/kernels/mamba/test_mamba_ssm.py"
 # Note: SDE can't be downloaded from CI host because of AWS WAF
 # - label: CPU-Compatibility Tests
@@ -18,7 +18,7 @@ steps:
  - label: "XPU example Test"
    depends_on:
      - image-build-xpu
-    timeout_in_minutes: 50
+    timeout_in_minutes: 30
    optional: true
    device: intel_gpu
    agent_tags:
@@ -39,7 +39,7 @@ steps:
  - label: "XPU V1 test"
    depends_on:
      - image-build-xpu
-    timeout_in_minutes: 70
+    timeout_in_minutes: 30
    optional: true
    device: intel_gpu
    agent_tags:
@@ -60,7 +60,7 @@ steps:
  - label: "XPU server test"
    depends_on:
      - image-build-xpu
-    timeout_in_minutes: 45
+    timeout_in_minutes: 30
    optional: true
    device: intel_gpu
    agent_tags:
@@ -3,7 +3,7 @@ depends_on:
  - image-build-xpu
 steps:
 - label: XPU Sleep Mode
-  timeout_in_minutes: 45
+  timeout_in_minutes: 30
  device: intel_gpu
  agent_tags:
    label: production
@@ -1,26 +0,0 @@
 group: Benchmarks
 depends_on: 
  - image-build-xpu
 steps:
 - label: Benchmarks CLI Test
  key: benchmarks-cli-test
  timeout_in_minutes: 40
  device: intel_gpu
  agent_tags:
    label: production
    gpu: 1+
    mem: 16+
  no_plugin: true
  working_dir: "."
  env:
    REGISTRY: "public.ecr.aws/q9t5s3a7"
    REPO: "vllm-ci-test-repo"
    VLLM_TEST_DEVICE: "xpu"
  source_file_dependencies:
  - vllm/
  - tests/benchmarks/
  commands:
    - >-
      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
      'cd tests &&
      pytest -v -s benchmarks/'
@@ -2,44 +2,6 @@ group: Engine Intel
 depends_on:
  - image-build-xpu
 steps:
 - label: Engine
  key: engine
  timeout_in_minutes: 40
  device: intel_gpu
  agent_tags:
    label: production
    gpu: 1+
    mem: 16+
  no_plugin: true
  working_dir: "."
  env:
    REGISTRY: "public.ecr.aws/q9t5s3a7"
    REPO: "vllm-ci-test-repo"
    VLLM_TEST_DEVICE: "xpu"
  source_file_dependencies:
  - vllm/compilation/
  - vllm/config/
  - vllm/engine/
  - vllm/entrypoints/logger.py
  - vllm/envs.py
  - vllm/logger.py
  - vllm/logging_utils/
  - vllm/platforms/
  - vllm/sequence.py
  - vllm/triton_utils/
  - vllm/utils/
  - tests/engine
  - tests/test_sequence
  - tests/test_config
  - tests/test_logger
  - tests/test_vllm_port
  - tests/test_jit_monitor.py
  commands:
    - >-
      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
      'cd tests &&
      pytest -v -s engine/test_arg_utils.py test_sequence.py test_logger.py test_vllm_port.py test_jit_monitor.py'
 - label: Engine (1 GPU)
  timeout_in_minutes: 30
  device: intel_gpu
@@ -61,41 +23,3 @@ steps:
      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
      'cd tests &&
      pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py'
 - label: V1 e2e (2 GPUs)
  timeout_in_minutes: 30
  device: intel_gpu
  agent_tags:
    label: production
    gpu: 2+
    mem: 16+
  no_plugin: true
  working_dir: "."
  env:
    REGISTRY: "public.ecr.aws/q9t5s3a7"
    REPO: "vllm-ci-test-repo"
    VLLM_TEST_DEVICE: "xpu"
  source_file_dependencies:
    - vllm/compilation/
    - vllm/config/
    - vllm/distributed/
    - vllm/engine/
    - vllm/envs.py
    - vllm/forward_context.py
    - vllm/inputs/
    - vllm/logger.py
    - vllm/logging_utils/
    - vllm/model_executor/
    - vllm/multimodal/
    - vllm/platforms/
    - vllm/sampling_params.py
    - vllm/transformers_utils/
    - vllm/triton_utils/
    - vllm/utils/
    - vllm/v1/
    - tests/v1/e2e/spec_decode
  commands:
  - >-
      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
      'cd tests &&
      pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "tensor_parallelism"'
@@ -86,7 +86,7 @@ steps:
      pytest -v -s lora/test_punica_ops.py::test_add_lora_fused_moe_early_exit'
 - label: LoRA Punica FP8/XPU Ops
-  timeout_in_minutes: 60
+  timeout_in_minutes: 45
  device: intel_gpu
  agent_tags:
    label: production
@@ -3,7 +3,7 @@ depends_on:
  - image-build-xpu
 steps:
 - label: V1 Core + KV + Metrics
-  timeout_in_minutes: 45
+  timeout_in_minutes: 30
  device: intel_gpu
  agent_tags:
    label: production
@@ -33,7 +33,7 @@ steps:
      pytest -v -s v1/executor'
 - label: V1 Sample + Logits
-  timeout_in_minutes: 90
+  timeout_in_minutes: 30
  device: intel_gpu
  agent_tags:
    label: production
@@ -125,13 +125,13 @@ steps:
      pytest -v -s v1/kv_offload &&
      pytest -v -s v1/kv_connector/unit/test_offloading_connector.py'
- label: NixlConnector PD accuracy (4 GPUs)
+- label: NixlConnector PD accuracy (2 GPUs)
  timeout_in_minutes: 60
-  num_devices: 4
+  num_devices: 2
  device: intel_gpu
  agent_tags:
    label: production
-    gpu: 4+
+    gpu: 2+
    mem: 16+
  no_plugin: true
  working_dir: "."
@@ -148,14 +148,11 @@ steps:
    - >-
      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
      'cd tests &&
-      bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh &&
+      bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh'
      PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=1 bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh &&
      PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh &&
      PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh'
 - label: Regression
  key: regression
-  timeout_in_minutes: 50
+  timeout_in_minutes: 30
  device: intel_gpu
  agent_tags:
    label: production
@@ -189,7 +186,7 @@ steps:
 - label: Metrics, Tracing (2 GPUs)
  key: metrics-tracing-2-gpus
-  timeout_in_minutes: 45
+  timeout_in_minutes: 30
  num_devices: 2
  device: intel_gpu
  agent_tags:
@@ -225,7 +222,7 @@ steps:
 - label: Async Engine, Inputs, Utils, Worker
  key: async-engine-inputs-utils-worker
-  timeout_in_minutes: 55
+  timeout_in_minutes: 30
  device: intel_gpu
  agent_tags:
    label: production
@@ -262,25 +259,3 @@ steps:
    pytest -v -s detokenizer &&
    pytest -v -s -m "not cpu_test" ./multimodal &&
    pytest -v -s utils_ --ignore=utils_/test_mem_utils.py'
 - label: Fusion Unit Tests
  timeout_in_minutes: 30
  device: intel_gpu
  agent_tags:
    label: production
    gpu: 1+
    mem: 16+
  no_plugin: true
  working_dir: "."
  env:
    REGISTRY: "public.ecr.aws/q9t5s3a7"
    REPO: "vllm-ci-test-repo"
    VLLM_TEST_DEVICE: "xpu"
  source_file_dependencies:
    - vllm/compilation/
    - tests/compile/passes/test_qk_norm_rope_fusion.py
  commands:
    - >-
      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
      'cd tests &&
      pytest -v -s compile/passes/test_qk_norm_rope_fusion.py'
@@ -1,33 +0,0 @@
 group: Model Executor Intel
 depends_on:
  - image-build-xpu
 steps:
 - label: Model Executor (Intel)
  key: model-executor-intel
  timeout_in_minutes: 45
  device: intel_gpu
  agent_tags:
    label: production
    gpu: 1+
    mem: 24+
  no_plugin: true
  working_dir: "."
  env:
    REGISTRY: "public.ecr.aws/q9t5s3a7"
    REPO: "vllm-ci-test-repo"
    VLLM_TEST_DEVICE: "xpu"
  source_file_dependencies:
  - vllm/engine/arg_utils.py
  - vllm/config/model.py
  - vllm/model_executor
  - tests/model_executor
  commands:
    - >-
      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
      'apt-get update && apt-get install -y curl libsodium23 &&
      pip3 install tensorizer==2.10.1 &&
      pip3 install runai-model-streamer[s3,gcs,azure]\>=0.15.7 &&
      export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
      export PYTHONFAULTHANDLER=1 &&
      cd tests &&
      pytest -v -s model_executor -m "not slow_test" --ignore="model_executor/layers/test_rocm_unquantized_gemm.py" --deselect="tests/model_executor/model_loader/test_reload.py::test_kv_scale_reload"'
@@ -8,7 +8,7 @@ steps:
  agent_tags:
    label: production
    gpu: 2+
-    mem: 24+
+    mem: 16+
  no_plugin: true
  working_dir: "."
  env:
@@ -28,9 +28,7 @@ steps:
      'export VLLM_USE_V2_MODEL_RUNNER=1 &&
      cd tests &&
      pytest -v -s v1/engine/test_llm_engine.py -k "not test_engine_metrics" &&
      pytest -v -s v1/e2e/general/test_context_length.py &&
      ENFORCE_EAGER=1 pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram" &&
      pytest -v -s entrypoints/llm/test_struct_output_generate.py -k "xgrammar and not speculative_config6 and not speculative_config7 and not speculative_config8 and not speculative_config0" &&
      pytest -v -s v1/e2e/general/test_min_tokens.py'
 - label: Model Runner V2 Examples (Intel)
@@ -62,55 +60,3 @@ steps:
      python3 basic/offline_inference/generate.py --model facebook/opt-125m &&
      python3 generate/multimodal/vision_language_offline.py --seed 0 &&
      python3 features/automatic_prefix_caching/prefix_caching_offline.py'
 - label: Model Runner V2 Distributed (2 GPUs)
  timeout_in_minutes: 50
  device: intel_gpu
  agent_tags:
    label: production
    gpu: 2+
    mem: 16+
  no_plugin: true
  working_dir: "."
  env:
    REGISTRY: "public.ecr.aws/q9t5s3a7"
    REPO: "vllm-ci-test-repo"
    VLLM_TEST_DEVICE: "xpu"
  source_file_dependencies:
    - vllm/v1/worker/gpu/
    - vllm/v1/worker/gpu_worker.py
    - tests/basic_correctness/test_basic_correctness.py
    - tests/v1/distributed/test_async_llm_dp.py
    - tests/v1/distributed/test_eagle_dp.py
  commands:
    - >-
      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
      'export VLLM_USE_V2_MODEL_RUNNER=1 &&
      cd tests &&
      TARGET_TEST_SUITE=L4 pytest -v -s basic_correctness/test_basic_correctness.py -m "distributed\(num_gpus=2\)" -k "not ray and not True"'
 - label: Model Runner V2 Spec Decode
  timeout_in_minutes: 50
  device: intel_gpu
  agent_tags:
    label: production
    gpu: 1+
    mem: 24+
  no_plugin: true
  working_dir: "."
  env:
    REGISTRY: "public.ecr.aws/q9t5s3a7"
    REPO: "vllm-ci-test-repo"
    VLLM_TEST_DEVICE: "xpu"
  source_file_dependencies:
  - vllm/v1/worker/gpu/
  - vllm/v1/worker/gpu_worker.py
  - tests/v1/spec_decode/test_max_len.py
  - tests/v1/spec_decode/test_rejection_sampler_utils.py
  - tests/v1/e2e/spec_decode/test_spec_decode.py
  commands:
    - >-
      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
      'export VLLM_USE_V2_MODEL_RUNNER=1 &&
      cd tests &&
      pytest -v -s v1/spec_decode/test_synthetic_rejection_sampler_utils.py'
@@ -4,7 +4,7 @@ depends_on:
 steps:
 - label: Distributed Model Tests (2 GPUs)
  key: distributed-model-tests-2-gpus
-  timeout_in_minutes: 65
+  timeout_in_minutes: 50
  device: intel_gpu
  agent_tags:
    label: production
@@ -4,7 +4,7 @@ depends_on:
 steps:
 - label: "Multi-Modal Models (Standard) 1: qwen2"
  key: multi-modal-models-standard-1-qwen2
-  timeout_in_minutes: 70
+  timeout_in_minutes: 45
  device: intel_gpu
  agent_tags:
    label: production
@@ -29,7 +29,7 @@ steps:
 - label: "Multi-Modal Models (Standard) 2: qwen3 + gemma"
  key: multi-modal-models-standard-2-qwen3-gemma
-  timeout_in_minutes: 70
+  timeout_in_minutes: 45
  device: intel_gpu
  agent_tags:
    label: production
@@ -52,7 +52,7 @@ steps:
 - label: "Multi-Modal Models (Standard) 3: llava + qwen2_vl"
  key: multi-modal-models-standard-3-llava-qwen2-vl
-  timeout_in_minutes: 65
+  timeout_in_minutes: 45
  device: intel_gpu
  agent_tags:
    label: production
@@ -100,7 +100,7 @@ steps:
 - label: Multi-Modal Processor # 44min
  key: multi-modal-processor
-  timeout_in_minutes: 60
+  timeout_in_minutes: 45
  device: intel_gpu
  agent_tags:
    label: production
@@ -1,29 +0,0 @@
 group: Samplers Intel
 depends_on:
  - image-build-xpu
 steps:
 - label: Samplers Test (FlashInfer)
  key: samplers-test-flashinfer-intel
  timeout_in_minutes: 40
  device: intel_gpu
  agent_tags:
    label: production
    gpu: 1+
    mem: 24+
  no_plugin: true
  working_dir: "."
  env:
    REGISTRY: "public.ecr.aws/q9t5s3a7"
    REPO: "vllm-ci-test-repo"
    VLLM_TEST_DEVICE: "xpu"
  source_file_dependencies:
  - vllm/model_executor/layers
  - vllm/sampling_metadata.py
  - tests/samplers
  - tests/conftest.py
  - vllm/entrypoints/generate/beam_search
  commands:
    - >-
      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
      'cd tests &&
      VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers'
@@ -17,7 +17,7 @@ steps:
  - label: "XPU example Test"
    depends_on:
      - image-build-xpu
-    timeout_in_minutes: 50
+    timeout_in_minutes: 30
    device: intel_gpu
    agent_tags:
      label: production
@@ -76,7 +76,7 @@ steps:
  - label: "XPU V1 test"
    depends_on:
      - image-build-xpu
-    timeout_in_minutes: 70
+    timeout_in_minutes: 30
    device: intel_gpu
    agent_tags:
      label: production
@@ -99,13 +99,12 @@ steps:
        pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py --ignore=v1/worker/test_worker_memory_snapshot.py &&
        pytest -v -s v1/structured_output &&
        pytest -v -s v1/test_serial_utils.py &&
        pytest -v -s v1/e2e/general/test_correctness_sliding_window.py --deselect="tests/v1/e2e/general/test_correctness_sliding_window.py::test_sliding_window_retrieval[True-1-5-google/gemma-3-1b-it]" &&
        pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py --ignore=v1/spec_decode/test_speculators_correctness.py &&
        pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py --ignore=v1/kv_connector/unit/test_offloading_connector.py'
  - label: "XPU server test"
    depends_on:
      - image-build-xpu
-    timeout_in_minutes: 45
+    timeout_in_minutes: 30
    device: intel_gpu
    agent_tags:
      label: production
@@ -145,32 +144,7 @@ steps:
      - >-
        bash .buildkite/scripts/hardware_ci/run-intel-test.sh
        'cd tests &&
-        pytest -v -s quantization/test_auto_round.py &&
+        pytest -v -s quantization/test_auto_round.py'
        pytest -v -s quantization/test_online.py'
  - label: "XPU GPQA Eval (GPT-OSS)"
    depends_on:
      - image-build-xpu
    timeout_in_minutes: 60
    device: intel_gpu
    agent_tags:
      label: production
      gpu: 1+
      mem: 24+
    no_plugin: true
    env:
      REGISTRY: "public.ecr.aws/q9t5s3a7"
      REPO: "vllm-ci-test-repo"
      VLLM_TEST_DEVICE: "xpu"
    source_file_dependencies:
      - vllm/
      - tests/evals/gpt_oss/
      - .buildkite/intel_jobs/test-intel.yaml
    commands:
      - >-
        bash .buildkite/scripts/hardware_ci/run-intel-test.sh
        'pip install "gpt-oss[eval]==0.0.5" &&
        cd tests &&
        pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-xpu.txt'
  - label: "XPU compressed tensors FP8 test"
    depends_on:
      - image-build-xpu
@@ -1,7 +1,6 @@
 # For hf script, without -t option (tensor parallel size).
 # bash .buildkite/lm-eval-harness/run-lm-eval-mmlupro-vllm-baseline.sh -m meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8 -l 250 -t 8 -f 5
 model_name: "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8"
 rocm_safetensors_load_strategy: lazy
 required_gpu_arch:
  - gfx942
  - gfx950
@@ -72,11 +72,6 @@ def launch_lm_eval(eval_config, tp_size):
    if moe_backend is not None:
        model_args += f"moe_backend={moe_backend},"
    if current_platform.is_rocm():
        rocm_load_strategy = eval_config.get("rocm_safetensors_load_strategy")
        if rocm_load_strategy is not None:
            model_args += f"safetensors_load_strategy={rocm_load_strategy},"
    env_vars = eval_config.get("env_vars", None)
    with scoped_env_vars(env_vars):
        results = lm_eval.simple_evaluate(
@@ -28,6 +28,11 @@
        "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json"
      }
    },
    {
        "dataset_name": "sharegpt",
        "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json"
      }
    },
    {
      "test_name": "serving_llama8B_tp1_random_128_128",
      "server_parameters": {
@@ -31,46 +31,8 @@ steps:
      - text: "What is the release version?"
        key: release-version
-  - group: "Build CUDA 13.0 Python wheels"
+  - group: "Build Python wheels"
    key: "build-wheels"
    steps:
      - label: "Build wheel - aarch64 - CUDA 13.0"
        depends_on: ~
        id: build-wheel-arm64-cuda-13-0
        agents:
          queue: arm64_cpu_queue_release
        commands:
          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg BUILD_OS=manylinux --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.0 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
          - "mkdir artifacts"
          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
          - "bash .buildkite/scripts/upload-nightly-wheels.sh"
          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)" release-wheels'
        env:
          DOCKER_BUILDKIT: "1"
      - label: "Build wheel - x86_64 - CUDA 13.0"
        depends_on: ~
        id: build-wheel-x86-cuda-13-0
        agents:
          queue: cpu_queue_release
        commands:
          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg BUILD_OS=manylinux --build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda13.0 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
          - "mkdir artifacts"
          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
          - "bash .buildkite/scripts/upload-nightly-wheels.sh"
          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)" release-wheels'
        env:
          DOCKER_BUILDKIT: "1"
  - block: "Unblock to build additional Python wheels"
    depends_on: ~
    key: block-build-additional-wheels
    if: build.env("NIGHTLY") != "1"
  - group: "Build additional Python wheels"
    key: "build-additional-wheels"
    depends_on: block-build-additional-wheels
    allow_dependency_failure: true
    steps:
      - label: "Build wheel - aarch64 - CUDA 12.9"
        depends_on: ~
@@ -82,7 +44,21 @@ steps:
          - "mkdir artifacts"
          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
          - "bash .buildkite/scripts/upload-nightly-wheels.sh"
-          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)" release-wheels'
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"'
        env:
          DOCKER_BUILDKIT: "1"
      - label: "Build wheel - aarch64 - CUDA 13.0"
        depends_on: ~
        id: build-wheel-arm64-cuda-13-0
        agents:
          queue: arm64_cpu_queue_release
        commands:
          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg BUILD_OS=manylinux --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.0 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
          - "mkdir artifacts"
          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
          - "bash .buildkite/scripts/upload-nightly-wheels.sh"
          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"'
        env:
          DOCKER_BUILDKIT: "1"
@@ -96,53 +72,10 @@ steps:
          - "mkdir artifacts"
          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
          - "bash .buildkite/scripts/upload-nightly-wheels.sh"
-          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)" release-wheels'
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"'
        env:
          DOCKER_BUILDKIT: "1"
      - label: "Build wheel - macOS arm64 - CPU"
        depends_on: ~
        id: build-wheel-macos-arm64-cpu
        agents:
          queue: macmini
        timeout_in_minutes: 90
        commands:
          - "bash .buildkite/scripts/build-macos-wheel.sh"
          - "mkdir -p artifacts/transfer"
          - "basename artifacts/dist/*.whl > artifacts/transfer/wheel.name"
          - "shasum -a 256 artifacts/dist/*.whl > artifacts/transfer/wheel.sha256"
          - "split -b 1m artifacts/dist/*.whl artifacts/transfer/wheel.part."
        env:
          BUILDKITE_ARTIFACT_UPLOAD_CONCURRENCY: "1"
          BUILDKITE_NO_HTTP2: "true"
        artifact_paths:
          - "artifacts/transfer/*"
      - label: "Upload wheel - macOS arm64 - CPU"
        depends_on: build-wheel-macos-arm64-cpu
        id: upload-wheel-macos-arm64-cpu
        agents:
          queue: small_cpu_queue_release
        timeout_in_minutes: 15
        commands:
          - 'buildkite-agent artifact download "artifacts/transfer/*" .'
          - "mkdir -p artifacts/dist artifacts/reassembled"
          - 'wheel_name=$$(cat artifacts/transfer/wheel.name)'
          - '[[ "$$wheel_name" =~ ^[A-Za-z0-9][A-Za-z0-9._+-]*\.whl$$ && "$$wheel_name" != *..* ]]'
          - 'expected_sha=$$(awk ''{print $$1}'' artifacts/transfer/wheel.sha256)'
          - '[[ "$$expected_sha" =~ ^[0-9a-f]{64}$$ ]]'
          - 'cat artifacts/transfer/wheel.part.* > artifacts/reassembled/wheel'
          - 'actual_sha=$$(shasum -a 256 artifacts/reassembled/wheel | awk ''{print $$1}'')'
          - '[[ "$$actual_sha" == "$$expected_sha" ]]'
          - 'mv artifacts/reassembled/wheel "artifacts/dist/$$wheel_name"'
          - "aws sts get-caller-identity"
          - "VLLM_WHEEL_PLATFORM=macos bash .buildkite/scripts/upload-nightly-wheels.sh"
          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)" release-wheels'
        plugins:
          - aws-assume-role-with-web-identity#v1.6.0:
              role-arn: arn:aws:iam::936637512419:role/vllm-release-macos-wheel-uploader
              region: us-west-2
      - label: "Build wheel - x86_64 - CUDA 12.9"
        depends_on: ~
        id: build-wheel-x86-cuda-12-9
@@ -153,7 +86,21 @@ steps:
          - "mkdir artifacts"
          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
          - "bash .buildkite/scripts/upload-nightly-wheels.sh"
-          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)" release-wheels'
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"'
        env:
          DOCKER_BUILDKIT: "1"
      - label: "Build wheel - x86_64 - CUDA 13.0"
        depends_on: ~
        id: build-wheel-x86-cuda-13-0
        agents:
          queue: cpu_queue_release
        commands:
          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg BUILD_OS=manylinux --build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda13.0 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
          - "mkdir artifacts"
          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
          - "bash .buildkite/scripts/upload-nightly-wheels.sh"
          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"'
        env:
          DOCKER_BUILDKIT: "1"
@@ -167,31 +114,17 @@ steps:
          - "mkdir artifacts"
          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
          - "bash .buildkite/scripts/upload-nightly-wheels.sh"
-          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)" release-wheels'
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"'
        env:
          DOCKER_BUILDKIT: "1"
  - label: "Generate and upload wheel indices"
    key: generate-wheel-indices
    depends_on: "build-wheels"
    allow_dependency_failure: true
    if: build.env("NIGHTLY") != "1"
    agents:
      queue: cpu_queue_release
    commands:
-      - "UPDATE_VERSION_INDEX=0 bash .buildkite/scripts/generate-and-upload-nightly-index.sh"
+      - "bash .buildkite/scripts/generate-and-upload-nightly-index.sh"
  - label: "Regenerate indices with additional wheels"
    key: generate-additional-wheel-indices
    depends_on:
      - build-wheels
      - build-additional-wheels
      - generate-wheel-indices
    allow_dependency_failure: true
    agents:
      queue: cpu_queue_release
    commands:
      - 'UPDATE_NIGHTLY_INDEX="$${NIGHTLY:-0}" bash .buildkite/scripts/generate-and-upload-nightly-index.sh'
  - block: "Unblock to build release Docker images"
    depends_on: ~
@@ -229,7 +162,7 @@ steps:
          # re-tag to default image tag and push, just in case arm64 build fails
          - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
-          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)" release-images'
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"'
      - label: "Build release image - aarch64 - CUDA 13.0"
        depends_on: ~
@@ -254,7 +187,7 @@ steps:
              --progress plain \
              -f docker/Dockerfile .
          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"
-          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)" release-images'
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"'
      - label: "Build release image - x86_64 - CUDA 12.9"
        depends_on: ~
@@ -281,7 +214,7 @@ steps:
          # re-tag to default image tag and push, just in case arm64 build fails
          - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"
          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"
-          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129" release-images'
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129"'
      - label: "Build release image - aarch64 - CUDA 12.9"
        depends_on: ~
@@ -305,7 +238,7 @@ steps:
              --progress plain \
              -f docker/Dockerfile .
          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129"
-          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129" release-images'
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129"'
      - label: "Build release image - x86_64 - CUDA 13.0 - Ubuntu 24.04"
        depends_on: ~
@@ -334,7 +267,7 @@ steps:
          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"
          - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"
          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"
-          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404" release-images'
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"'
      - label: "Build release image - aarch64 - CUDA 13.0 - Ubuntu 24.04"
        depends_on: ~
@@ -361,7 +294,7 @@ steps:
              --progress plain \
              -f docker/Dockerfile .
          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"
-          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404" release-images'
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"'
      - label: "Build release image - x86_64 - CUDA 12.9 - Ubuntu 24.04"
        depends_on: ~
@@ -389,7 +322,7 @@ steps:
          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"
          - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"
          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"
-          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404" release-images'
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"'
      - label: "Build release image - aarch64 - CUDA 12.9 - Ubuntu 24.04"
        depends_on: ~
@@ -415,7 +348,7 @@ steps:
              --progress plain \
              -f docker/Dockerfile .
          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"
-          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404" release-images'
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"'
      - label: ":docker: Build release image - x86_64 - XPU"
        depends_on: ~
@@ -432,7 +365,7 @@ steps:
              --progress plain \
              -f docker/Dockerfile.xpu .
          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-xpu"
-          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-xpu" release-images'
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-xpu"'
        env:
          DOCKER_BUILDKIT: "1"
@@ -452,7 +385,7 @@ steps:
          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_CPU_X86=true --tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version) --tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:latest --progress plain --target vllm-openai -f docker/Dockerfile.cpu ."
          - "docker push public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:latest"
          - "docker push public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version)"
-          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version)" release-images'
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version)"'
        env:
          DOCKER_BUILDKIT: "1"
@@ -472,7 +405,7 @@ steps:
          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --tag public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:$(buildkite-agent meta-data get release-version) --tag public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:latest --progress plain --target vllm-openai -f docker/Dockerfile.cpu ."
          - "docker push public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:latest"
          - "docker push public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:$(buildkite-agent meta-data get release-version)"
-          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:$(buildkite-agent meta-data get release-version)" release-images'
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:$(buildkite-agent meta-data get release-version)"'
        env:
          DOCKER_BUILDKIT: "1"
@@ -490,7 +423,7 @@ steps:
          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
          - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64 --amend"
          - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
-          - 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: CUDA 13.0" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT" release-manifests'
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: CUDA 13.0" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"'
      - label: "Create multi-arch manifest - CUDA 12.9"
        depends_on:
@@ -503,7 +436,7 @@ steps:
          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
          - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu129 --amend"
          - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"
-          - 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: CUDA 12.9" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129" release-manifests'
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: CUDA 12.9" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"'
      - label: "Create multi-arch manifest - CUDA 13.0 - Ubuntu 24.04"
        depends_on:
@@ -516,7 +449,7 @@ steps:
          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
          - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-ubuntu2404 --amend"
          - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"
-          - 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: CUDA 13.0 Ubuntu 24.04" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404" release-manifests'
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: CUDA 13.0 Ubuntu 24.04" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"'
      - label: "Create multi-arch manifest - CUDA 12.9 - Ubuntu 24.04"
        depends_on:
@@ -529,7 +462,7 @@ steps:
          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
          - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu129-ubuntu2404 --amend"
          - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"
-          - 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: CUDA 12.9 Ubuntu 24.04" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404" release-manifests'
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: CUDA 12.9 Ubuntu 24.04" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"'
      - label: "Create manifest - XPU"
        depends_on:
@@ -539,7 +472,7 @@ steps:
          queue: small_cpu_queue_release
        commands:
          - "bash .buildkite/scripts/xpu/create-xpu-ecr-manifest.sh"
-          - 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: XPU" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-xpu" release-manifests'
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: XPU" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-xpu"'
      - label: "Publish nightly multi-arch image to DockerHub"
        depends_on:
@@ -590,370 +523,366 @@ steps:
  #
  # =============================================================================
-  - group: "Build ROCm Wheel / Image "
+  # ROCm Job 1: Build ROCm Base Wheels (with S3 caching)
-    key: "build-rocm-wheel-image"
+  - label: ":rocm: Build ROCm Base Image & Wheels"
    id: build-rocm-base-wheels
    depends_on: ~
-    steps:
+    agents:
-      # ROCm Job 1: Build ROCm Base Wheels (with S3 caching)
+      queue: cpu_queue_release
-      - label: ":rocm: Build ROCm Base Image & Wheels"
+    commands:
-        id: build-rocm-base-wheels
+      - |
-        depends_on: ~
+        set -euo pipefail
        agents:
          queue: cpu_queue_release
        commands:
          - |
            set -euo pipefail
-            # Generate cache key
+        # Generate cache key
-            CACHE_KEY=$$(.buildkite/scripts/cache-rocm-base-wheels.sh key)
+        CACHE_KEY=$$(.buildkite/scripts/cache-rocm-base-wheels.sh key)
-            ECR_CACHE_TAG="public.ecr.aws/q9t5s3a7/vllm-release-repo:$${CACHE_KEY}-rocm-base"
+        ECR_CACHE_TAG="public.ecr.aws/q9t5s3a7/vllm-release-repo:$${CACHE_KEY}-rocm-base"
-            echo "========================================"
+        echo "========================================"
-            echo "ROCm Base Build Configuration"
+        echo "ROCm Base Build Configuration"
-            echo "========================================"
+        echo "========================================"
-            echo "  CACHE_KEY: $${CACHE_KEY}"
+        echo "  CACHE_KEY: $${CACHE_KEY}"
-            echo "  ECR_CACHE_TAG: $${ECR_CACHE_TAG}"
+        echo "  ECR_CACHE_TAG: $${ECR_CACHE_TAG}"
-            echo "========================================"
+        echo "========================================"
-            # Login to ECR
+        # Login to ECR
-            aws ecr-public get-login-password --region us-east-1 | \
+        aws ecr-public get-login-password --region us-east-1 | \
-              docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
+          docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
-            IMAGE_EXISTS=false
+        IMAGE_EXISTS=false
-            WHEELS_EXIST=false
+        WHEELS_EXIST=false
-            # Check ECR for Docker image
+        # Check ECR for Docker image
-            if docker manifest inspect "$${ECR_CACHE_TAG}" > /dev/null 2>&1; then
+        if docker manifest inspect "$${ECR_CACHE_TAG}" > /dev/null 2>&1; then
-              IMAGE_EXISTS=true
+          IMAGE_EXISTS=true
-              echo "ECR image cache HIT"
+          echo "ECR image cache HIT"
-            fi
+        fi
-            # Check S3 for wheels
+        # Check S3 for wheels
-            WHEEL_CACHE_STATUS=$(.buildkite/scripts/cache-rocm-base-wheels.sh check)
+        WHEEL_CACHE_STATUS=$(.buildkite/scripts/cache-rocm-base-wheels.sh check)
-            if [ "$${WHEEL_CACHE_STATUS}" = "hit" ]; then
+        if [ "$${WHEEL_CACHE_STATUS}" = "hit" ]; then
-              WHEELS_EXIST=true
+          WHEELS_EXIST=true
-              echo "S3 wheels cache HIT"
+          echo "S3 wheels cache HIT"
-            fi
+        fi
-            # Scenario 1: Both cached (best case)
+        # Scenario 1: Both cached (best case)
-            if [ "$${IMAGE_EXISTS}" = "true" ] && [ "$${WHEELS_EXIST}" = "true" ]; then
+        if [ "$${IMAGE_EXISTS}" = "true" ] && [ "$${WHEELS_EXIST}" = "true" ]; then
-              echo ""
+          echo ""
-              echo "FULL CACHE HIT - Reusing both image and wheels"
+          echo "FULL CACHE HIT - Reusing both image and wheels"
-              echo ""
+          echo ""
-              # Download wheels
+          # Download wheels
-              .buildkite/scripts/cache-rocm-base-wheels.sh download
+          .buildkite/scripts/cache-rocm-base-wheels.sh download
-              # Save ECR tag for downstream jobs
+          # Save ECR tag for downstream jobs
-              buildkite-agent meta-data set "rocm-base-image-tag" "$${ECR_CACHE_TAG}"
+          buildkite-agent meta-data set "rocm-base-image-tag" "$${ECR_CACHE_TAG}"
-            # Scenario 2: Full rebuild needed
+        # Scenario 2: Full rebuild needed
-            else
+        else
-              echo ""
+          echo ""
-              echo " CACHE MISS - Building from scratch..."
+          echo " CACHE MISS - Building from scratch..."
-              echo ""
+          echo ""
-              # Build full base image and push to ECR
+          # Build full base image and push to ECR
-              DOCKER_BUILDKIT=1 docker buildx build \
+          DOCKER_BUILDKIT=1 docker buildx build \
-                --file docker/Dockerfile.rocm_base \
+            --file docker/Dockerfile.rocm_base \
-                --tag "$${ECR_CACHE_TAG}" \
+            --tag "$${ECR_CACHE_TAG}" \
-                --build-arg USE_SCCACHE=1 \
+            --build-arg USE_SCCACHE=1 \
-                --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
+            --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
-                --build-arg SCCACHE_REGION_NAME=us-west-2 \
+            --build-arg SCCACHE_REGION_NAME=us-west-2 \
-                --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
+            --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
-                --push \
+            --push \
-                .
+            .
-              # Build wheel extraction stage
+          # Build wheel extraction stage
-              DOCKER_BUILDKIT=1 docker buildx build \
+          DOCKER_BUILDKIT=1 docker buildx build \
-                --file docker/Dockerfile.rocm_base \
+            --file docker/Dockerfile.rocm_base \
-                --tag rocm-base-debs:$${BUILDKITE_BUILD_NUMBER} \
+            --tag rocm-base-debs:$${BUILDKITE_BUILD_NUMBER} \
-                --target debs_wheel_release \
+            --target debs_wheel_release \
-                --build-arg USE_SCCACHE=1 \
+            --build-arg USE_SCCACHE=1 \
-                --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
+            --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
-                --build-arg SCCACHE_REGION_NAME=us-west-2 \
+            --build-arg SCCACHE_REGION_NAME=us-west-2 \
-                --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
+            --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
-                --load \
+            --load \
-                .
+            .
-              # Extract and upload wheels
+          # Extract and upload wheels
-              mkdir -p artifacts/rocm-base-wheels
+          mkdir -p artifacts/rocm-base-wheels
-              cid=$(docker create rocm-base-debs:$${BUILDKITE_BUILD_NUMBER})
+          cid=$(docker create rocm-base-debs:$${BUILDKITE_BUILD_NUMBER})
-              docker cp $${cid}:/app/debs/. artifacts/rocm-base-wheels/
+          docker cp $${cid}:/app/debs/. artifacts/rocm-base-wheels/
-              docker rm $${cid}
+          docker rm $${cid}
-              .buildkite/scripts/cache-rocm-base-wheels.sh upload
+          .buildkite/scripts/cache-rocm-base-wheels.sh upload
-              # Cache base docker image to ECR
+          # Cache base docker image to ECR
-              docker push "$${ECR_CACHE_TAG}"
+          docker push "$${ECR_CACHE_TAG}"
-              buildkite-agent meta-data set "rocm-base-image-tag" "$${ECR_CACHE_TAG}"
+          buildkite-agent meta-data set "rocm-base-image-tag" "$${ECR_CACHE_TAG}"
-              echo ""
+          echo ""
-              echo " Build complete - Image and wheels cached"
+          echo " Build complete - Image and wheels cached"
-            fi
+        fi
-        artifact_paths:
+    artifact_paths:
-          - "artifacts/rocm-base-wheels/*.whl"
+      - "artifacts/rocm-base-wheels/*.whl"
-        env:
+    env:
-          DOCKER_BUILDKIT: "1"
+      DOCKER_BUILDKIT: "1"
-          S3_BUCKET: "vllm-wheels"
+      S3_BUCKET: "vllm-wheels"
-      # ROCm Job 2: Build vLLM ROCm Wheel
+  # ROCm Job 2: Build vLLM ROCm Wheel
-      - label: ":python: Build vLLM ROCm Wheel - x86_64"
+  - label: ":python: Build vLLM ROCm Wheel - x86_64"
-        id: build-rocm-vllm-wheel
+    id: build-rocm-vllm-wheel
-        depends_on:
+    depends_on:
-          - step: build-rocm-base-wheels
+      - step: build-rocm-base-wheels
-            allow_failure: false
+        allow_failure: false
-        agents:
+    agents:
-          queue: cpu_queue_release
+      queue: cpu_queue_release
-        timeout_in_minutes: 180
+    timeout_in_minutes: 180
-        commands:
+    commands:
-          # Download artifacts and prepare Docker image
+      # Download artifacts and prepare Docker image
-          - |
+      - |
-            set -euo pipefail
+        set -euo pipefail
-            # Ensure git tags are up-to-date (Buildkite's default fetch doesn't update tags)
+        # Ensure git tags are up-to-date (Buildkite's default fetch doesn't update tags)
-            # This fixes version detection when tags are moved/force-pushed
+        # This fixes version detection when tags are moved/force-pushed
-            echo "Fetching latest tags from origin..."
+        echo "Fetching latest tags from origin..."
-            git fetch --tags --force origin
+        git fetch --tags --force origin
-            # Log tag information for debugging version detection
+        # Log tag information for debugging version detection
-            echo "========================================"
+        echo "========================================"
-            echo "Git Tag Verification"
+        echo "Git Tag Verification"
-            echo "========================================"
+        echo "========================================"
-            echo "Current HEAD: $(git rev-parse HEAD)"
+        echo "Current HEAD: $(git rev-parse HEAD)"
-            echo "git describe --tags: $(git describe --tags 2>/dev/null || echo 'No tags found')"
+        echo "git describe --tags: $(git describe --tags 2>/dev/null || echo 'No tags found')"
-            echo ""
+        echo ""
-            echo "Recent tags (pointing to commits near HEAD):"
+        echo "Recent tags (pointing to commits near HEAD):"
-            git tag -l --sort=-creatordate | head -5
+        git tag -l --sort=-creatordate | head -5
-            echo "setuptools_scm version detection:"
+        echo "setuptools_scm version detection:"
-            pip install -q setuptools_scm 2>/dev/null || true
+        pip install -q setuptools_scm 2>/dev/null || true
-            python3 -c "import setuptools_scm; print('  Detected version:', setuptools_scm.get_version())" 2>/dev/null || echo "  (setuptools_scm not available in this environment)"
+        python3 -c "import setuptools_scm; print('  Detected version:', setuptools_scm.get_version())" 2>/dev/null || echo "  (setuptools_scm not available in this environment)"
-            echo "========================================"
+        echo "========================================"
-            # Download wheel artifacts from current build
+        # Download wheel artifacts from current build
-            echo "Downloading wheel artifacts from current build"
+        echo "Downloading wheel artifacts from current build"
-            buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" .
+        buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" .
-            # Get ECR image tag from metadata (set by build-rocm-base-wheels)
+        # Get ECR image tag from metadata (set by build-rocm-base-wheels)
-            ECR_IMAGE_TAG="$$(buildkite-agent meta-data get rocm-base-image-tag 2>/dev/null || echo '')"
+        ECR_IMAGE_TAG="$$(buildkite-agent meta-data get rocm-base-image-tag 2>/dev/null || echo '')"
-            if [ -z "$${ECR_IMAGE_TAG}" ]; then
+        if [ -z "$${ECR_IMAGE_TAG}" ]; then
-              echo "ERROR: rocm-base-image-tag metadata not found"
+          echo "ERROR: rocm-base-image-tag metadata not found"
-              echo "This should have been set by the build-rocm-base-wheels job"
+          echo "This should have been set by the build-rocm-base-wheels job"
-              exit 1
+          exit 1
-            fi
+        fi
-            echo "Pulling base Docker image from ECR: $${ECR_IMAGE_TAG}"
+        echo "Pulling base Docker image from ECR: $${ECR_IMAGE_TAG}"
-            # Login to ECR
+        # Login to ECR
-            aws ecr-public get-login-password --region us-east-1 | \
+        aws ecr-public get-login-password --region us-east-1 | \
-              docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
+          docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
-            # Pull base Docker image from ECR
+        # Pull base Docker image from ECR
-            docker pull "$${ECR_IMAGE_TAG}"
+        docker pull "$${ECR_IMAGE_TAG}"
-            echo "Loaded base image: $${ECR_IMAGE_TAG}"
+        echo "Loaded base image: $${ECR_IMAGE_TAG}"
-            # Prepare base wheels for Docker build context
+        # Prepare base wheels for Docker build context
-            mkdir -p docker/context/base-wheels
+        mkdir -p docker/context/base-wheels
-            touch docker/context/base-wheels/.keep
+        touch docker/context/base-wheels/.keep
-            cp artifacts/rocm-base-wheels/*.whl docker/context/base-wheels/
+        cp artifacts/rocm-base-wheels/*.whl docker/context/base-wheels/
-            echo "Base wheels for vLLM build:"
+        echo "Base wheels for vLLM build:"
-            ls -lh docker/context/base-wheels/
+        ls -lh docker/context/base-wheels/
-            echo "========================================"
+        echo "========================================"
-            echo "Building vLLM wheel with:"
+        echo "Building vLLM wheel with:"
-            echo "  BUILDKITE_COMMIT: $${BUILDKITE_COMMIT}"
+        echo "  BUILDKITE_COMMIT: $${BUILDKITE_COMMIT}"
-            echo "  BUILDKITE_BRANCH: $${BUILDKITE_BRANCH}"
+        echo "  BUILDKITE_BRANCH: $${BUILDKITE_BRANCH}"
-            echo "  BASE_IMAGE: $${ECR_IMAGE_TAG}"
+        echo "  BASE_IMAGE: $${ECR_IMAGE_TAG}"
-            echo "========================================"
+        echo "========================================"
-            # Build vLLM wheel using local checkout (REMOTE_VLLM=0)
+        # Build vLLM wheel using local checkout (REMOTE_VLLM=0)
-            DOCKER_BUILDKIT=1 docker build \
+        DOCKER_BUILDKIT=1 docker build \
-              --file docker/Dockerfile.rocm \
+          --file docker/Dockerfile.rocm \
-              --target export_vllm_wheel_release \
+          --target export_vllm_wheel_release \
-              --output type=local,dest=rocm-dist \
+          --output type=local,dest=rocm-dist \
-              --build-arg BASE_IMAGE="$${ECR_IMAGE_TAG}" \
+          --build-arg BASE_IMAGE="$${ECR_IMAGE_TAG}" \
-              --build-arg REMOTE_VLLM=0 \
+          --build-arg REMOTE_VLLM=0 \
-              --build-arg GIT_REPO_CHECK=1 \
+          --build-arg GIT_REPO_CHECK=1 \
-              --build-arg USE_SCCACHE=1 \
+          --build-arg USE_SCCACHE=1 \
-              --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
+          --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
-              --build-arg SCCACHE_REGION_NAME=us-west-2 \
+          --build-arg SCCACHE_REGION_NAME=us-west-2 \
-              --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
+          --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
-              .
+          .
-            echo "Built vLLM wheel:"
+        echo "Built vLLM wheel:"
-            ls -lh rocm-dist/*.whl
+        ls -lh rocm-dist/*.whl
-            # Copy wheel to artifacts directory
+        # Copy wheel to artifacts directory
-            mkdir -p artifacts/rocm-vllm-wheel
+        mkdir -p artifacts/rocm-vllm-wheel
-            cp rocm-dist/*.whl artifacts/rocm-vllm-wheel/
+        cp rocm-dist/*.whl artifacts/rocm-vllm-wheel/
-            echo "Final vLLM wheel:"
+        echo "Final vLLM wheel:"
-            ls -lh artifacts/rocm-vllm-wheel/
+        ls -lh artifacts/rocm-vllm-wheel/
-        artifact_paths:
+    artifact_paths:
-          - "artifacts/rocm-vllm-wheel/*.whl"
+      - "artifacts/rocm-vllm-wheel/*.whl"
-        env:
+    env:
-          DOCKER_BUILDKIT: "1"
+      DOCKER_BUILDKIT: "1"
-          S3_BUCKET: "vllm-wheels"
+      S3_BUCKET: "vllm-wheels"
-      # ROCm Job 3: Upload Wheels to S3
+  # ROCm Job 3: Upload Wheels to S3
-      - label: ":s3: Upload ROCm Wheels to S3"
+  - label: ":s3: Upload ROCm Wheels to S3"
-        id: upload-rocm-wheels
+    id: upload-rocm-wheels
-        depends_on:
+    depends_on:
-          - step: build-rocm-vllm-wheel
+      - step: build-rocm-vllm-wheel
-            allow_failure: false
+        allow_failure: false
-        agents:
+    agents:
-          queue: cpu_queue_release
+      queue: cpu_queue_release
-        timeout_in_minutes: 60
+    timeout_in_minutes: 60
-        commands:
+    commands:
-          # Download all wheel artifacts and run upload
+      # Download all wheel artifacts and run upload
-          - |
+      - |
-            set -euo pipefail
+        set -euo pipefail
-            # Download artifacts from current build
+        # Download artifacts from current build
-            echo "Downloading artifacts from current build"
+        echo "Downloading artifacts from current build"
-            buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" .
+        buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" .
-            buildkite-agent artifact download "artifacts/rocm-vllm-wheel/*.whl" .
+        buildkite-agent artifact download "artifacts/rocm-vllm-wheel/*.whl" .
-            # # Run upload script
+        # Run upload script
-            bash .buildkite/scripts/upload-rocm-wheels.sh
+        bash .buildkite/scripts/upload-rocm-wheels.sh
-        env:
+    env:
-          DOCKER_BUILDKIT: "1"
+      DOCKER_BUILDKIT: "1"
-          S3_BUCKET: "vllm-wheels"
+      S3_BUCKET: "vllm-wheels"
-      # ROCm Job 4: Annotate ROCm Wheel Release
+  # ROCm Job 4: Annotate ROCm Wheel Release
-      - label: ":memo: Annotate ROCm wheel release"
+  - label: ":memo: Annotate ROCm wheel release"
-        id: annotate-rocm-release
+    id: annotate-rocm-release
-        depends_on:
+    depends_on:
-          - upload-rocm-wheels
+      - upload-rocm-wheels
-        agents:
+    agents:
-          queue: cpu_queue_release
+      queue: cpu_queue_release
-        commands:
+    commands:
-          - "bash .buildkite/scripts/annotate-rocm-release.sh"
+      - "bash .buildkite/scripts/annotate-rocm-release.sh"
-        env:
+    env:
-          S3_BUCKET: "vllm-wheels"
+      S3_BUCKET: "vllm-wheels"
-      # ROCm Job 5: Generate Root Index for ROCm Wheels (for release only)
+  # ROCm Job 5: Generate Root Index for ROCm Wheels (for release only)
-      # This is the job to create https://wheels.vllm.ai/rocm/ index allowing
+  # This is the job to create https://wheels.vllm.ai/rocm/ index allowing
-      # users to install with `uv pip install vllm --extra-index-url https://wheels.vllm.ai/rocm/`
+  # users to install with `uv pip install vllm --extra-index-url https://wheels.vllm.ai/rocm/`
-      - block: "Generate Root Index for ROCm Wheels for Release"
+  - block: "Generate Root Index for ROCm Wheels for Release"
-        key: block-generate-root-index-rocm-wheels
+    key: block-generate-root-index-rocm-wheels
-        depends_on: upload-rocm-wheels
+    depends_on: upload-rocm-wheels
-      - label: ":package: Generate Root Index for ROCm Wheels for Release"
+  - label: ":package: Generate Root Index for ROCm Wheels for Release"
-        depends_on: block-generate-root-index-rocm-wheels
+    depends_on: block-generate-root-index-rocm-wheels
-        id: generate-root-index-rocm-wheels
+    id: generate-root-index-rocm-wheels
-        agents:
+    agents:
-          queue: cpu_queue_release
+      queue: cpu_queue_release
-        commands:
+    commands:
-          - "bash tools/vllm-rocm/generate-rocm-wheels-root-index.sh"
+      - "bash tools/vllm-rocm/generate-rocm-wheels-root-index.sh"
-        env:
+    env:
-          S3_BUCKET: "vllm-wheels"
+      S3_BUCKET: "vllm-wheels"
-          VARIANT: "rocm723"
+      VARIANT: "rocm723"
-      # ROCm Job 6: Build ROCm Release Docker Image
+  # ROCm Job 6: Build ROCm Release Docker Image
-      - label: ":docker: Build release image - x86_64 - ROCm"
+  - label: ":docker: Build release image - x86_64 - ROCm"
-        id: build-rocm-release-image
+    id: build-rocm-release-image
-        depends_on:
+    depends_on:
-          - step: block-build-release-images
+      - step: block-build-release-images
-            allow_failure: true
+        allow_failure: true
-          - step: build-rocm-base-wheels
+      - step: build-rocm-base-wheels
-            allow_failure: false
+        allow_failure: false
-        agents:
+    agents:
-          queue: cpu_queue_release
+      queue: cpu_queue_release
-        timeout_in_minutes: 60
+    timeout_in_minutes: 60
-        commands:
+    commands:
-          - |
+      - |
-            set -euo pipefail
+        set -euo pipefail
-            # Login to ECR
+        # Login to ECR
-            aws ecr-public get-login-password --region us-east-1 | \
+        aws ecr-public get-login-password --region us-east-1 | \
-              docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
+          docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
-            # Get ECR image tag from metadata (set by build-rocm-base-wheels)
+        # Get ECR image tag from metadata (set by build-rocm-base-wheels)
-            ECR_IMAGE_TAG="$$(buildkite-agent meta-data get rocm-base-image-tag 2>/dev/null || echo '')"
+        ECR_IMAGE_TAG="$$(buildkite-agent meta-data get rocm-base-image-tag 2>/dev/null || echo '')"
-            if [ -z "$${ECR_IMAGE_TAG}" ]; then
+        if [ -z "$${ECR_IMAGE_TAG}" ]; then
-              echo "ERROR: rocm-base-image-tag metadata not found"
+          echo "ERROR: rocm-base-image-tag metadata not found"
-              echo "This should have been set by the build-rocm-base-wheels job"
+          echo "This should have been set by the build-rocm-base-wheels job"
-              exit 1
+          exit 1
-            fi
+        fi
-            echo "Pulling base Docker image from ECR: $${ECR_IMAGE_TAG}"
+        echo "Pulling base Docker image from ECR: $${ECR_IMAGE_TAG}"
-            # Pull base Docker image from ECR
+        # Pull base Docker image from ECR
-            docker pull "$${ECR_IMAGE_TAG}"
+        docker pull "$${ECR_IMAGE_TAG}"
-            echo "Loaded base image: $${ECR_IMAGE_TAG}"
+        echo "Loaded base image: $${ECR_IMAGE_TAG}"
-            # Pass the base image ECR tag to downstream steps (nightly publish)
+        # Pass the base image ECR tag to downstream steps (nightly publish)
-            buildkite-agent meta-data set "rocm-base-ecr-tag" "$${ECR_IMAGE_TAG}"
+        buildkite-agent meta-data set "rocm-base-ecr-tag" "$${ECR_IMAGE_TAG}"
-            echo "========================================"
+        echo "========================================"
-            echo "Building vLLM ROCm release image with:"
+        echo "Building vLLM ROCm release image with:"
-            echo "  BASE_IMAGE: $${ECR_IMAGE_TAG}"
+        echo "  BASE_IMAGE: $${ECR_IMAGE_TAG}"
-            echo "  BUILDKITE_COMMIT: $${BUILDKITE_COMMIT}"
+        echo "  BUILDKITE_COMMIT: $${BUILDKITE_COMMIT}"
-            echo "========================================"
+        echo "========================================"
-            # Build vLLM ROCm release image using cached base
+        # Build vLLM ROCm release image using cached base
-            DOCKER_BUILDKIT=1 docker build \
+        DOCKER_BUILDKIT=1 docker build \
-              --build-arg max_jobs=16 \
+          --build-arg max_jobs=16 \
-              --build-arg BASE_IMAGE="$${ECR_IMAGE_TAG}" \
+          --build-arg BASE_IMAGE="$${ECR_IMAGE_TAG}" \
-              --build-arg USE_SCCACHE=1 \
+          --build-arg USE_SCCACHE=1 \
-              --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
+          --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
-              --build-arg SCCACHE_REGION_NAME=us-west-2 \
+          --build-arg SCCACHE_REGION_NAME=us-west-2 \
-              --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
+          --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
-              --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm \
+          --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm \
-              --target vllm-openai \
+          --target vllm-openai \
-              --progress plain \
+          --progress plain \
-              -f docker/Dockerfile.rocm .
+          -f docker/Dockerfile.rocm .
-            # Push to ECR
+        # Push to ECR
-            docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm
+        docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm
-            echo ""
+        echo ""
-            echo " Successfully built and pushed ROCm release image"
+        echo " Successfully built and pushed ROCm release image"
-            echo "   Image: public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm"
+        echo "   Image: public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm"
-            echo ""
+        echo ""
-        env:
+    env:
-          DOCKER_BUILDKIT: "1"
+      DOCKER_BUILDKIT: "1"
-          S3_BUCKET: "vllm-wheels"
+      S3_BUCKET: "vllm-wheels"
-      - label: "Publish nightly XPU image to DockerHub"
+  - label: "Publish nightly XPU image to DockerHub"
-        depends_on:
+    depends_on:
-          - create-manifest-xpu
+      - create-manifest-xpu
-        if: build.env("NIGHTLY") == "1"
+    if: build.env("NIGHTLY") == "1"
-        agents:
+    agents:
-          queue: small_cpu_queue_release
+      queue: small_cpu_queue_release
-        commands:
+    commands:
-          - "bash .buildkite/scripts/xpu/push-nightly-builds-xpu.sh"
+      - "bash .buildkite/scripts/xpu/push-nightly-builds-xpu.sh"
-          - "bash .buildkite/scripts/cleanup-nightly-builds.sh nightly- vllm/vllm-openai-xpu"
+      - "bash .buildkite/scripts/cleanup-nightly-builds.sh nightly- vllm/vllm-openai-xpu"
-        plugins:
+    plugins:
-          - docker-login#v3.0.0:
+      - docker-login#v3.0.0:
-              username: vllmbot
+          username: vllmbot
-              password-env: DOCKERHUB_TOKEN
+          password-env: DOCKERHUB_TOKEN
-        env:
+    env:
-          DOCKER_BUILDKIT: "1"
+      DOCKER_BUILDKIT: "1"
-          DOCKERHUB_USERNAME: "vllmbot"
+      DOCKERHUB_USERNAME: "vllmbot"
-      - label: "Publish nightly ROCm image to DockerHub"
+  - label: "Publish nightly ROCm image to DockerHub"
-        depends_on:
+    depends_on:
-          - build-rocm-release-image
+      - build-rocm-release-image
-        if: build.env("NIGHTLY") == "1"
+    if: build.env("NIGHTLY") == "1"
-        agents:
+    agents:
-          queue: small_cpu_queue_release
+      queue: small_cpu_queue_release
-        commands:
+    commands:
-          - "bash .buildkite/scripts/push-nightly-builds-rocm.sh"
+      - "bash .buildkite/scripts/push-nightly-builds-rocm.sh"
-          # Clean up old nightly builds (keep only last 14)
+      # Clean up old nightly builds (keep only last 14)
-          - "bash .buildkite/scripts/cleanup-nightly-builds.sh nightly- vllm/vllm-openai-rocm"
+      - "bash .buildkite/scripts/cleanup-nightly-builds.sh nightly- vllm/vllm-openai-rocm"
-          - "bash .buildkite/scripts/cleanup-nightly-builds.sh base-nightly- vllm/vllm-openai-rocm"
+      - "bash .buildkite/scripts/cleanup-nightly-builds.sh base-nightly- vllm/vllm-openai-rocm"
-        plugins:
+    plugins:
-          - docker-login#v3.0.0:
+      - docker-login#v3.0.0:
-              username: vllmbot
+          username: vllmbot
-              password-env: DOCKERHUB_TOKEN
+          password-env: DOCKERHUB_TOKEN
-        env:
+    env:
-          DOCKER_BUILDKIT: "1"
+      DOCKER_BUILDKIT: "1"
-          DOCKERHUB_USERNAME: "vllmbot"
+      DOCKERHUB_USERNAME: "vllmbot"
  # =============================================================================
  # Publish to DockerHub and PyPI (at the end so all builds complete first)
@@ -1002,8 +931,6 @@ steps:
        depends_on:
          - input-release-version
          - build-wheels
          - build-additional-wheels
          - generate-additional-wheel-indices
      - label: "Upload release wheels to PyPI"
        depends_on:
@@ -3,8 +3,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 #
 # Append a build artifact line to the Buildkite annotation.
-# Usage: annotate-build-artifact.sh <label> <value> <context>
+# Usage: annotate-build-artifact.sh <label> <value>
 set -e
 echo "- **${1}**: \`${2}\`" | \
-  buildkite-agent annotate --append --style 'info' \
+  buildkite-agent annotate --append --style 'info' --context 'release-artifacts'
    --context "${3:?context is required}"
@@ -1,35 +0,0 @@
 #!/usr/bin/env bash
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 #
 # Build the macOS arm64 CPU wheel natively on a macOS agent (the `macmini`
 # queue) into artifacts/dist/ for upload-nightly-wheels.sh.
 set -euo pipefail
 # The macmini queue uses persistent checkouts, so refresh tags for setuptools-scm.
 git fetch --tags --force origin
 # The Rust frontend build needs protoc.
 if ! command -v protoc >/dev/null 2>&1; then
  brew install protobuf
 fi
 # upload-nightly-wheels.sh expects exactly one wheel.
 rm -rf artifacts/dist
 mkdir -p artifacts/dist
 export VLLM_TARGET_DEVICE=cpu
 export VLLM_REQUIRE_RUST_FRONTEND=1
 export MACOSX_DEPLOYMENT_TARGET=11.0
 # uv's CPython is universal2; force an arm64-only build and tag so the wheel
 # isn't mislabelled universal2 and installed on Intel Macs where import fails.
 export ARCHFLAGS="-arch arm64"
 export _PYTHON_HOST_PLATFORM="macosx-11.0-arm64"
 export CMAKE_BUILD_PARALLEL_LEVEL="${CMAKE_BUILD_PARALLEL_LEVEL:-4}"
 uv venv --python 3.12
 uv pip install -r requirements/build/cpu.txt --index-strategy unsafe-best-match
 uv build --wheel --no-build-isolation -o artifacts/dist
 ls -l artifacts/dist/*.whl
@@ -15,9 +15,9 @@ set -euo pipefail
 DEFAULT_REPO_SLUG="vllm-project/vllm"
 DEFAULT_CI_HCL_SOURCE="docker/ci-rocm.hcl"
-DEFAULT_CI_BASE_CONTENT_FILES="requirements/common.txt requirements/rocm.txt requirements/test/rocm.txt docker/Dockerfile.rocm_base docker/ci-rocm.hcl docker/docker-bake-rocm.hcl tools/install_torchcodec_rocm.sh tools/install_protoc.sh rust-toolchain.toml tests/vllm_test_utils .buildkite/scripts/ci-bake-rocm.sh .buildkite/scripts/rocm/build-ci-base.sh"
+DEFAULT_CI_BASE_CONTENT_FILES="requirements/common.txt requirements/rocm.txt requirements/test/rocm.txt docker/Dockerfile.rocm_base docker/ci-rocm.hcl docker/docker-bake-rocm.hcl tools/install_torchcodec_rocm.sh tests/vllm_test_utils .buildkite/scripts/ci-bake-rocm.sh .buildkite/scripts/rocm/build-ci-base.sh"
 DEFAULT_CI_BASE_DOCKERFILE="docker/Dockerfile.rocm"
-DEFAULT_CI_BASE_DOCKERFILE_STAGES="base rust_toolchain_input_0 rust_toolchain_input_1 rust-toolchain-input rust-toolchain build_nixl build_rocshmem build_deepep mori_base ci_base"
+DEFAULT_CI_BASE_DOCKERFILE_STAGES="base build_rixl build_rocshmem build_deepep mori_base ci_base"
 DEFAULT_CI_BASE_METADATA_VERSION="1"
 IMAGE_EXISTED_BEFORE_BUILD=0
@@ -285,7 +285,7 @@ get_content_arg_names() {
    fi | awk 'NF && !seen[$0]++'
 }
-compute_ci_base_content_hash_once() {
+compute_ci_base_content_hash() {
    local -a content_paths=()
    local -a content_args=()
    local dockerfile="${CI_BASE_DOCKERFILE:-}"
@@ -301,8 +301,7 @@ compute_ci_base_content_hash_once() {
        if [[ -n "${dockerfile}" ]]; then
            printf 'dockerfile:%s\n' "${dockerfile}"
            printf 'resolved-build-args:\n'
-            hash_dockerfile_arg_values "${dockerfile}" "${content_args[@]}" \
+            hash_dockerfile_arg_values "${dockerfile}" "${content_args[@]}"
                || return 1
            if [[ -n "${stages}" ]]; then
                printf 'dockerfile-stages:%s\n' "${stages}"
                if [[ -f "${dockerfile}" ]]; then
@@ -315,53 +314,6 @@ compute_ci_base_content_hash_once() {
    } | sha256sum | cut -d' ' -f1
 }
 compute_ci_base_content_hash() {
    local attempts="${CI_BASE_HASH_ATTEMPTS:-3}"
    local delay_secs="${CI_BASE_HASH_RETRY_DELAY:-5}"
    local attempt=0
    local hash=""
    local failed=0
    local -a hashes=()
    if [[ ! "${attempts}" =~ ^[1-9][0-9]*$ ]]; then
        echo "Invalid CI_BASE_HASH_ATTEMPTS: ${attempts}" >&2
        return 1
    fi
    if [[ ! "${delay_secs}" =~ ^[0-9]+$ ]]; then
        echo "Invalid CI_BASE_HASH_RETRY_DELAY: ${delay_secs}" >&2
        return 1
    fi
    for ((attempt = 1; attempt <= attempts; attempt++)); do
        if ! hash=$(compute_ci_base_content_hash_once); then
            echo "ci_base content hash calculation ${attempt}/${attempts} failed" >&2
            failed=1
        else
            hashes+=("${hash}")
            echo "ci_base content hash calculation ${attempt}/${attempts}: ${hash}" >&2
        fi
        if ((attempt < attempts)); then
            sleep "${delay_secs}"
        fi
    done
    if ((failed)) || ((${#hashes[@]} != attempts)); then
        echo "Could not calculate a reliable ci_base content hash" >&2
        return 1
    fi
    for hash in "${hashes[@]:1}"; do
        if [[ "${hash}" != "${hashes[0]}" ]]; then
            echo "ci_base content hash changed between calculations" >&2
            printf '  observed: %s\n' "${hashes[@]}" >&2
            return 1
        fi
    done
    printf '%s\n' "${hashes[0]}"
 }
 extract_dockerfile_arg_default() {
    local dockerfile="$1"
    local arg_name="$2"
@@ -414,11 +366,7 @@ hash_dockerfile_arg_values() {
        printf 'arg:%s=%s\n' "${arg_name}" "${arg_value:-<empty>}"
        if [[ "${arg_name}" == "BASE_IMAGE" && -n "${arg_value}" ]]; then
            digest=$(resolve_image_digest "${arg_value}")
-            if [[ -z "${digest}" ]]; then
+            printf 'arg:%s.digest=%s\n' "${arg_name}" "${digest:-unknown}"
                echo "Failed to resolve digest for BASE_IMAGE=${arg_value}" >&2
                return 1
            fi
            printf 'arg:%s.digest=%s\n' "${arg_name}" "${digest}"
        fi
    done
 }
@@ -816,7 +764,7 @@ configure_ci_base_image_refs() {
        fi
        set_buildkite_metadata "rocm-ci-base-image" "${CI_BASE_IMAGE_TAG}"
        set_buildkite_metadata "rocm-ci-base-image-content" "${content_tag}"
-        set_buildkite_metadata "rocm-ci-base-image-commit" "${CI_BASE_IMAGE_TAG_COMMIT_REF:-}"
+        set_buildkite_metadata "rocm-ci-base-image-commit" "${CI_BASE_IMAGE_TAG_COMMIT:-}"
        set_buildkite_metadata "rocm-ci-base-image-stable" "${CI_BASE_IMAGE_TAG_STABLE:-}"
        return 0
    fi
@@ -1159,8 +1107,8 @@ ci_base_metadata_pairs() {
    metadata_pair "vllm.rocm.nic_backend" "$(resolve_dockerfile_arg_value "${dockerfile}" "NIC_BACKEND")"
    metadata_pair "vllm.rocm.ainic_version" "$(resolve_dockerfile_arg_value "${dockerfile}" "AINIC_VERSION")"
    metadata_pair "vllm.rocm.ubuntu_codename" "$(resolve_dockerfile_arg_value "${dockerfile}" "UBUNTU_CODENAME")"
-    metadata_pair "vllm.rocm.nixl_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "NIXL_REPO")"
+    metadata_pair "vllm.rocm.rixl_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "RIXL_REPO")"
-    metadata_pair "vllm.rocm.nixl_commit" "${NIXL_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "NIXL_BRANCH")}"
+    metadata_pair "vllm.rocm.rixl_commit" "${RIXL_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "RIXL_BRANCH")}"
    metadata_pair "vllm.rocm.ucx_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "UCX_REPO")"
    metadata_pair "vllm.rocm.ucx_commit" "${UCX_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "UCX_BRANCH")}"
    metadata_pair "vllm.rocm.rocshmem_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "ROCSHMEM_REPO")"
@@ -1169,7 +1117,7 @@ ci_base_metadata_pairs() {
    metadata_pair "vllm.rocm.deepep_commit" "${DEEPEP_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_BRANCH")}"
    metadata_pair "vllm.rocm.deepep_nic" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_NIC")"
    metadata_pair "vllm.rocm.deepep_rocm_arch" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_ROCM_ARCH")"
-    metadata_pair "vllm.rocm.nixl_cache_key" "${NIXL_CACHE_KEY:-}"
+    metadata_pair "vllm.rocm.rixl_cache_key" "${RIXL_CACHE_KEY:-}"
    metadata_pair "vllm.rocm.rocshmem_cache_key" "${ROCSHMEM_CACHE_KEY:-}"
    metadata_pair "vllm.rocm.deepep_cache_key" "${DEEPEP_CACHE_KEY:-}"
@@ -1263,24 +1211,12 @@ uses_rocm_csrc_cache() {
    esac
 }
 uses_rocm_rust_cache() {
    case "${TARGET}" in
        rust-rocm-ci|test-rocm-ci|test-rocm-ci-with-wheel|test-rocm-ci-with-artifacts|export-wheel-rocm)
            return 0
            ;;
        *)
            return 1
            ;;
    esac
 }
 compute_rocm_csrc_content_hash() {
    local bake_dir=""
    local dockerfile_rocm=""
    local -a content_paths=(
        "requirements/common.txt"
        "requirements/rocm.txt"
        "pyproject.toml"
        "setup.py"
        "CMakeLists.txt"
        "cmake"
@@ -1324,56 +1260,6 @@ compute_rocm_csrc_content_hash_if_needed() {
    echo "ROCm csrc content cache ref: ${ROCM_CSRC_CONTENT_CACHE_REF}"
 }
 compute_rocm_rust_content_hash() {
    local bake_dir=""
    local dockerfile_rocm=""
    local -a content_paths=(
        "requirements/build/rust.txt"
        "rust/Cargo.lock"
        "rust/Cargo.toml"
        "rust/proto"
        "rust/src"
        "rust-toolchain.toml"
        "tools/build_rust.py"
        "tools/install_protoc.sh"
        "build_rust.sh"
    )
    local -a content_args=()
    bake_dir=$(dirname "${VLLM_BAKE_FILE}")
    dockerfile_rocm="${bake_dir}/Dockerfile.rocm"
    mapfile -t content_args < <(
        get_content_arg_names "${dockerfile_rocm}" "base rust_toolchain_input_0 rust_toolchain_input_1 rust-toolchain-input rust_input_0 rust_input_1 rust-input rust-toolchain rust-build" "${ROCM_RUST_CONTENT_ARGS:-}"
    )
    {
        printf 'rust-input-files-hash:%s\n' "$(compute_content_hash "${content_paths[@]}")"
        printf 'dockerfile:%s\n' "${dockerfile_rocm}"
        printf 'resolved-build-args:\n'
        hash_dockerfile_arg_values "${dockerfile_rocm}" "${content_args[@]}"
        printf 'dockerfile-stages:base rust_toolchain_input_0 rust_toolchain_input_1 rust-toolchain-input rust_input_0 rust_input_1 rust-input rust-toolchain rust-build\n'
        if [[ -f "${dockerfile_rocm}" ]]; then
            hash_dockerfile_stages "${dockerfile_rocm}" "base rust_toolchain_input_0 rust_toolchain_input_1 rust-toolchain-input rust_input_0 rust_input_1 rust-input rust-toolchain rust-build"
        else
            printf 'missing:%s\n' "${dockerfile_rocm}"
        fi
    } | sha256sum | cut -d' ' -f1
 }
 compute_rocm_rust_content_hash_if_needed() {
    local cache_repo="${DOCKERHUB_CACHE_REPO:-rocm/vllm-ci-cache}"
    if [[ "${ROCM_RUST_CONTENT_CACHE:-1}" == "0" ]] || ! uses_rocm_rust_cache; then
        return 0
    fi
    ROCM_RUST_CONTENT_HASH=$(compute_rocm_rust_content_hash)
    ROCM_RUST_CONTENT_CACHE_REF="${cache_repo}:rust-rocm-input-${ROCM_RUST_CONTENT_HASH}"
    export ROCM_RUST_CONTENT_HASH
    export ROCM_RUST_CONTENT_CACHE_REF
    echo "ROCm Rust content cache ref: ${ROCM_RUST_CONTENT_CACHE_REF}"
 }
 write_hcl_string_list_entries() {
    local indent="$1"
    local value=""
@@ -1431,7 +1317,6 @@ write_rocm_build_arg_override() {
                "${CI_BASE_DOCKERFILE_STAGES:-${DEFAULT_CI_BASE_DOCKERFILE_STAGES}}" \
                "${CI_BASE_CONTENT_ARGS:-}"
            get_content_arg_names "${dockerfile_rocm}" "base csrc-build" "${ROCM_CSRC_CONTENT_ARGS:-}"
            get_content_arg_names "${dockerfile_rocm}" "base rust_toolchain_input_0 rust_toolchain_input_1 rust-toolchain-input rust_input_0 rust_input_1 rust-input rust-toolchain rust-build" "${ROCM_RUST_CONTENT_ARGS:-}"
        } | awk 'NF && !seen[$0]++'
    )
@@ -1480,133 +1365,46 @@ validate_cache_export_mode() {
    esac
 }
 validate_content_cache_export_mode() {
    local mode="$1"
    local env_name="$2"
    case "${mode}" in
        missing|always|never)
            ;;
        *)
            echo "Error: ${env_name} must be one of: missing, always, never"
            exit 1
            ;;
    esac
 }
 should_export_content_cache_ref() {
    local cache_ref="$1"
    local cache_name="$2"
    local mode="${ROCM_CONTENT_CACHE_EXPORT_MODE:-missing}"
    case "${mode}" in
        always)
            echo "${cache_name} content cache export mode is always; exporting ${cache_ref}"
            return 0
            ;;
        never)
            echo "${cache_name} content cache export mode is never; not exporting ${cache_ref}"
            return 1
            ;;
        missing|"")
            if docker buildx imagetools inspect "${cache_ref}" >/dev/null 2>&1; then
                echo "${cache_name} content cache exists; not re-exporting ${cache_ref}"
                return 1
            fi
            echo "${cache_name} content cache missing; will export ${cache_ref}"
            return 0
            ;;
        *)
            echo "Error: ROCM_CONTENT_CACHE_EXPORT_MODE must be one of: missing, always, never"
            exit 1
            ;;
    esac
 }
 write_rocm_cache_override() {
    local cache_repo="${DOCKERHUB_CACHE_REPO:-rocm/vllm-ci-cache}"
    local content_cache_export_mode="${ROCM_CONTENT_CACHE_EXPORT_MODE:-missing}"
    local csrc_cache_to_mode="${ROCM_CSRC_CACHE_TO_MODE:-max}"
    local rust_cache_to_mode="${ROCM_RUST_CACHE_TO_MODE:-max}"
    local rocm_cache_to_mode="${ROCM_FINAL_CACHE_TO_MODE:-min}"
-    local -a csrc_content_cache_from=()
+    local -a content_cache_from=()
    local -a rust_content_cache_from=()
    local -a combined_content_cache_from=()
    local -a csrc_cache_to=()
    local -a rust_cache_to=()
    local -a rocm_cache_to=()
    local -a export_wheel_cache_to=()
    local export_csrc_cache=1
    local export_rust_cache=1
-    if ! uses_rocm_csrc_cache && ! uses_rocm_rust_cache; then
+    if ! uses_rocm_csrc_cache; then
        return 0
    fi
    validate_content_cache_export_mode \
        "${content_cache_export_mode}" \
        "ROCM_CONTENT_CACHE_EXPORT_MODE"
    validate_cache_export_mode "${csrc_cache_to_mode}" "ROCM_CSRC_CACHE_TO_MODE"
    validate_cache_export_mode "${rust_cache_to_mode}" "ROCM_RUST_CACHE_TO_MODE"
    validate_cache_export_mode "${rocm_cache_to_mode}" "ROCM_FINAL_CACHE_TO_MODE"
    echo "ROCm content cache export mode: ${content_cache_export_mode}"
    echo "ROCm csrc cache export mode: ${csrc_cache_to_mode}"
    echo "ROCm Rust cache export mode: ${rust_cache_to_mode}"
    echo "ROCm final image cache export mode: ${rocm_cache_to_mode}"
    if [[ -n "${ROCM_CSRC_CONTENT_CACHE_REF:-}" ]]; then
-        csrc_content_cache_from+=("type=registry,ref=${ROCM_CSRC_CONTENT_CACHE_REF}")
+        content_cache_from+=("type=registry,ref=${ROCM_CSRC_CONTENT_CACHE_REF}")
-        if should_export_content_cache_ref "${ROCM_CSRC_CONTENT_CACHE_REF}" "ROCm csrc"; then
+        csrc_cache_to+=(
-            csrc_cache_to+=(
+            "type=registry,ref=${ROCM_CSRC_CONTENT_CACHE_REF},mode=${csrc_cache_to_mode},ignore-error=true"
-                "type=registry,ref=${ROCM_CSRC_CONTENT_CACHE_REF},mode=${csrc_cache_to_mode},ignore-error=true"
+        )
            )
        else
            export_csrc_cache=0
        fi
    fi
    if [[ -n "${ROCM_RUST_CONTENT_CACHE_REF:-}" ]]; then
        rust_content_cache_from+=("type=registry,ref=${ROCM_RUST_CONTENT_CACHE_REF}")
        if should_export_content_cache_ref "${ROCM_RUST_CONTENT_CACHE_REF}" "ROCm Rust"; then
            rust_cache_to+=(
                "type=registry,ref=${ROCM_RUST_CONTENT_CACHE_REF},mode=${rust_cache_to_mode},ignore-error=true"
            )
        else
            export_rust_cache=0
        fi
    fi
    combined_content_cache_from=("${csrc_content_cache_from[@]}" "${rust_content_cache_from[@]}")
    # Docker Hub cache exports are best-effort. A cache-only target failure can
    # otherwise cancel the sibling image target before its manifest is pushed.
    if [[ -n "${BUILDKITE_COMMIT:-}" ]]; then
-        if [[ ${export_csrc_cache} -eq 1 ]]; then
+        csrc_cache_to+=(
-            csrc_cache_to+=(
+            "type=registry,ref=${cache_repo}:csrc-rocm-${BUILDKITE_COMMIT},mode=${csrc_cache_to_mode},ignore-error=true"
-                "type=registry,ref=${cache_repo}:csrc-rocm-${BUILDKITE_COMMIT},mode=${csrc_cache_to_mode},ignore-error=true"
+        )
            )
        fi
        if [[ ${export_rust_cache} -eq 1 ]]; then
            rust_cache_to+=(
                "type=registry,ref=${cache_repo}:rust-rocm-${BUILDKITE_COMMIT},mode=${rust_cache_to_mode},ignore-error=true"
            )
        fi
        rocm_cache_to+=(
            "type=registry,ref=${cache_repo}:rocm-${BUILDKITE_COMMIT},mode=${rocm_cache_to_mode},ignore-error=true"
        )
    fi
    if [[ -n "${ROCM_CACHE_BRANCH_TAG:-}" ]]; then
-        if [[ ${export_csrc_cache} -eq 1 ]]; then
+        csrc_cache_to+=(
-            csrc_cache_to+=(
+            "type=registry,ref=${cache_repo}:csrc-rocm-branch-${ROCM_CACHE_BRANCH_TAG},mode=${csrc_cache_to_mode},ignore-error=true"
-                "type=registry,ref=${cache_repo}:csrc-rocm-branch-${ROCM_CACHE_BRANCH_TAG},mode=${csrc_cache_to_mode},ignore-error=true"
+        )
            )
        fi
        if [[ ${export_rust_cache} -eq 1 ]]; then
            rust_cache_to+=(
                "type=registry,ref=${cache_repo}:rust-rocm-branch-${ROCM_CACHE_BRANCH_TAG},mode=${rust_cache_to_mode},ignore-error=true"
            )
        fi
        rocm_cache_to+=(
            "type=registry,ref=${cache_repo}:rocm-branch-${ROCM_CACHE_BRANCH_TAG},mode=${rocm_cache_to_mode},ignore-error=true"
        )
@@ -1624,7 +1422,7 @@ target "csrc-rocm-ci" {
  cache-from = concat(
    get_cache_from_rocm_csrc(),
 EOF
-        write_hcl_string_list "    " "${csrc_content_cache_from[@]}"
+        write_hcl_string_list "    " "${content_cache_from[@]}"
        cat <<EOF
  )
 EOF
@@ -1632,23 +1430,11 @@ EOF
        cat <<EOF
 }
 target "rust-rocm-ci" {
  cache-from = concat(
    get_cache_from_rocm_rust(),
 EOF
        write_hcl_string_list "    " "${rust_content_cache_from[@]}"
        cat <<EOF
  )
 EOF
        write_hcl_string_list_attr "  " "cache-to" "${rust_cache_to[@]}"
        cat <<EOF
 }
 target "test-rocm-ci" {
  cache-from = concat(
    get_cache_from_rocm(),
 EOF
-        write_hcl_string_list "    " "${combined_content_cache_from[@]}"
+        write_hcl_string_list "    " "${content_cache_from[@]}"
        cat <<EOF
  )
 EOF
@@ -1660,7 +1446,7 @@ target "export-wheel-rocm" {
  cache-from = concat(
    get_cache_from_rocm(),
 EOF
-        write_hcl_string_list "    " "${combined_content_cache_from[@]}"
+        write_hcl_string_list "    " "${content_cache_from[@]}"
        cat <<EOF
  )
 EOF
@@ -1686,7 +1472,7 @@ extract_dependency_pins() {
        return 0
    fi
-    for var in NIXL_BRANCH UCX_BRANCH ROCSHMEM_BRANCH DEEPEP_BRANCH; do
+    for var in RIXL_BRANCH UCX_BRANCH ROCSHMEM_BRANCH DEEPEP_BRANCH; do
        if [[ -n "${!var:-}" ]]; then
            echo "Using provided ${var}: ${!var}"
            continue
@@ -1706,30 +1492,30 @@ extract_dependency_pins() {
 compute_dependency_cache_keys() {
    local bake_dir=""
    local dockerfile_rocm=""
-    local nixl_branch=""
+    local rixl_branch=""
    local ucx_branch=""
    local rocshmem_branch=""
    local deepep_branch=""
-    local nixl_material=""
+    local rixl_material=""
    local rocshmem_material=""
    local deepep_material=""
    bake_dir=$(dirname "${VLLM_BAKE_FILE}")
    dockerfile_rocm="${bake_dir}/Dockerfile.rocm"
-    nixl_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "NIXL_BRANCH")
+    rixl_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "RIXL_BRANCH")
    ucx_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "UCX_BRANCH")
    rocshmem_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "ROCSHMEM_BRANCH")
    deepep_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "DEEPEP_BRANCH")
-    if [[ -n "${nixl_branch}" && -n "${ucx_branch}" ]]; then
+    if [[ -n "${rixl_branch}" && -n "${ucx_branch}" ]]; then
-        nixl_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_nixl")
+        rixl_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_rixl")
-        NIXL_CACHE_KEY=$(
+        RIXL_CACHE_KEY=$(
            compose_dependency_cache_key \
-                "${nixl_branch}-ucx-${ucx_branch}" \
+                "${rixl_branch}-ucx-${ucx_branch}" \
-                "${nixl_material}"
+                "${rixl_material}"
        )
-        export NIXL_CACHE_KEY
+        export RIXL_CACHE_KEY
-        echo "NIXL dependency cache key: ${NIXL_CACHE_KEY}"
+        echo "RIXL dependency cache key: ${RIXL_CACHE_KEY}"
    fi
    if [[ -n "${rocshmem_branch}" ]]; then
@@ -1780,11 +1566,11 @@ dependency_cache_ref_for_target() {
    local cache_repo="${DOCKERHUB_CACHE_REPO:-rocm/vllm-ci-cache}"
    case "${target}" in
-        nixl-rocm-ci)
+        rixl-rocm-ci)
-            if [[ -n "${NIXL_CACHE_KEY:-}" ]]; then
+            if [[ -n "${RIXL_CACHE_KEY:-}" ]]; then
-                printf '%s\n' "${cache_repo}:nixl-rocm-${NIXL_CACHE_KEY}"
+                printf '%s\n' "${cache_repo}:rixl-rocm-${RIXL_CACHE_KEY}"
-            elif [[ -n "${NIXL_BRANCH:-}" ]]; then
+            elif [[ -n "${RIXL_BRANCH:-}" ]]; then
-                printf '%s\n' "${cache_repo}:nixl-rocm-${NIXL_BRANCH}-ucx-${UCX_BRANCH:-}"
+                printf '%s\n' "${cache_repo}:rixl-rocm-${RIXL_BRANCH}-ucx-${UCX_BRANCH:-}"
            fi
            ;;
        rocshmem-rocm-ci)
@@ -1815,7 +1601,7 @@ add_dependency_cache_target() {
 resolve_ci_base_dependency_targets() {
    local mode="${ROCM_DEP_CACHE_EXPORT_MODE:-missing}"
-    local nixl_ref=""
+    local rixl_ref=""
    local rocshmem_ref=""
    local deepep_ref=""
@@ -1824,7 +1610,7 @@ resolve_ci_base_dependency_targets() {
    case "${mode}" in
        always)
            echo "ROCM_DEP_CACHE_EXPORT_MODE=always; exporting all dependency caches serially"
-            for target in nixl-rocm-ci rocshmem-rocm-ci deepep-rocm-ci; do
+            for target in rixl-rocm-ci rocshmem-rocm-ci deepep-rocm-ci; do
                if [[ -n "$(dependency_cache_ref_for_target "${target}")" ]]; then
                    add_dependency_cache_target "${target}"
                fi
@@ -1844,13 +1630,13 @@ resolve_ci_base_dependency_targets() {
            ;;
    esac
-    if [[ "${mode}" != "always" && -n "${NIXL_CACHE_KEY:-}" ]]; then
+    if [[ "${mode}" != "always" && -n "${RIXL_CACHE_KEY:-}" ]]; then
-        nixl_ref=$(dependency_cache_ref_for_target "nixl-rocm-ci")
+        rixl_ref=$(dependency_cache_ref_for_target "rixl-rocm-ci")
-        if dependency_cache_ref_exists "${nixl_ref}"; then
+        if dependency_cache_ref_exists "${rixl_ref}"; then
-            echo "NIXL dependency cache exists: ${nixl_ref}"
+            echo "RIXL dependency cache exists: ${rixl_ref}"
        else
-            echo "NIXL dependency cache missing; will seed: ${nixl_ref}"
+            echo "RIXL dependency cache missing; will seed: ${rixl_ref}"
-            add_dependency_cache_target "nixl-rocm-ci"
+            add_dependency_cache_target "rixl-rocm-ci"
        fi
    fi
@@ -1950,8 +1736,8 @@ confirm_remote_image_push() {
        fi
        if [[ -z "${remote_revision}" \
-              && ${IMAGE_EXISTED_BEFORE_BUILD} -eq 0 ]] \
+              && ${IMAGE_EXISTED_BEFORE_BUILD} -eq 0 \
-            && image_tag_is_commit_scoped; then
+              && image_tag_is_commit_scoped ]]; then
            echo "Remote image exists under a commit-scoped tag; accepting push despite missing revision label."
            return 0
        fi
@@ -2081,57 +1867,36 @@ upload_wheel_artifacts_if_present() {
    local wheel_dir="./wheel-export"
    local artifact_dir="artifacts/vllm-rocm-install"
    local archive_name="vllm-rocm-install.tar.gz"
    local metadata_dir="${wheel_dir}/.vllm-ci-artifact"
    local native_base_image=""
    local whl=""
    local whl_name=""
    local -a wheels=()
    if ! should_upload_wheel_artifacts; then
        return 0
    fi
-    if [[ -d "${wheel_dir}" ]]; then
+    if [[ ! -d "${wheel_dir}" ]] || ! ls "${wheel_dir}"/*.whl >/dev/null 2>&1; then
-        mapfile -t wheels < <(find "${wheel_dir}" -maxdepth 1 -type f -name '*.whl' -print)
+        echo "No ROCm wheel artifacts found in ${wheel_dir}"
-    fi
+        return 0
    if [[ ${#wheels[@]} -ne 1 ]]; then
        echo "Expected exactly one ROCm wheel in ${wheel_dir}; found ${#wheels[@]}" >&2
        return 1
    fi
    whl="${wheels[0]}"
    whl_name=$(basename "${whl}")
    native_base_image="${CI_BASE_IMAGE_TAG_COMMIT_REF:-${CI_BASE_IMAGE:-}}"
    if [[ -z "${native_base_image}" ]]; then
        echo "Native ROCm artifact requires a ci_base image reference" >&2
        return 1
    fi
    echo "--- :package: Uploading ROCm vLLM install artifact"
-    rm -rf "${artifact_dir}" "${metadata_dir}"
+    mkdir -p "${artifact_dir}"
    mkdir -p "${artifact_dir}" "${metadata_dir}"
    printf '%s\n' "${BUILDKITE_COMMIT:-local}" > "${metadata_dir}/commit.txt"
    printf '%s\n' "${native_base_image}" > "${metadata_dir}/native-base-image.txt"
    printf '%s\n' "${CI_BASE_IMAGE:-}" > "${metadata_dir}/ci-base-image.txt"
    printf '%s\n' "${IMAGE_TAG:-}" > "${metadata_dir}/fallback-image.txt"
    printf '%s\n' "${whl_name}" > "${metadata_dir}/wheel-filename.txt"
    tar -C "${wheel_dir}" -czf "${artifact_dir}/${archive_name}" .
    (
        cd "${artifact_dir}"
        sha256sum "${archive_name}" > "${archive_name}.sha256"
    )
    echo "Created ${archive_name}: $(du -sh "${artifact_dir}/${archive_name}" | cut -f1)"
-    cp "${metadata_dir}"/*.txt "${artifact_dir}/"
+    printf '%s\n' "${CI_BASE_IMAGE:-}" > "${artifact_dir}/ci-base-image.txt"
-    cp "${whl}" "${artifact_dir}/${whl_name}"
+    printf '%s\n' "${IMAGE_TAG:-}" > "${artifact_dir}/fallback-image.txt"
-    echo "Copied ${whl_name}: $(du -sh "${artifact_dir}/${whl_name}" | cut -f1)"
+
    for whl in "${wheel_dir}"/*.whl; do
        [[ -f "${whl}" ]] || continue
        whl_name=$(basename "${whl}")
        cp "${whl}" "${artifact_dir}/${whl_name}"
        echo "Copied ${whl_name}: $(du -sh "${artifact_dir}/${whl_name}" | cut -f1)"
    done
    if command -v buildkite-agent >/dev/null 2>&1; then
-        buildkite-agent artifact upload "${artifact_dir}/*" || return 1
+        buildkite-agent artifact upload "${artifact_dir}/*"
        echo "ROCm vLLM install artifacts uploaded to ${artifact_dir}/"
    elif [[ "${BUILDKITE:-false}" == "true" ]]; then
        echo "buildkite-agent not found; cannot upload required ROCm artifacts" >&2
        return 1
    else
        echo "Not in Buildkite, skipping artifact upload"
    fi
@@ -2155,7 +1920,6 @@ main() {
    compute_dependency_cache_keys
    write_ci_base_label_override
    compute_rocm_csrc_content_hash_if_needed
    compute_rocm_rust_content_hash_if_needed
    write_rocm_cache_override
    resolve_ci_base_dependency_targets
    print_bake_config
@@ -2163,11 +1927,6 @@ main() {
        echo "BAKE_PRINT_ONLY=1 set; skipping build"
        return 0
    fi
    if should_upload_wheel_artifacts; then
        # wheel-export is an output directory, not a BuildKit cache. Starting
        # clean prevents a failed/retried export from packaging a stale wheel.
        rm -rf ./wheel-export
    fi
    seed_dependency_caches_if_needed
    run_bake
    upload_wheel_artifacts_if_present
@@ -45,10 +45,8 @@ $PYTHON .buildkite/scripts/generate-nightly-index.py --version "$SUBPATH" --curr
 echo "Uploading indices to $S3_COMMIT_PREFIX"
 aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "$S3_COMMIT_PREFIX"
-# copy to /nightly/ only when enabled for a main branch build that is not a PR
+# copy to /nightly/ only if it is on the main branch and not a PR
-if [[ "${UPDATE_NIGHTLY_INDEX:-1}" == "1" && \
+if [[ "$BUILDKITE_BRANCH" == "main" && "$BUILDKITE_PULL_REQUEST" == "false" ]]; then
      "$BUILDKITE_BRANCH" == "main" && \
      "$BUILDKITE_PULL_REQUEST" == "false" ]]; then
    echo "Uploading indices to overwrite /nightly/"
    aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/nightly/"
 fi
@@ -69,7 +67,7 @@ pure_version="${version%%+*}"
 echo "Pure version (without variant): $pure_version"
 # re-generate and copy to /<pure_version>/ only if it does not have "dev" in the version
-if [[ "${UPDATE_VERSION_INDEX:-1}" == "1" && "$version" != *"dev"* ]]; then
+if [[ "$version" != *"dev"* ]]; then
    echo "Re-generating indices for /$pure_version/"
    rm -rf "${INDICES_OUTPUT_DIR:?}"
    mkdir -p "$INDICES_OUTPUT_DIR"
@@ -1,7 +1,7 @@
 #!/bin/bash
-# This script runs ROCm tests either directly in a native CI pod or inside the
+# This script runs tests inside the corresponding ROCm docker container.
-# corresponding Docker container. Multi-node tests continue to use Docker.
+# It handles both single-node and multi-node test configurations.
 #
 # Multi-node detection: Instead of matching on fragile group names, we detect
 # multi-node jobs structurally by looking for the bracket command syntax
@@ -34,27 +34,10 @@ set -o pipefail
 : "${CLICOLOR_FORCE:=1}"
 : "${PY_COLORS:=1}"
 : "${ROCM_DOCKER_TTY:=1}"
 : "${PYTHONFAULTHANDLER:=1}"
 : "${PYTEST_TIMEOUT:=2400}"
 if [[ " ${PYTEST_ADDOPTS:-} " != *" --color"* ]]; then
  PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }--color=yes"
 fi
-if [[ " ${PYTEST_ADDOPTS:-} " != *" --durations="* ]]; then
+export BUILDKIT_PROGRESS TERM FORCE_COLOR CLICOLOR_FORCE PY_COLORS PYTEST_ADDOPTS ROCM_DOCKER_TTY
  PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }--durations=25"
 fi
 if [[ " ${PYTEST_ADDOPTS:-} " != *" --durations-min="* ]]; then
  PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }--durations-min=1.0"
 fi
 # Dump stacks after 25 minutes, then stop an individual test after 40 minutes.
 if [[ " ${PYTEST_ADDOPTS:-} " != *" faulthandler_timeout="* ]]; then
  PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }-o faulthandler_timeout=1500"
 fi
 if [[ " ${PYTEST_ADDOPTS:-} " != *" --timeout-method="* &&
  " ${PYTEST_ADDOPTS:-} " != *" --timeout-method "* ]]; then
  PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }--timeout-method=thread"
 fi
 export BUILDKIT_PROGRESS TERM FORCE_COLOR CLICOLOR_FORCE PY_COLORS PYTEST_ADDOPTS PYTEST_TIMEOUT ROCM_DOCKER_TTY
 export PYTHONFAULTHANDLER
 # Export Python path for commands that run directly on the host. Containerized
 # tests set this to /vllm-workspace below so spawned Python processes do not
@@ -70,28 +53,6 @@ report_docker_usage() {
  docker system df || true
 }
 clear_ci_orchestration_env() {
  unset -v \
    VLLM_TEST_GROUP_NAME \
    VLLM_CI_REQUIRE_PERSISTENT_HF_CACHE \
    VLLM_CI_ARTIFACT_STEP \
    VLLM_TEST_CACHE \
    VLLM_CI_EXECUTION_MODE \
    VLLM_CI_WORKSPACE \
    VLLM_CI_REQUIRE_WORKSPACE_MOUNT \
    VLLM_TEST_COMMANDS \
    VLLM_CI_BRANCH \
    VLLM_CI_BASE_IMAGE \
    VLLM_CI_FALLBACK_IMAGE \
    VLLM_CI_DOCKER_DISABLED \
    VLLM_CI_ARTIFACT_GLOB \
    VLLM_CI_ARTIFACT_CHECKSUM_GLOB \
    VLLM_CI_EXPECTED_GPU_COUNT \
    VLLM_CI_USE_ARTIFACTS \
    VLLM_CI_RESULTS_ROOT \
    VLLM_ALLOW_DEPRECATED_BEAM_SEARCH
 }
 cleanup_network() {
  local max_nodes=${NUM_NODES:-2}
  for node in $(seq 0 $((max_nodes - 1))); do
@@ -184,11 +145,7 @@ prepare_artifact_image() {
  fi
  cp "${wheel_dir}"/*.whl "${context_dir}/wheels/" || return 1
-  tar -C "${wheel_dir}" \
+  tar -C "${wheel_dir}" --exclude='*.whl' -cf - . \
    --exclude='*.whl' \
    --exclude='.vllm-ci-artifact' \
    --exclude='./.vllm-ci-artifact' \
    -cf - . \
    | tar -C "${workspace_dir}" -xf - || return 1
  cat > "${context_dir}/Dockerfile" <<'EOF'
 ARG BASE_IMAGE
@@ -211,276 +168,6 @@ EOF
  return 0
 }
 is_native_runtime() {
  [[ "${AMD_CI_RUNTIME:-}" == "native" || "${NATIVE_CI:-}" == "true" ]]
 }
 validate_native_workspace() {
  local workspace_dir="${VLLM_CI_WORKSPACE:-/vllm-workspace}"
  local workspace_real=""
  local checkout_real=""
  local workspace_mount=""
  mkdir -p "${workspace_dir}" || return 1
  workspace_real=$(readlink -m "${workspace_dir}") || return 1
  if [[ -n "${BUILDKITE_BUILD_CHECKOUT_PATH:-}" ]]; then
    checkout_real=$(readlink -m "${BUILDKITE_BUILD_CHECKOUT_PATH}") || return 1
    if [[ "${checkout_real}" == "${workspace_real}" \
      || "${checkout_real}" == "${workspace_real}/"* \
      || "${workspace_real}" == "${checkout_real}/"* ]]; then
      echo "Refusing to replace ${workspace_real}; it overlaps the Buildkite checkout ${checkout_real}" >&2
      return 1
    fi
  fi
  if [[ "${VLLM_CI_REQUIRE_WORKSPACE_MOUNT:-1}" == "1" ]]; then
    if ! command -v findmnt >/dev/null 2>&1; then
      echo "findmnt is required to verify the native workspace mount" >&2
      return 1
    fi
    workspace_mount=$(findmnt -n -T "${workspace_real}" -o TARGET 2>/dev/null || true)
    if [[ "$(readlink -m "${workspace_mount:-/}")" != "${workspace_real}" ]]; then
      echo "Native CI requires a dedicated volume mounted at ${workspace_real}" >&2
      return 1
    fi
  fi
 }
 prepare_native_workspace() {
  if [[ "${VLLM_CI_USE_ARTIFACTS:-0}" != "1" ]]; then
    echo "Native CI requires VLLM_CI_USE_ARTIFACTS=1"
    return 1
  fi
  if ! command -v buildkite-agent >/dev/null 2>&1; then
    echo "buildkite-agent not found; cannot download ROCm wheel artifact"
    return 1
  fi
  validate_native_workspace || return 1
  local artifact_glob="${VLLM_CI_ARTIFACT_GLOB:-artifacts/vllm-rocm-install/vllm-rocm-install.tar.gz}"
  local artifact_checksum_glob="${VLLM_CI_ARTIFACT_CHECKSUM_GLOB:-${artifact_glob}.sha256}"
  local artifact_step="${VLLM_CI_ARTIFACT_STEP:-image-build-amd}"
  local archive=""
  local checksum=""
  local download_dir=""
  local metadata_dir=""
  local recorded_base=""
  local recorded_commit=""
  local recorded_wheel=""
  local workspace_dir="${VLLM_CI_WORKSPACE:-/vllm-workspace}"
  local wheel_dir=""
  local attempt=0
  local attempt_dir=""
  local -a archives=()
  local -a checksums=()
  local -a wheels=()
  artifact_work_dir=$(mktemp -d -t vllm-rocm-artifact.XXXXXX) || return 1
  wheel_dir="${artifact_work_dir}/wheels"
  mkdir -p "${wheel_dir}" || return 1
  echo "--- Downloading ROCm wheel artifact from ${artifact_step} (native in-pod)"
  for attempt in 1 2 3; do
    attempt_dir="${artifact_work_dir}/download-${attempt}"
    rm -rf "${attempt_dir}" || return 1
    mkdir -p "${attempt_dir}" || return 1
    if buildkite-agent artifact download \
      "${artifact_glob}" "${attempt_dir}" --step "${artifact_step}" \
      && buildkite-agent artifact download \
        "${artifact_checksum_glob}" "${attempt_dir}" --step "${artifact_step}"; then
      download_dir="${attempt_dir}"
      break
    fi
    echo "Artifact download attempt ${attempt}/3 failed"
    if [[ "${attempt}" -lt 3 ]]; then
      sleep $((attempt * 2))
    fi
  done
  if [[ -z "${download_dir}" ]]; then
    echo "Failed to download ${artifact_glob} and ${artifact_checksum_glob} from ${artifact_step}"
    return 1
  fi
  mapfile -t archives < <(
    find "${download_dir}" -name "vllm-rocm-install.tar.gz" -type f -print
  )
  mapfile -t checksums < <(
    find "${download_dir}" -name "vllm-rocm-install.tar.gz.sha256" -type f -print
  )
  if [[ ${#archives[@]} -ne 1 || ${#checksums[@]} -ne 1 ]]; then
    echo "Expected exactly one ROCm archive and checksum; found ${#archives[@]} archive(s) and ${#checksums[@]} checksum(s)" >&2
    return 1
  fi
  archive="${archives[0]}"
  checksum="${checksums[0]}"
  if [[ "$(dirname "${archive}")" != "$(dirname "${checksum}")" ]]; then
    echo "ROCm archive and checksum were downloaded to different directories" >&2
    return 1
  fi
  (
    cd "$(dirname "${archive}")"
    sha256sum -c "$(basename "${checksum}")"
  ) || return 1
  tar --no-same-owner -xzf "${archive}" -C "${wheel_dir}" || return 1
  mapfile -t wheels < <(
    find "${wheel_dir}" -maxdepth 1 -type f -name '*.whl' -print
  )
  if [[ ${#wheels[@]} -ne 1 ]]; then
    echo "ROCm artifact must contain exactly one top-level wheel; found ${#wheels[@]}" >&2
    return 1
  fi
  metadata_dir="${wheel_dir}/.vllm-ci-artifact"
  for metadata_file in commit.txt native-base-image.txt wheel-filename.txt; do
    if [[ ! -s "${metadata_dir}/${metadata_file}" ]]; then
      echo "ROCm artifact metadata is missing ${metadata_file}" >&2
      return 1
    fi
  done
  for metadata_file in ci-base-image.txt fallback-image.txt; do
    if [[ ! -f "${metadata_dir}/${metadata_file}" ]]; then
      echo "ROCm artifact metadata is missing ${metadata_file}" >&2
      return 1
    fi
  done
  recorded_commit=$(tr -d '\r\n' < "${metadata_dir}/commit.txt")
  recorded_base=$(tr -d '\r\n' < "${metadata_dir}/native-base-image.txt")
  recorded_wheel=$(tr -d '\r\n' < "${metadata_dir}/wheel-filename.txt")
  if [[ -z "${BUILDKITE_COMMIT:-}" || "${recorded_commit}" != "${BUILDKITE_COMMIT}" ]]; then
    echo "ROCm artifact commit ${recorded_commit} does not match ${BUILDKITE_COMMIT:-unset}" >&2
    return 1
  fi
  if [[ -z "${VLLM_CI_BASE_IMAGE:-}" || "${recorded_base}" != "${VLLM_CI_BASE_IMAGE}" ]]; then
    echo "ROCm artifact base ${recorded_base} does not match ${VLLM_CI_BASE_IMAGE:-unset}" >&2
    return 1
  fi
  if [[ "${recorded_wheel}" != "$(basename "${wheels[0]}")" ]]; then
    echo "ROCm artifact wheel manifest ${recorded_wheel} does not match $(basename "${wheels[0]}")" >&2
    return 1
  fi
  for required_dir in tests .buildkite requirements; do
    if [[ ! -d "${wheel_dir}/${required_dir}" ]]; then
      echo "ROCm wheel artifact did not contain ${required_dir}/" >&2
      return 1
    fi
  done
  echo "--- Installing ROCm wheel into pod environment"
  python3 -m pip install --no-deps --force-reinstall "${wheels[0]}" || return 1
  echo "--- Preparing ${workspace_dir} from artifact"
  find "${workspace_dir}" -mindepth 1 -maxdepth 1 -exec rm -rf -- {} + || return 1
  tar -C "${wheel_dir}" \
    --exclude='*.whl' \
    --exclude='.vllm-ci-artifact' \
    --exclude='./.vllm-ci-artifact' \
    -cf - . | tar --no-same-owner -C "${workspace_dir}" -xf - || return 1
  if [[ ! -d "${workspace_dir}/tests" ]]; then
    echo "Failed to stage the native test workspace" >&2
    return 1
  fi
  return 0
 }
 initialize_native_environment() {
  local job_id="${BUILDKITE_JOB_ID:-${BUILDKITE_PARALLEL_JOB:-local}}"
  local job_id_suffix=""
  local native_root=""
  local hf_fstype=""
  local hf_mount=""
  if [[ "$(id -u)" -ne 0 ]]; then
    echo "Native ROCm CI currently requires the ci_base container to run as root" >&2
    return 1
  fi
  job_id="${job_id//[^A-Za-z0-9_.-]/_}"
  job_id_suffix="${job_id##*-}"
  job_id_suffix="${job_id_suffix:0:12}"
  native_root="/tmp/vllm-native-${job_id}"
  TMPDIR="/tmp/vllm-${job_id_suffix}/tmp"
  VLLM_RPC_BASE_PATH="/tmp"
  TORCHINDUCTOR_CACHE_DIR="${native_root}/cache/torchinductor"
  TRITON_CACHE_DIR="${native_root}/cache/triton"
  VLLM_CACHE_ROOT="${native_root}/cache/vllm"
  XDG_CACHE_HOME="${native_root}/cache/xdg"
  : "${HF_HOME:=/home/buildkite-agent/huggingface}"
  # datasets uses POSIX locks that are unsupported by the shared HF NFS cache.
  # Keep processed datasets job-local while retaining the persistent Hub cache.
  HF_DATASETS_CACHE="${native_root}/cache/huggingface/datasets"
  : "${HF_HUB_DOWNLOAD_TIMEOUT:=300}"
  : "${HF_HUB_ETAG_TIMEOUT:=60}"
  export TMPDIR VLLM_RPC_BASE_PATH
  export TORCHINDUCTOR_CACHE_DIR TRITON_CACHE_DIR VLLM_CACHE_ROOT XDG_CACHE_HOME
  export HF_HOME HF_DATASETS_CACHE HF_HUB_DOWNLOAD_TIMEOUT HF_HUB_ETAG_TIMEOUT
  export PYTORCH_ROCM_ARCH=""
  mkdir -p "${TMPDIR}" \
    "${TORCHINDUCTOR_CACHE_DIR}" \
    "${TRITON_CACHE_DIR}" \
    "${VLLM_CACHE_ROOT}" \
    "${XDG_CACHE_HOME}" \
    "${HF_HOME}" \
    "${HF_DATASETS_CACHE}" || return 1
  echo "Native compile caches: VLLM_CACHE_ROOT=${VLLM_CACHE_ROOT} TORCHINDUCTOR_CACHE_DIR=${TORCHINDUCTOR_CACHE_DIR}"
  if [[ "${VLLM_CI_REQUIRE_PERSISTENT_HF_CACHE:-0}" == "1" ]]; then
    if ! command -v findmnt >/dev/null 2>&1; then
      echo "findmnt is required to verify the native Hugging Face cache mount" >&2
      return 1
    fi
    hf_mount=$(findmnt -n -T "${HF_HOME}" -o TARGET 2>/dev/null || true)
    if [[ -z "${hf_mount}" || "${hf_mount}" == "/" ]]; then
      echo "Native CI requires a persistent volume mounted at or above ${HF_HOME}" >&2
      return 1
    fi
  fi
  if command -v findmnt >/dev/null 2>&1; then
    hf_fstype=$(findmnt -n -T "${HF_HOME}" -o FSTYPE 2>/dev/null || true)
  fi
  if [[ "${hf_fstype}" == nfs || "${hf_fstype}" == nfs4 ]]; then
    # Keep hf-xet state local and avoid vectored writes on shared NFS.
    export HF_XET_CACHE="${native_root}/cache/hf-xet"
    export HF_XET_HIGH_PERFORMANCE=0
    export HF_XET_RECONSTRUCTION_USE_VECTORED_WRITE=0
    mkdir -p "${HF_XET_CACHE}" || return 1
    echo "Configured hf-xet for shared ${hf_fstype} cache at ${HF_HOME}"
  fi
 }
 run_native_preflight() {
  local expected_gpus="${VLLM_CI_EXPECTED_GPU_COUNT:-1}"
  if [[ ! "${expected_gpus}" =~ ^[0-9]+$ ]]; then
    echo "Invalid VLLM_CI_EXPECTED_GPU_COUNT=${expected_gpus}" >&2
    return 1
  fi
  python3 -c "import encodings, importlib.metadata as im, importlib.util as iu; [im.version(d) for d in ('transformers', 'torch', 'ray', 'sympy', 'markupsafe', 'vllm')]; missing=[m for m in ('torch.utils.model_zoo', 'transformers.models.nomic_bert', 'ray.dag', 'sympy.physics', 'markupsafe._speedups') if iu.find_spec(m) is None]; assert not missing, missing" || return 1
  if [[ "${expected_gpus}" == "0" ]]; then
    echo "Native CPU-only AMD job: skipping ROCm device validation"
    return 0
  fi
  echo "--- ROCm info"
  rocminfo || return 1
  VLLM_CI_EXPECTED_GPU_COUNT="${expected_gpus}" python3 - <<'PY'
 import os
 import torch
 expected = int(os.environ["VLLM_CI_EXPECTED_GPU_COUNT"])
 assert torch.version.hip, "PyTorch is not a ROCm build"
 assert torch.cuda.is_available(), "ROCm GPU is not available to PyTorch"
 actual = torch.cuda.device_count()
 assert actual == expected, f"Expected {expected} ROCm GPU(s), found {actual}"
 PY
 }
 is_multi_node() {
  local cmds="$1"
  # Primary signal: NUM_NODES environment variable set by the pipeline
@@ -663,58 +350,7 @@ re_quote_pytest_markers() {
 # Main
 ###############################################################################
-if is_native_runtime; then
+# --- GPU initialization ---
  echo "--- Native in-pod ROCm CI (AMD_CI_RUNTIME=${AMD_CI_RUNTIME:-unset}, NATIVE_CI=${NATIVE_CI:-unset})"
  artifact_work_dir=""
  cleanup_native_workspace() {
    if [[ -n "${artifact_work_dir}" ]]; then
      rm -rf "${artifact_work_dir}"
    fi
  }
  trap cleanup_native_workspace EXIT
  if [[ -n "${VLLM_TEST_COMMANDS:-}" ]]; then
    commands="${VLLM_TEST_COMMANDS}"
    commands_source="env"
  else
    commands="$*"
    commands_source="argv"
    if [[ -z "$commands" ]]; then
      echo "Error: No test commands provided for native CI." >&2
      exit 1
    fi
  fi
  if [[ "$commands_source" == "argv" ]]; then
    commands=$(re_quote_pytest_markers "$commands")
  fi
  if is_multi_node "$commands"; then
    echo "Native CI does not support multi-node jobs yet."
    exit 1
  fi
  if ! initialize_native_environment; then
    echo "Failed to initialize the native test environment"
    exit 1
  fi
  if ! prepare_native_workspace; then
    echo "Failed to prepare native test workspace"
    exit 1
  fi
  export PYTHONPATH="${VLLM_CI_WORKSPACE:-/vllm-workspace}"
  echo "Native test commands: $commands"
  run_native_preflight || exit 1
  # Keep AMD CI orchestration variables out of vLLM's runtime environment.
  clear_ci_orchestration_env
  /bin/bash -o pipefail -c "${commands}"
  handle_pytest_exit "$?"
 fi
 # --- GPU initialization for legacy Docker execution ---
 echo "--- ROCm info"
 rocminfo
@@ -816,27 +452,25 @@ fi
 echo "Final commands: $commands"
-standalone_merge_base_env=()
+# The ROCm test image often ships /vllm-workspace without .git (artifact tarball unpack).
-if [[ "$commands" == *python_only_compile.sh* ]]; then
+# tests/standalone_tests/python_only_compile.sh uses merge-base(HEAD, origin/main) for
-  # The ROCm test image often ships /vllm-workspace without .git. Resolve the
+# wheels.vllm.ai; compute on the agent (full git checkout) and pass into the container.
-  # wheels.vllm.ai commit from the agent checkout for this test only.
+vllm_standalone_merge_base=""
-  vllm_standalone_merge_base=""
+checkout="${BUILDKITE_BUILD_CHECKOUT_PATH:-}"
-  checkout="${BUILDKITE_BUILD_CHECKOUT_PATH:-}"
+if [[ -z "${checkout}" || ! -d "${checkout}" ]]; then
-  if [[ -z "${checkout}" || ! -d "${checkout}" ]]; then
+  checkout="."
    checkout="."
  fi
  # Pass safe.directory per-command because Buildkite uses mixed user IDs.
  if git -c "safe.directory=${checkout}" -C "${checkout}" rev-parse --is-inside-work-tree >/dev/null 2>&1; then
    vllm_standalone_merge_base="$(
      git -c "safe.directory=${checkout}" -C "${checkout}" merge-base HEAD origin/main 2>/dev/null || true
    )"
  fi
  if [[ -z "${vllm_standalone_merge_base}" ]]; then
    vllm_standalone_merge_base="${BUILDKITE_COMMIT:-}"
  fi
  echo "INFO: passing CI_STANDALONE_MERGE_BASE into container: ${vllm_standalone_merge_base}"
  standalone_merge_base_env=(-e "CI_STANDALONE_MERGE_BASE=${vllm_standalone_merge_base}")
 fi
 # Pass safe.directory per-command (-c) because buildkite runs will always fail
 # the next check on git 2.35.2+ due to mixed uses of root and buildkite-agent/uids.
 if git -c "safe.directory=${checkout}" -C "${checkout}" rev-parse --is-inside-work-tree >/dev/null 2>&1; then
  vllm_standalone_merge_base="$(
    git -c "safe.directory=${checkout}" -C "${checkout}" merge-base HEAD origin/main 2>/dev/null || true
  )"
 fi
 if [[ -z "${vllm_standalone_merge_base}" ]]; then
  vllm_standalone_merge_base="${BUILDKITE_COMMIT:-}"
 fi
 echo "INFO: passing VLLM_STANDALONE_MERGE_BASE into container: ${vllm_standalone_merge_base}"
 MYPYTHONPATH="/vllm-workspace"
@@ -867,7 +501,6 @@ else
 fi
 # --- Route: multi-node vs single-node ---
 clear_ci_orchestration_env
 if is_multi_node "$commands"; then
  echo "--- Multi-node job detected"
  export DCKR_VER=$(docker --version | sed 's/Docker version \(.*\), build .*/\1/')
@@ -956,9 +589,7 @@ else
    -e FORCE_COLOR \
    -e CLICOLOR_FORCE \
    -e PY_COLORS \
    -e PYTHONFAULTHANDLER \
    -e PYTEST_ADDOPTS \
    -e PYTEST_TIMEOUT \
    -v "${HF_CACHE}:${HF_MOUNT}" \
    -e "HF_HOME=${HF_MOUNT}" \
    -e "PYTHONPATH=${MYPYTHONPATH}" \
@@ -968,7 +599,7 @@ else
    -e "VLLM_CACHE_ROOT=${CONTAINER_CACHE_ROOT}/vllm" \
    -e "XDG_CACHE_HOME=${CONTAINER_CACHE_ROOT}/xdg" \
    -e "PYTORCH_ROCM_ARCH=" \
-    "${standalone_merge_base_env[@]}" \
+    -e "VLLM_STANDALONE_MERGE_BASE=${vllm_standalone_merge_base}" \
    --name "${container_name}" \
    "${image_name}" \
    /bin/bash -c "${CONTAINER_PREFLIGHT} && ${commands}"
@@ -3,9 +3,8 @@ set -euox pipefail
 export VLLM_CPU_KVCACHE_SPACE=1 
 export VLLM_CPU_CI_ENV=1
-# Skip torch.compile via vLLM's --enforce-eager flag (passed below) instead of
+# Reduce sub-processes for acceleration
-# TORCH_COMPILE_DISABLE=1, which torch 2.12 no longer treats as a silent no-op
+export TORCH_COMPILE_DISABLE=1 
 # when callers specify fullgraph=True.
 export VLLM_ENABLE_V1_MULTIPROCESSING=0
 SDE_ARCHIVE="sde-external-10.7.0-2026-02-18-lin.tar.xz"
@@ -50,15 +49,15 @@ wait_for_pid_and_check_log() {
 }
 # Test Sky Lake (AVX512F)
-./sde/sde64 -skl -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 --enforce-eager > test_0.log 2>&1 &
+./sde/sde64 -skl -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 > test_0.log 2>&1 &
 PID_TEST_0=$!
 # Test Cascade Lake (AVX512F + VNNI)
-./sde/sde64 -clx -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 --enforce-eager > test_1.log 2>&1 &
+./sde/sde64 -clx -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 > test_1.log 2>&1 &
 PID_TEST_1=$!
 # Test Cooper Lake (AVX512F + VNNI + BF16)
-./sde/sde64 -cpx -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 --enforce-eager > test_2.log 2>&1 &
+./sde/sde64 -cpx -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 > test_2.log 2>&1 &
 PID_TEST_2=$!
 wait_for_pid_and_check_log $PID_TEST_0 test_0.log
@@ -40,9 +40,7 @@ function cpu_tests() {
    pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py
    pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
    pytest -x -v -s tests/kernels/moe/test_cpu_int4_moe.py
-    pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py
+    pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py"
    pytest -x -v -s tests/kernels/mamba/test_causal_conv1d.py
    pytest -x -v -s tests/kernels/mamba/test_mamba_ssm.py"
  # skip tests requiring model downloads if HF_TOKEN is not set
  # due to rate-limits
@@ -99,4 +97,3 @@ function cpu_tests() {
 # All of CPU tests are expected to be finished less than 40 mins.
 export -f cpu_tests
 timeout 2h bash -c cpu_tests
@@ -35,7 +35,6 @@ case "${test_suite}" in
    pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py --ignore=v1/worker/test_worker_memory_snapshot.py
    pytest -v -s v1/structured_output
    pytest -v -s v1/test_serial_utils.py
    pytest -v -s v1/e2e/general/test_correctness_sliding_window.py --deselect="tests/v1/e2e/general/test_correctness_sliding_window.py::test_sliding_window_retrieval[True-1-5-google/gemma-3-1b-it]"
    pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py --ignore=v1/spec_decode/test_speculators_correctness.py
    pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py --ignore=v1/kv_connector/unit/test_offloading_connector.py
    ;;
@@ -369,7 +369,7 @@ export HF_TOKEN ZE_AFFINITY_MASK
    -e CMDS \
    --name "${container_name}" \
    "${IMAGE}" \
-    bash -c 'set -e; echo "ZE_AFFINITY_MASK is ${ZE_AFFINITY_MASK:-}"; eval "$CMDS"' \
+    bash -c 'set -e; source /opt/intel/oneapi/setvars.sh --force; source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force; echo "ZE_AFFINITY_MASK is ${ZE_AFFINITY_MASK:-}"; eval "$CMDS"' \
    >/dev/null
 } 9>/tmp/docker-pull.lock
@@ -13,18 +13,6 @@ metadata_get() {
    fi
 }
 use_ci_base_if_present() {
    local ci_base_image=""
    ci_base_image="$(metadata_get rocm-ci-base-image)"
    if [[ -z "${ci_base_image}" ]]; then
        return 1
    fi
    export CI_BASE_IMAGE="${ci_base_image}"
    echo "Using ROCm ci_base image selected by the preceding build step: ${CI_BASE_IMAGE}"
 }
 use_refreshed_base_if_present() {
    local base_refreshed=""
@@ -34,12 +22,15 @@ use_refreshed_base_if_present() {
    fi
    export BASE_IMAGE
    export CI_BASE_IMAGE
    export IMAGE_TAG_LATEST
    BASE_IMAGE="$(metadata_get rocm-base-image)"
    CI_BASE_IMAGE="$(metadata_get rocm-ci-base-image)"
    IMAGE_TAG_LATEST="$(metadata_get rocm-ci-image-descriptive)"
    echo "Using refreshed ROCm base image for test image: ${BASE_IMAGE}"
    echo "Using refreshed ROCm ci_base image for test image: ${CI_BASE_IMAGE}"
    if [[ -n "${IMAGE_TAG_LATEST}" ]]; then
        echo "Also tagging full ROCm CI image as: ${IMAGE_TAG_LATEST}"
    fi
@@ -50,8 +41,6 @@ use_refreshed_base_if_present() {
 main() {
    local base_refreshed=0
    use_ci_base_if_present || true
    if use_refreshed_base_if_present; then
        base_refreshed=1
    fi
@@ -6,14 +6,8 @@ set -ex
 # manylinux platform tag with auditwheel.
 # Index generation is handled separately by generate-and-upload-nightly-index.sh.
-# auditwheel is Linux-only; macOS wheels already carry a valid tag, so skip the
+# shellcheck source=lib/manylinux.sh
-# manylinux retag for them.
+source .buildkite/scripts/lib/manylinux.sh
 WHEEL_PLATFORM="${VLLM_WHEEL_PLATFORM:-linux}"
 if [[ "$WHEEL_PLATFORM" == "linux" ]]; then
  # shellcheck source=lib/manylinux.sh
  source .buildkite/scripts/lib/manylinux.sh
 fi
 BUCKET="vllm-wheels"
 SUBPATH=$BUILDKITE_COMMIT
@@ -33,10 +27,8 @@ wheel="${wheel_files[0]}"
 # ========= detect manylinux tag and rename ==========
-if [[ "$WHEEL_PLATFORM" == "linux" ]]; then
+wheel="$(apply_manylinux_tag "$wheel")"
-  wheel="$(apply_manylinux_tag "$wheel")"
+echo "Renamed wheel to: $wheel"
  echo "Renamed wheel to: $wheel"
 fi
 # Extract the version from the wheel
 version=$(unzip -p "$wheel" '**/METADATA' | grep '^Version: ' | cut -d' ' -f2)
@@ -113,8 +113,8 @@ $PYTHON .buildkite/scripts/generate-nightly-index.py \
 echo "Uploading indices to $S3_COMMIT_PREFIX"
 aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "$S3_COMMIT_PREFIX"
-# Only scheduled nightly builds should update the moving nightly index.
+# Update rocm/nightly/ if on main branch and not a PR
-if [[ "${NIGHTLY:-0}" == "1" ]]; then
+if [[ "$BUILDKITE_BRANCH" == "main" && "$BUILDKITE_PULL_REQUEST" == "false" ]] || [[ "$NIGHTLY" == "1" ]]; then
    echo "Updating rocm/nightly/ index..."
    aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/rocm/nightly/"
 fi
@@ -147,7 +147,7 @@ echo ""
 echo "Install command (by commit):"
 echo "  pip install vllm --extra-index-url https://${BUCKET}.s3.amazonaws.com/$ROCM_SUBPATH/"
 echo ""
-if [[ "${NIGHTLY:-0}" == "1" ]]; then
+if [[ "$BUILDKITE_BRANCH" == "main" ]] || [[ "$NIGHTLY" == "1" ]]; then
    echo "Install command (nightly):"
    echo "  pip install vllm --extra-index-url https://${BUCKET}.s3.amazonaws.com/rocm/nightly/"
 fi
@@ -16,9 +16,8 @@ steps:
  parallelism: 2
  mirror:
    amd:
-      dind: false
+      device: mi325_1
-      device: mi300_1
+      timeout_in_minutes: 95
      timeout_in_minutes: 125
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -4,7 +4,7 @@ depends_on:
 steps:
 - label: Basic Correctness
  key: basic-correctness
-  timeout_in_minutes: 68
+  timeout_in_minutes: 45
  device: h200_18gb
  source_file_dependencies:
  - vllm/
@@ -18,8 +18,7 @@ steps:
  - pytest -v -s basic_correctness/test_cpu_offload.py
  mirror:
    amd:
-      dind: false
+      device: mi325_1
-      device: mi300_1
+      timeout_in_minutes: 70
      timeout_in_minutes: 60
      depends_on:
      - image-build-amd
@@ -4,7 +4,7 @@ depends_on:
 steps:
 - label: Benchmarks CLI Test
  key: benchmarks-cli-test
-  timeout_in_minutes: 45
+  timeout_in_minutes: 30
  device: h200_18gb
  source_file_dependencies:
  - vllm/
@@ -13,9 +13,7 @@ steps:
  - pytest -v -s benchmarks/
  mirror:
    amd:
      dind: false
      device: mi300_1
      timeout_in_minutes: 40
      depends_on:
      - image-build-amd
@@ -16,10 +16,8 @@ steps:
  commands:
    - pytest -v -s cuda/test_cuda_context.py
    - pytest -v -s cuda/test_platform_no_cuda_init.py
    - pytest -v -s cuda/test_cuda_compatibility_path.py
 - label: Cudagraph
  device: h200_35gb
  key: cudagraph
  timeout_in_minutes: 30
  source_file_dependencies:
@@ -27,10 +25,7 @@ steps:
  - vllm/v1/cudagraph_dispatcher.py
  - vllm/config/compilation.py
  - vllm/compilation
  - vllm/v1/worker/encoder_cudagraph.py
  - vllm/v1/worker/encoder_cudagraph_defs.py
  commands:
    - pytest -v -s v1/cudagraph/test_cudagraph_dispatch.py
    - pytest -v -s v1/cudagraph/test_cudagraph_mode.py
    - pytest -v -s v1/cudagraph/test_breakable_cudagraph.py
    - pytest -v -s v1/cudagraph/test_encoder_cudagraph.py
@@ -15,9 +15,8 @@ steps:
    - bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
  mirror:
    amd:
      dind: false
      device: mi300_4
-      timeout_in_minutes: 60
+      timeout_in_minutes: 85
      depends_on:
        - image-build-amd
      source_file_dependencies:
@@ -66,9 +65,8 @@ steps:
    - DP_EP=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
  mirror:
    amd:
      dind: false
      device: mi300_4
-      timeout_in_minutes: 40
+      timeout_in_minutes: 60
      depends_on:
        - image-build-amd
      source_file_dependencies:
@@ -92,9 +90,8 @@ steps:
    - CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
  mirror:
    amd:
      dind: false
      device: mi300_4
-      timeout_in_minutes: 60
+      timeout_in_minutes: 85
      depends_on:
        - image-build-amd
      source_file_dependencies:
@@ -118,9 +115,8 @@ steps:
    - HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
  mirror:
    amd:
      dind: false
      device: mi300_4
-      timeout_in_minutes: 55
+      timeout_in_minutes: 80
      depends_on:
        - image-build-amd
      source_file_dependencies:
@@ -131,22 +127,6 @@ steps:
        - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
        - HYBRID_SSM=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
 - label: NixlConnector PD edge case test (2 GPUs)
  key: nixlconnector-pd-edge-cases-2-gpus
  timeout_in_minutes: 40
  working_dir: "/vllm-workspace/tests"
  num_devices: 2
  source_file_dependencies:
    - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
    - vllm/v1/core/sched/
    - tests/v1/kv_connector/nixl_integration/
  env:
    PREFILL_GPU_ID: "0"
    DECODE_GPU_ID: "1"
  commands:
    - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
    - bash v1/kv_connector/nixl_integration/run_edge_case_test.sh
 - label: Hybrid SSM NixlConnector PD prefix cache test (2 GPUs)
  key: hybrid-ssm-nixlconnector-pd-prefix-cache-2-gpus
  timeout_in_minutes: 25
@@ -191,9 +171,8 @@ steps:
    - bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh
  mirror:
    amd:
      dind: false
      device: mi300_2
-      timeout_in_minutes: 45
+      timeout_in_minutes: 70
      depends_on:
        - image-build-amd
      source_file_dependencies:
@@ -203,7 +182,7 @@ steps:
        - vllm/platforms/rocm.py
      commands:
        - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
-        - KV_CACHE_MEMORY_BYTES=8G ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh
+        - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh
 - label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs)
  key: multiconnector-nixl-offloading-pd-edge-cases-2-gpus
@@ -219,25 +198,3 @@ steps:
  commands:
    - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
    - bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
 # P TP 4 - D DPEP 4 test case for DSv4-Flash
 - label: DSv4-Flash Disaggregated DP EP
  key: dsv4-flash-disaggregated
  timeout_in_minutes: 60
  device: h200
  optional: true
  working_dir: "/vllm-workspace/tests"
  num_devices: 8
  env:
    ENABLE_HMA_FLAG: "1"
    DP_EP: "1"
    GPU_MEMORY_UTILIZATION: "0.85"
    PREFILLER_TP_SIZE: "4"
    DECODER_TP_SIZE: "4"
    PREFILL_BLOCK_SIZE: "256"
    DECODE_BLOCK_SIZE: "256"
    MODEL_NAMES: "deepseek-ai/DeepSeek-V4-Flash"
    VLLM_SERVE_EXTRA_ARGS: "--trust-remote-code,--kv-cache-dtype,fp8"
  commands:
    - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
    - bash v1/kv_connector/nixl_integration/run_accuracy_test.sh
@@ -39,9 +39,7 @@ steps:
  - DP_SIZE=2 pytest -v -s entrypoints/openai/test_multi_api_servers.py
  mirror:
    amd:
      dind: false
      device: mi300_2
      timeout_in_minutes: 45
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -28,9 +28,8 @@ steps:
  - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py test_jit_monitor.py
  mirror:
    amd:
-      dind: false
+      device: mi325_1
-      device: mi300_1
+      timeout_in_minutes: 50
      timeout_in_minutes: 40
      depends_on:
      - image-build-amd
@@ -40,21 +39,19 @@ steps:
  source_file_dependencies:
    - vllm/v1/engine/
    - tests/v1/engine/
    - tests/v1/test_tensor_ipc_queue.py
  commands:
    - pytest -v -s v1/engine/test_preprocess_error_handling.py
    - pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py
    - pytest -v -s v1/test_tensor_ipc_queue.py
  mirror:
    amd:
-      device: mi250_1
+      device: mi325_1
-      timeout_in_minutes: 45
+      timeout_in_minutes: 55
      depends_on:
      - image-build-amd
 - label: e2e Scheduling (1 GPU)
  key: e2e-scheduling-1-gpu
-  timeout_in_minutes: 53
+  timeout_in_minutes: 35
  device: h200_18gb
  source_file_dependencies:
    - vllm/v1/
@@ -63,8 +60,8 @@ steps:
    - pytest -v -s v1/e2e/general/test_async_scheduling.py
  mirror:
    amd:
-      device: mi250_1
+      device: mi325_1
-      timeout_in_minutes: 55
+      timeout_in_minutes: 70
      depends_on:
      - image-build-amd
@@ -79,8 +76,8 @@ steps:
    - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py
  mirror:
    amd:
-      device: mi250_1
+      device: mi325_1
-      timeout_in_minutes: 50
+      timeout_in_minutes: 60
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -117,9 +114,7 @@ steps:
    - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "tensor_parallelism"
  mirror:
    amd:
      dind: false
      device: mi300_2
      timeout_in_minutes: 30
      depends_on:
      - image-build-amd
@@ -3,7 +3,6 @@ depends_on:
  - image-build
 steps:
 - label: Entrypoints Unit Tests
  device: h200_35gb
  key: entrypoints-unit-tests
  timeout_in_minutes: 25
  working_dir: "/vllm-workspace/tests"
@@ -16,7 +15,6 @@ steps:
  - pytest -v -s entrypoints/weight_transfer
 - label: Entrypoints Integration (LLM)
  device: h200_35gb
  key: entrypoints-integration-llm
  timeout_in_minutes: 60
  working_dir: "/vllm-workspace/tests"
@@ -30,16 +28,16 @@ steps:
  - pytest -v -s entrypoints/llm/offline_mode # Needs to avoid interference with other tests
  mirror:
    amd:
-      dind: false
+      device: mi325_1
-      device: mi300_1
+      # TODO(akaratza): Test after Torch >= 2.12 bump
-      timeout_in_minutes: 55
+      soft_fail: true
      depends_on:
      - image-build-amd
 - label: Entrypoints Integration (API Server)
  key: entrypoints-integration-api-server
  device: h200_35gb
-  timeout_in_minutes: 75
+  timeout_in_minutes: 50
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -52,16 +50,13 @@ steps:
  - pytest -v -s entrypoints/scale_out
  mirror:
    amd:
-      dind: false
+      device: mi325_1
      device: mi300_1
      timeout_in_minutes: 65
      depends_on:
      - image-build-amd
 - label: Entrypoints Integration (API Server OpenAI - Part 1)
  device: h200_35gb
  key: entrypoints-integration-api-server-openai-part-1
-  timeout_in_minutes: 68
+  timeout_in_minutes: 45
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -72,16 +67,14 @@ steps:
  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness
  mirror:
    amd:
-      dind: false
+      device: mi325_1
      device: mi300_1
      timeout_in_minutes: 65
      depends_on:
      - image-build-amd
 - label: Entrypoints Integration (API Server OpenAI - Part 2)
  device: h200_35gb
  key: entrypoints-integration-api-server-openai-part-2
-  timeout_in_minutes: 83
+  timeout_in_minutes: 45
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -93,14 +86,12 @@ steps:
  - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
  mirror:
    amd:
-      dind: false
+      device: mi325_1
-      device: mi300_1
+      timeout_in_minutes: 80
      timeout_in_minutes: 70
      depends_on:
      - image-build-amd
 - label: Entrypoints Integration (API Server Generate)
  device: h200_35gb
  key: entrypoints-integration-api-server-generate
  timeout_in_minutes: 50
  working_dir: "/vllm-workspace/tests"
@@ -117,14 +108,12 @@ steps:
  - pytest -v -s entrypoints/anthropic
  mirror:
    amd:
-      dind: false
+      device: mi325_1
      device: mi300_1
      timeout_in_minutes: 65
      depends_on:
      - image-build-amd
 - label: Entrypoints Integration (Responses API)
  device: h200_35gb
  key: entrypoints-integration-responses-api
  timeout_in_minutes: 50
  working_dir: "/vllm-workspace/tests"
@@ -159,9 +148,8 @@ steps:
  - pytest -v -s entrypoints/multimodal
 - label: Entrypoints Integration (Pooling)
  device: h200_35gb
  key: entrypoints-integration-pooling
-  timeout_in_minutes: 75
+  timeout_in_minutes: 50
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -181,9 +169,7 @@ steps:
  - pytest -s entrypoints/openai/correctness/
  mirror:
    amd:
-      dind: false
+      device: mi325_1
      device: mi300_1
      timeout_in_minutes: 30
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -16,9 +16,7 @@ steps:
  - pytest -v -s distributed/test_eplb_utils.py
  mirror:
    amd:
      dind: false
      device: mi300_1
      timeout_in_minutes: 30
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -52,5 +50,4 @@ steps:
  - vllm/compilation/
  - tests/distributed/
  commands:
  - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
  - pytest -v -s distributed/test_elastic_ep.py
@@ -1,26 +0,0 @@
 group: Fault Tolerance
 depends_on:
  - image-build
 steps:
 - label: Fault Tolerance E2E (2xH100)
  key: fault-tolerance-e2e-2xh100
  timeout_in_minutes: 35
  device: h100
  num_devices: 2
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/v1/fault_tolerance/
  - vllm/v1/worker/sentinel/
  - vllm/entrypoints/serve/fault_tolerance/
  - vllm/distributed/elastic_ep/
  - vllm/distributed/device_communicators/
  - vllm/v1/engine/
  - vllm/v1/worker/
  - tests/v1/fault_tolerance/
  - tests/v1/distributed/test_external_lb_dp.py
  commands:
  # Base image has no nixl; install it or has_nixl_ep() skips the tests.
  - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
  # https://github.com/NVIDIA/nccl/issues/1838
  - export NCCL_CUMEM_HOST_ENABLE=0
  - pytest -v -s v1/fault_tolerance/test_fault_tolerance_e2e.py
@@ -15,7 +15,6 @@ steps:
    - pytest -v -s tests/kernels/ir
 - label: Kernels Core Operation Test
  device: h200_35gb
  key: kernels-core-operation-test
  timeout_in_minutes: 120
  source_file_dependencies:
@@ -61,45 +60,9 @@ steps:
  source_file_dependencies:
  - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
  - vllm/models/deepseek_v4/common/ops/
  - vllm/models/deepseek_v4/nvidia/
  - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
  - tests/models/test_deepseek_v4_mega_moe.py
  commands:
    - pytest -v -s kernels/test_fused_deepseek_v4_*.py
    - pytest -v -s models/test_deepseek_v4_mega_moe.py
 # Catch-all for test files at the tests/kernels root. This job collects
 # the whole root so new files are wired by default.
 # Files with dedicated jobs elsewhere in this file are excluded via --ignore
 # (test_kda, test_bf16x3_router_gemm_cutedsl and test_ll_bf16_gemm run in
 # their own jobs / Kernels (B200)).
 - label: Kernels Root Misc Test (B200)
  key: kernels-root-misc-test-b200
  timeout_in_minutes: 45
  device: b200-k8s
  source_file_dependencies:
  - csrc/
  - vllm/
  - tests/kernels/
  commands:
    - pytest -v -s kernels/
        --ignore=kernels/attention
        --ignore=kernels/core
        --ignore=kernels/helion
        --ignore=kernels/ir
        --ignore=kernels/mamba
        --ignore=kernels/moe
        --ignore=kernels/quantization
        --ignore=kernels/test_concat_mla_q.py
        --ignore=kernels/test_fused_qk_norm_rope_gate.py
        --ignore=kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
        --ignore=kernels/test_top_k_per_row.py
        --ignore=kernels/test_kda.py
        --ignore=kernels/test_bf16x3_router_gemm_cutedsl.py
        --ignore=kernels/test_ll_bf16_gemm.py
        --ignore=kernels/test_shuffle_rows.py
    # BROKEN on main, pending kernel fixes (B200):
    #   test_shuffle_rows.py (1: test_shuffle_rows_edge_cases)
 - label: Kernels Attention Test %N
  key: kernels-attention-test
@@ -116,8 +79,7 @@ steps:
  parallelism: 2
  mirror:
    amd:
-      dind: false
+      device: mi325_1
      device: mi300_1
      timeout_in_minutes: 90
      depends_on:
      - image-build-amd
@@ -155,9 +117,7 @@ steps:
  parallelism: 2
  mirror:
    amd:
-      dind: false
+      device: mi325_1
      device: mi300_1
      timeout_in_minutes: 120
      source_file_dependencies:
      - csrc/quantization/
      - vllm/model_executor/layers/quantization
@@ -187,9 +147,8 @@ steps:
  parallelism: 5
  mirror:
    amd:
-      dind: false
+      device: mi325_1
-      device: mi300_1
+      timeout_in_minutes: 65
      timeout_in_minutes: 55
      source_file_dependencies:
      - csrc/quantization/cutlass_w8a8/moe/
      - csrc/moe/
@@ -204,7 +163,6 @@ steps:
      - image-build-amd
 - label: Kernels Mamba Test
  device: h200_35gb
  key: kernels-mamba-test
  timeout_in_minutes: 40
  source_file_dependencies:
@@ -214,6 +172,17 @@ steps:
  commands:
    - pytest -v -s kernels/mamba
 - label: Kernels KDA Test
  timeout_in_minutes: 25
  device: h200_18gb
  source_file_dependencies:
  - vllm/model_executor/layers/fla/ops/kda.py
  - vllm/model_executor/layers/fla/ops/chunk_delta_h.py
  - vllm/model_executor/layers/fla/ops/l2norm.py
  - tests/kernels/test_kda.py
  commands:
    - pytest -v -s kernels/test_kda.py
 - label: Kernels DeepGEMM Test (H100)
  key: kernels-deepgemm-test-h100
  timeout_in_minutes: 35
@@ -263,15 +232,6 @@ steps:
  - vllm/v1/attention/backends/mla/flashinfer_mla.py
  - vllm/v1/attention/selector.py
  - vllm/platforms/cuda.py
  - vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py
  - vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py
  - vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_splitk.py
  - vllm/cute_utils/
  - vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/
  - vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_cutedsl.py
  - tests/kernels/mamba/test_gdn_prefill_cutedsl.py
  - tests/kernels/test_bf16x3_router_gemm_cutedsl.py
  - tests/kernels/test_ll_bf16_gemm.py
  - tests/kernels/test_top_k_per_row.py
  commands:
    - nvidia-smi
@@ -300,9 +260,6 @@ steps:
    - pytest -v -s tests/kernels/moe/test_flashinfer_moe.py
    - pytest -v -s tests/kernels/moe/test_trtllm_nvfp4_moe.py
    - pytest -v -s tests/kernels/moe/test_cutedsl_moe.py
    - pytest -v -s tests/kernels/mamba/test_gdn_prefill_cutedsl.py
    - pytest -v -s tests/kernels/test_bf16x3_router_gemm_cutedsl.py
    - pytest -v -s tests/kernels/test_ll_bf16_gemm.py
    # e2e
    - pytest -v -s tests/models/quantization/test_nvfp4.py
@@ -14,9 +14,8 @@ steps:
  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.txt
  mirror:
    amd:
-      dind: false
+      device: mi325_1
-      device: mi300_1
+      timeout_in_minutes: 55
      timeout_in_minutes: 45
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -79,28 +78,6 @@ steps:
  commands:
  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small-tp.txt
 - label: LM Eval PCP (4xB200)
  key: lm-eval-pcp-4xb200
  timeout_in_minutes: 360
  device: b200-k8s
  num_devices: 4
  optional: true
  source_file_dependencies:
  - csrc/
  - tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml
  - tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-EP.yaml
  - tests/evals/gsm8k/configs/models-pcp.txt
  - vllm/model_executor/layers/quantization
  - vllm/config/parallel.py
  - vllm/distributed/parallel_state.py
  - vllm/model_executor/layers/attention/mla_attention.py
  - vllm/model_executor/layers/attention/pcp.py
  - vllm/v1/worker/gpu/model_runner.py
  - vllm/v1/worker/gpu/pcp_manager.py
  autorun_on_main: true
  commands:
  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-pcp.txt
 - label: LM Eval Large Models EP (2xB200)
  key: lm-eval-large-models-ep-2xb200
  timeout_in_minutes: 60
@@ -126,7 +103,7 @@ steps:
  - vllm/transformers_utils/configs/qwen3_5_moe.py
  - vllm/model_executor/models/qwen3_next.py
  - vllm/model_executor/models/qwen3_next_mtp.py
-  - vllm/third_party/flash_linear_attention/ops/
+  - vllm/model_executor/layers/fla/ops/
  commands:
  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-qwen35-blackwell.txt
@@ -140,9 +117,8 @@ steps:
    - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-h200.txt
  mirror:
    amd:
      dind: false
      device: mi300_8
-      timeout_in_minutes: 40
+      timeout_in_minutes: 60
      depends_on:
      - image-build-amd
      commands:
@@ -337,7 +313,7 @@ steps:
 - label: LM Eval KV-Offload (2xH100)
  key: kv-offload-medium
-  timeout_in_minutes: 45
+  timeout_in_minutes: 30
  device: h100
  num_devices: 2
  source_file_dependencies:
@@ -347,7 +323,7 @@ steps:
  - vllm/v1/simple_kv_offload/
  - tests/evals/gsm8k/test_gsm8k_offloading.py
  commands:
-  - pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "qwen3.5-35b or deepseek-v2-lite"
+  - pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "qwen3.5-35b"
 - label: LM Eval KV-Offload (4xH100)
  key: kv-offload-large
@@ -14,10 +14,9 @@ steps:
  parallelism: 4
  mirror:
    amd:
-      dind: false
+      device: mi325_1
      device: mi300_1
      working_dir: "/vllm-workspace/tests"
-      timeout_in_minutes: 85
+      timeout_in_minutes: 65
      source_file_dependencies:
      - vllm/lora
      - tests/lora
@@ -23,15 +23,14 @@ steps:
    - pytest -v -s -m 'not slow_test' v1/spec_decode
  mirror:
    amd:
      dind: false
      device: mi300_1
-      timeout_in_minutes: 50
+      timeout_in_minutes: 75
      depends_on:
      - image-build-amd
 - label: V1 Sample + Logits
  key: v1-sample-logits
-  timeout_in_minutes: 83
+  timeout_in_minutes: 45
  device: h200_18gb
  source_file_dependencies:
    - vllm/config/
@@ -59,16 +58,13 @@ steps:
    - pytest -v -s v1/test_outputs.py
  mirror:
    amd:
-      dind: false
+      device: mi325_1
      device: mi300_1
      timeout_in_minutes: 70
      depends_on:
      - image-build-amd
 - label: V1 Core + KV + Metrics
  device: h200_35gb
  key: v1-core-kv-metrics
-  timeout_in_minutes: 80
+  timeout_in_minutes: 60
  source_file_dependencies:
    - vllm/config/
    - vllm/distributed/
@@ -92,7 +88,6 @@ steps:
    - tests/v1/kv_offload
    - tests/v1/simple_kv_offload
    - tests/v1/worker
    - tests/v1/streaming_input
    - tests/v1/kv_connector/unit
    - tests/v1/ec_connector/unit
    - tests/v1/metrics
@@ -106,7 +101,6 @@ steps:
    - pytest -v -s v1/kv_offload
    - pytest -v -s v1/simple_kv_offload
    - pytest -v -s v1/worker
    - pytest -v -s v1/streaming_input
    - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
    - pytest -v -s -m 'not cpu_test' v1/ec_connector/unit
    - pytest -v -s -m 'not cpu_test' v1/metrics
@@ -115,9 +109,8 @@ steps:
    - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
  mirror:
    amd:
-      dind: false
+      device: mi325_1
-      device: mi300_1
+      timeout_in_minutes: 75
      timeout_in_minutes: 65
      depends_on:
      - image-build-amd
@@ -148,8 +141,6 @@ steps:
    - pytest -v -s -m 'cpu_test' v1/core
    - pytest -v -s v1/structured_output
    - pytest -v -s v1/test_serial_utils.py
    - pytest -v -s v1/test_kv_cache_spec_registry.py
    - pytest -v -s v1/cudagraph/test_cudagraph_manager.py
    - pytest -v -s -m 'cpu_test' v1/kv_connector/unit
    - pytest -v -s -m 'cpu_test' v1/metrics
@@ -213,7 +204,7 @@ steps:
  - vllm/multimodal
  - examples/
  commands:
-    - pip install --no-deps tensorizer # for tensorizer test
+    - pip install tensorizer # for tensorizer test
     # for basic
    - python3 basic/offline_inference/chat.py
    - python3 basic/offline_inference/generate.py --model facebook/opt-125m
@@ -237,9 +228,7 @@ steps:
    - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
  mirror:
    amd:
-      dind: false
+      device: mi325_1
      device: mi300_1
      timeout_in_minutes: 75
      source_file_dependencies:
      - vllm/entrypoints
      - vllm/multimodal
@@ -266,7 +255,6 @@ steps:
  - vllm/utils/
  - vllm/v1/
  - tests/v1/tracing
  - tests/tracing/
  commands:
  - "pip install \
      'opentelemetry-sdk>=1.26.0' \
@@ -274,14 +262,12 @@ steps:
      'opentelemetry-exporter-otlp>=1.26.0' \
      'opentelemetry-semantic-conventions-ai>=0.4.1'"
  - pytest -v -s v1/tracing
  - pytest -v -s tracing
  mirror:
    amd:
-      dind: false
+      device: mi325_2
      device: mi300_2
      timeout_in_minutes: 30
      depends_on:
      - image-build-amd
      optional: true
 - label: Python-only Installation
  key: python-only-installation
@@ -296,9 +282,8 @@ steps:
  - bash standalone_tests/python_only_compile.sh
  mirror:
    amd:
-      device: mi250_1
+      device: mi325_1
-      timeout_in_minutes: 55
+      timeout_in_minutes: 45
      soft_fail: true
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -398,7 +383,7 @@ steps:
 - label: Batch Invariance (A100)
  key: batch-invariance-a100
-  timeout_in_minutes: 60
+  timeout_in_minutes: 40
  device: a100
  source_file_dependencies:
    - vllm/v1/attention
@@ -408,11 +393,11 @@ steps:
    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
    - pip install pytest-timeout pytest-forked
    - pytest -v -s v1/determinism/test_batch_invariance.py
-    - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA
+    - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]
 - label: Batch Invariance (H100)
  key: batch-invariance-h100
-  timeout_in_minutes: 60
+  timeout_in_minutes: 40
  device: h100
  source_file_dependencies:
    - vllm/v1/attention
@@ -423,12 +408,12 @@ steps:
    - pip install pytest-timeout pytest-forked
    - pytest -v -s v1/determinism/test_batch_invariance.py
    - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
-    - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA
+    - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]
-    - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k FLASH_ATTN
+    - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN]
 - label: Batch Invariance (B200)
  key: batch-invariance-b200
-  timeout_in_minutes: 45
+  timeout_in_minutes: 35
  device: b200-k8s
  source_file_dependencies:
    - vllm/v1/attention
@@ -439,13 +424,10 @@ steps:
    - pip install pytest-timeout pytest-forked
    - pytest -v -s v1/determinism/test_batch_invariance.py
    - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
-    - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA
+    - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]
-    - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k FLASH_ATTN
+    - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN]
    - pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py
    - pytest -v -s v1/determinism/test_nvfp4_batch_invariant_scaled_mm.py
    - pytest -v -s v1/determinism/test_matmul_batch_invariant.py
    - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py
    - pytest -v -s v1/determinism/test_online_batch_invariance.py
 - label: Acceptance Length Test (Large Models) # optional
  device: h200_35gb
@@ -3,16 +3,13 @@ depends_on:
  - image-build
 steps:
 - label: Model Executor
  device: h200_35gb
  key: model-executor
-  timeout_in_minutes: 60
+  timeout_in_minutes: 45
  source_file_dependencies:
  - vllm/engine/arg_utils.py
  - vllm/config/model.py
  - vllm/model_executor
  - vllm/model_executor/warmup
  - tests/model_executor
  - tests/model_executor/test_jit_warmup.py
  - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
  commands:
    - apt-get update && apt-get install -y curl libsodium23
@@ -28,18 +25,14 @@ steps:
    - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py --timeout=900 --timeout-method=thread
  mirror:
    amd:
      dind: false
      device: mi300_1
      timeout_in_minutes: 60
      depends_on:
      - image-build-amd
      source_file_dependencies:
      - vllm/engine/arg_utils.py
      - vllm/config/model.py
      - vllm/model_executor
      - vllm/model_executor/warmup
      - tests/model_executor
      - tests/model_executor/test_jit_warmup.py
      - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
      - vllm/_aiter_ops.py
      - vllm/platforms/rocm.py
@@ -41,7 +41,7 @@ steps:
  commands:
    - set -x
    - export VLLM_USE_V2_MODEL_RUNNER=1
-    - pip install --no-deps tensorizer # for tensorizer test
+    - pip install tensorizer # for tensorizer test
    - python3 basic/offline_inference/chat.py # for basic
    - python3 basic/offline_inference/generate.py --model facebook/opt-125m
    #- python3 basic/offline_inference/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10  # TODO
@@ -42,39 +42,10 @@ steps:
    - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py
  mirror:
    amd:
-      dind: false
+      device: mi325_1
      device: mi300_1
      timeout_in_minutes: 50
      depends_on:
      - image-build-amd
 - label: Inkling Unit Tests (B200)
  key: inkling-unit-tests-b200
  timeout_in_minutes: 40
  device: b200-k8s
  source_file_dependencies:
  - vllm/models/inkling/
  - vllm/cute_utils/
  - cmake/external_projects/tml_fa4.cmake
  - tests/models/inkling/
  commands:
    # FA4 kernel tests require SM100; the suite skips them elsewhere.
    - pytest -v -s models/inkling
 - label: Kimi K3 Unit Tests (B200)
  key: kimi-k3-unit-tests-b200
  timeout_in_minutes: 40
  device: b200-k8s
  source_file_dependencies:
  - vllm/models/kimi_k3/
  - csrc/libtorch_stable/kimi_k3/
  - tests/models/kimi_k3/
  - tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py
  - tests/kernels/test_bf16_skinny_gemm.py
  commands:
    # The native NVIDIA Kimi K3 kernels require the SM100 family.
    - pytest -v -s models/kimi_k3 kernels/attention/test_kimi_k3_mla_fused_epilogue.py kernels/test_bf16_skinny_gemm.py
 - label: Basic Models Test (Other CPU) # 5min
  key: basic-models-test-other-cpu
  depends_on:
@@ -84,8 +55,7 @@ steps:
  - vllm/
  - tests/models/test_utils.py
  - tests/models/test_vision.py
  - tests/models/test_adapters.py
  - tests/models/transformers/fusers/
  device: cpu-small
  commands:
-    - pytest -v -s models/test_utils.py models/test_vision.py models/test_adapters.py models/transformers/fusers/
+    - pytest -v -s models/test_utils.py models/test_vision.py models/transformers/fusers/
@@ -15,14 +15,11 @@ steps:
    - pytest -v -s models/language -m 'core_model and (not slow_test)'
  mirror:
    amd:
      dind: false
      device: mi300_1
      timeout_in_minutes: 45
      depends_on:
      - image-build-amd
 - label: Language Models Tests (Extra Standard) %N
  device: h200_35gb
  key: language-models-tests-extra-standard
  timeout_in_minutes: 40
  source_file_dependencies:
@@ -38,9 +35,7 @@ steps:
  parallelism: 2
  mirror:
    amd:
      dind: false
      device: mi300_1
      timeout_in_minutes: 40
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -54,8 +49,8 @@ steps:
      - tests/models/language/pooling/test_classification.py
      - vllm/_aiter_ops.py
      - vllm/platforms/rocm.py
 - label: Language Models Tests (Hybrid) %N
  device: h200_35gb
  key: language-models-tests-hybrid
  timeout_in_minutes: 65
  source_file_dependencies:
@@ -63,16 +58,16 @@ steps:
  - tests/models/language/generation
  commands:
    # Install fast path packages for testing against transformers
    # Note: also needed to run plamo2 model in vLLM
    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
    - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
-    # Shard the hybrid language model tests that are numerically stable on Hopper.
+    # Shard hybrid language model tests
-    - pytest -v -s models/language/generation -m hybrid_model -k 'not granite-4.0-tiny-preview' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
+    - pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
  parallelism: 2
  mirror:
    amd:
-      dind: false
+      device: mi325_1
-      device: mi300_1
+      timeout_in_minutes: 70
      timeout_in_minutes: 60
      depends_on:
      - image-build-amd
      commands:
@@ -80,20 +75,6 @@ steps:
      - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
      - pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
 # Granite 4 hybrid generation is sensitive to hardware-specific Triton SSD
 # autotuning (https://github.com/vllm-project/vllm/issues/25194). Keep this one
 # correctness test on L4 until its H200 output matches the Transformers reference.
 - label: Language Models Tests (Granite L4 Compatibility)
  key: language-models-tests-granite-l4-compatibility
  timeout_in_minutes: 65
  source_file_dependencies:
  - vllm/
  - tests/models/language/generation
  commands:
    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
    - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
    - pytest -v -s models/language/generation -m hybrid_model -k 'granite-4.0-tiny-preview'
 - label: Language Models Test (Extended Generation) # 80min
  device: h200_35gb
  key: language-models-test-extended-generation
@@ -104,6 +85,7 @@ steps:
  - tests/models/language/generation
  commands:
    # Install fast path packages for testing against transformers
    # Note: also needed to run plamo2 model in vLLM
    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
    - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
    - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'
@@ -119,10 +101,10 @@ steps:
  commands:
    - pytest -v -s models/language/generation_ppl_test
- label: Language Models Test (Extended Pooling)
+- label: Language Models Test (Extended Pooling)  # 36min
  device: h200_35gb
  key: language-models-test-extended-pooling
-  timeout_in_minutes: 120
+  timeout_in_minutes: 70
  optional: true
  source_file_dependencies:
  - vllm/
@@ -131,15 +113,14 @@ steps:
    - pytest -v -s models/language/pooling -m 'not core_model'
  mirror:
    amd:
-      dind: false
+      device: mi325_1
-      device: mi300_1
+      timeout_in_minutes: 100
      timeout_in_minutes: 95
      depends_on:
      - image-build-amd
 - label: Language Models Test (MTEB)
  key: language-models-test-mteb
-  timeout_in_minutes: 68
+  timeout_in_minutes: 45
  device: h200_18gb
  optional: true
  source_file_dependencies:
@@ -4,7 +4,7 @@ depends_on:
 steps:
 - label: "Multi-Modal Models (Standard) 1: qwen2"
  key: multi-modal-models-standard-1-qwen2
-  timeout_in_minutes: 68
+  timeout_in_minutes: 45
  device: h200_18gb
  source_file_dependencies:
  - vllm/
@@ -14,15 +14,13 @@ steps:
    - pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model
  mirror:
    amd:
-      dind: false
+      device: mi325_1
      device: mi300_1
      timeout_in_minutes: 65
      depends_on:
      - image-build-amd
 - label: "Multi-Modal Models (Standard) 2: qwen3 + gemma"
  key: multi-modal-models-standard-2-qwen3-gemma
-  timeout_in_minutes: 75
+  timeout_in_minutes: 50
  device: h200_18gb
  source_file_dependencies:
  - vllm/
@@ -33,9 +31,7 @@ steps:
    - pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model
  mirror:
    amd:
-      dind: false
+      device: mi325_1
      device: mi300_1
      timeout_in_minutes: 55
      depends_on:
      - image-build-amd
@@ -51,15 +47,14 @@ steps:
    - pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model
  mirror:
    amd:
-      device: mi250_1
+      device: mi325_1
      timeout_in_minutes: 55
      depends_on:
      - image-build-amd
 - label: "Multi-Modal Models (Standard) 4: other + whisper"
  device: h200_35gb
  key: multi-modal-models-standard-4-other-whisper
-  timeout_in_minutes: 75
+  timeout_in_minutes: 50
  source_file_dependencies:
  - vllm/
  - tests/models/multimodal
@@ -70,9 +65,7 @@ steps:
    - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model  # Otherwise, mp_method="spawn" doesn't work
  mirror:
    amd:
-      dind: false
+      device: mi325_1
      device: mi300_1
      timeout_in_minutes: 50
      depends_on:
      - image-build-amd
@@ -92,7 +85,7 @@ steps:
 - label: Multi-Modal Processor # 44min
  key: multi-modal-processor
-  timeout_in_minutes: 98
+  timeout_in_minutes: 65
  device: h200_18gb
  source_file_dependencies:
  - vllm/
@@ -114,9 +107,7 @@ steps:
  - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-mm-small.txt --tp-size=1
  mirror:
    amd:
      dind: false
      device: mi300_1
      timeout_in_minutes: 35
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -127,7 +118,6 @@ steps:
      - vllm/model_executor/model_loader/
 - label: Multi-Modal Models (Extended Generation 1)
  device: h200_35gb
  key: multi-modal-models-extended-generation-1
  optional: true
  source_file_dependencies:
@@ -139,9 +129,7 @@ steps:
    - pytest -v -s models/multimodal/test_mapping.py
  mirror:
    amd:
-      dind: false
+      device: mi325_1
      device: mi300_1
      timeout_in_minutes: 90
      depends_on:
      - image-build-amd
@@ -176,9 +164,8 @@ steps:
    - pytest -v -s models/multimodal/pooling -m 'not core_model'
  mirror:
    amd:
-      dind: false
+      device: mi325_1
-      device: mi300_1
+      timeout_in_minutes: 75
      timeout_in_minutes: 60
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -5,7 +5,7 @@ steps:
 - label: PyTorch Compilation Unit Tests
  device: h200_35gb
  key: pytorch-compilation-unit-tests
-  timeout_in_minutes: 150
+  timeout_in_minutes: 90
  source_file_dependencies:
    - vllm/__init__.py
    - vllm/_aiter_ops.py
@@ -107,11 +107,16 @@ steps:
    - tests/compile/passes
  commands:
  - pytest -s -v compile/passes --ignore compile/passes/distributed
  mirror:
    amd:
      device: mi300_1
      timeout_in_minutes: 65
      depends_on:
      - image-build-amd
 - label: PyTorch Fullgraph Smoke Test
  device: h200_35gb
  key: pytorch-fullgraph-smoke-test
-  timeout_in_minutes: 90
+  timeout_in_minutes: 60
  source_file_dependencies:
  - vllm/__init__.py
  - vllm/_aiter_ops.py
@@ -143,42 +148,7 @@ steps:
  # as it is a heavy test that is covered in other steps.
  # Use `find` to launch multiple instances of pytest so that
  # they do not suffer from https://github.com/vllm-project/vllm/issues/28965
-  - "find compile/fullgraph/ -name 'test_*.py' -not -name 'test_full_cudagraph.py' -not -name 'test_full_graph.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
+  - "find compile/fullgraph/ -name 'test_*.py' -not -name 'test_full_graph.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
 # Hopper-only DeepSeek-V2-Lite cases in this file require two 29.3-GiB model
 # instances and cannot fit a 35GB MIG slice. L4 retains the original coverage:
 # those SM90 cases skip while the architecture-compatible cases still run.
 - label: PyTorch Fullgraph CUDAGraph (L4 Compatibility)
  key: pytorch-fullgraph-cudagraph-l4-compatibility
  timeout_in_minutes: 60
  source_file_dependencies:
  - vllm/__init__.py
  - vllm/_aiter_ops.py
  - vllm/_custom_ops.py
  - vllm/compilation/
  - vllm/config/
  - vllm/distributed/
  - vllm/engine/
  - vllm/env_override.py
  - vllm/envs.py
  - vllm/forward_context.py
  - vllm/inputs/
  - vllm/ir/
  - vllm/kernels/
  - vllm/logger.py
  - vllm/model_executor/
  - vllm/multimodal/
  - vllm/platforms/
  - vllm/plugins/
  - vllm/sampling_params.py
  - vllm/sequence.py
  - vllm/transformers_utils/
  - vllm/triton_utils/
  - vllm/utils/
  - vllm/v1/
  - tests/compile
  commands:
  - pytest -s -v compile/fullgraph/test_full_cudagraph.py
 - label: PyTorch Fullgraph
  key: pytorch-fullgraph
@@ -227,9 +197,7 @@ steps:
  - bash standalone_tests/pytorch_nightly_dependency.sh
  mirror:
    amd:
      dind: false
      device: mi300_1
      timeout_in_minutes: 30
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -3,11 +3,8 @@ depends_on:
  - image-build
 steps:
 - label: Quantization
  device: h200_35gb
  key: quantization
-  timeout_in_minutes: 75
+  timeout_in_minutes: 60
  env:
    VLLM_USE_V2_MODEL_RUNNER: "0"
  source_file_dependencies:
  - csrc/
  - vllm/model_executor/layers/quantization
@@ -22,12 +19,9 @@ steps:
  # TODO(jerryzh168): resolve the above comment
  - uv pip install --system torchao==0.17.0 --index-url https://download.pytorch.org/whl/cu130
  - uv pip install --system conch-triton-kernels
-  # The SM90-only checkpoint currently contains a removed weight_chan_scale
+  - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py
  # parameter. It was not exercised by the previous L4 job.
  - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py -k 'not test_compressed_tensors_w4a8_fp8'
 - label: Quantized Fusions
  device: h200_35gb
  key: quantized-fusions
  timeout_in_minutes: 20
  source_file_dependencies:
@@ -58,11 +52,8 @@ steps:
    - pytest -s -v tests/quantization/test_blackwell_moe.py
 - label: Quantized Models Test
  device: h200_35gb
  key: quantized-models-test
-  timeout_in_minutes: 65
+  timeout_in_minutes: 50
  env:
    VLLM_USE_V2_MODEL_RUNNER: "0"
  source_file_dependencies:
  - vllm/model_executor/layers/quantization
  - tests/models/quantization
@@ -26,7 +26,7 @@ steps:
  - export VLLM_USE_RUST_FRONTEND=1
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
  - pytest -v -s benchmarks/test_serve_cli.py -k "not insecure and not (test_bench_serve and not test_bench_serve_chat)"
-  - pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py -k "not test_invalid_json_schema and not test_invalid_regex and not test_kv_transfer_prompt_token_ids_round_trip and not test_kv_transfer_prompt_token_ids_streaming"
+  - pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py -k "not test_invalid_json_schema and not test_invalid_regex"
  - pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not multiple"
  # - pytest -v -s entrypoints/openai/completion/test_prompt_validation.py -k "not prompt_embeds"
@@ -81,7 +81,6 @@ steps:
  - pytest -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
 - label: Rust Frontend Tool Use
  device: h200_35gb
  timeout_in_minutes: 25
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -19,18 +19,8 @@ steps:
    - VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers
  mirror:
    amd:
-      device: mi250_1
+      device: mi325_1
      timeout_in_minutes: 40
      depends_on:
      - image-build-amd
      source_file_dependencies:
      - vllm/model_executor/layers
      - vllm/sampling_metadata.py
      - vllm/v1/sample/
      - vllm/entrypoints/generate/beam_search/
      - tests/samplers
      - tests/conftest.py
      - vllm/_aiter_ops.py
      - vllm/platforms/rocm.py
      commands:
      - pytest -v -s samplers
@@ -14,9 +14,8 @@ steps:
    - pytest -v -s v1/e2e/spec_decode -k "eagle_correctness"
  mirror:
    amd:
-      dind: false
+      device: mi325_1
-      device: mi300_1
+      timeout_in_minutes: 60
      timeout_in_minutes: 55
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -54,9 +53,8 @@ steps:
    - pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
  mirror:
    amd:
-      dind: false
+      device: mi325_1
-      device: mi300_1
+      timeout_in_minutes: 65
      timeout_in_minutes: 75
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -90,15 +88,14 @@ steps:
    - vllm/v1/spec_decode/
    - vllm/v1/worker/gpu/spec_decode/
    - tests/v1/e2e/spec_decode/
    - tests/spec_decode/
  commands:
    - pytest -v -s v1/e2e/spec_decode -k "ngram or suffix"
    - python3 spec_decode/test_custom_proposer.py
  mirror:
    amd:
-      dind: false
+      device: mi325_1
-      device: mi300_1
+      timeout_in_minutes: 55
-      timeout_in_minutes: 35
+      # TODO(akaratza): Test after Torch >= 2.12 bump
      soft_fail: true
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -122,8 +119,7 @@ steps:
    - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
  mirror:
    amd:
-      dind: false
+      device: mi325_1
      device: mi300_1
      timeout_in_minutes: 55
      depends_on:
      - image-build-amd
@@ -174,31 +170,3 @@ steps:
    - tests/v1/e2e/spec_decode/
  commands:
    - pytest -v -s v1/e2e/spec_decode -k "qwen3_5-hybrid"
 - label: Spec Decode DeepSeek MTP Parallel Load (B200)
  key: spec-decode-deepseek-mtp-parallel-load-b200
  timeout_in_minutes: 30
  device: b200-k8s
  optional: true
  num_devices: 2
  source_file_dependencies:
    - vllm/v1/spec_decode/llm_base_proposer.py
    - vllm/v1/spec_decode/eagle.py
    - vllm/v1/worker/gpu/spec_decode/eagle/
    - vllm/model_executor/models/deepseek_mtp.py
    - vllm/model_executor/models/deepseek_v2.py
    - tests/v1/e2e/spec_decode/test_mtp_parallel_load.py
  commands:
    - pytest -v -s v1/e2e/spec_decode/test_mtp_parallel_load.py
 - label: Spec Decode Acceptance Rates Nightly
  key: spec-decode-acceptance-rates-nightly
  timeout_in_minutes: 60
  device: h200_35gb
  optional: true
  source_file_dependencies:
    - vllm/v1/spec_decode/
    - vllm/v1/worker/gpu/spec_decode/
    - tests/v1/e2e/spec_decode/
  commands:
    - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "acceptance_rates"
@@ -1,14 +0,0 @@
 group: Torch ABI
 depends_on:
  - image-build
 steps:
 - label: Torch Stable ABI Audit
  key: torch-stable-abi-audit
  timeout_in_minutes: 5
  source_file_dependencies:
  - .buildkite/check-torch-abi.py
  - csrc/
  - cmake/
  - setup.py
  commands:
  - python3 /vllm-workspace/.buildkite/check-torch-abi.py
@@ -15,9 +15,7 @@ steps:
    - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models.txt
  mirror:
    amd:
      dind: false
      device: mi300_2
      timeout_in_minutes: 35
      depends_on:
      - image-build-amd
      commands:
@@ -47,7 +47,6 @@
 # Rust Frontend
 /rust/ @BugenZhao @njhill
 /rust/src/bench @esmeetu
 /build_rust.sh @BugenZhao @njhill
 /rust-toolchain.toml @BugenZhao @njhill
 /.buildkite/test_areas/rust* @BugenZhao @njhill
@@ -173,7 +172,7 @@ mkdocs.yaml @hmellor
 # Kernels
 /vllm/v1/attention/ops/chunked_prefill_paged_decode.py @tdoublep
 /vllm/v1/attention/ops/triton_unified_attention.py @tdoublep
-/vllm/third_party/flash_linear_attention @ZJY0516 @vadiklyutiy
+/vllm/model_executor/layers/fla @ZJY0516 @vadiklyutiy
 # ROCm related: specify owner with write access to notify AMD folks for careful code review
 /vllm/**/*rocm* @tjtanaa @dllehr-amd
@@ -19,7 +19,6 @@ pull_request_rules:
  description: Comment on PR when pre-commit check fails
  conditions:
    - check-failure=pre-commit
    - -check-cancelled=pre-commit
    - -closed
    - -draft
    - or:
@@ -182,18 +181,6 @@ pull_request_rules:
      add:
        - performance
 - name: label-quantization
  description: Automatically apply quantization label
  conditions:
    - label != stale
    - or:
      - files~=^vllm/model_executor/layers/quantization/
      - title~=(?i)quant
  actions:
    label:
      add:
        - quantization
 - name: label-qwen
  description: Automatically apply qwen label
  conditions:
@@ -233,31 +220,6 @@ pull_request_rules:
      add:
        - gpt-oss
 - name: label-kimi
  description: Automatically apply kimi label
  conditions:
    - label != stale
    - or:
      - files~=(?i)kimi
      - files~=(?i)moonshot
      - title~=(?i)(?:kimi|moonshot)
  actions:
    label:
      add:
        - kimi
 - name: label-k3
  description: Automatically apply k3 label (launch triage; retire after ramp-down)
  conditions:
    - label != stale
    - or:
      - files~=(?i)kimi[-_]?k3
      - title~=(?i)(?:kimi[-\s]?k3|\bk3\b)
  actions:
    label:
      add:
        - k3
 - name: label-nvidia
  description: Automatically apply nvidia label
  conditions:
@@ -130,66 +130,6 @@ jobs:
                  },
                ],
              },
              kimi: {
                keywords: [
                  { term: "Kimi",     searchIn: "both" },
                  { term: "Moonshot", searchIn: "both" },
                ],
                substrings: [
                  { term: "moonshotai/", searchIn: "both" },
                  { term: "kimi",        searchIn: "title" },
                ],
              },
              k3: {
                keywords: [
                  { term: "Kimi K3", searchIn: "both" },
                  { term: "K3",      searchIn: "title" },
                ],
                substrings: [
                  { term: "moonshotai/kimi-k3", searchIn: "both" },
                ],
              },
              quantization: {
                keywords: [
                  {
                    term: "quantization",
                    searchIn: "both"
                  },
                  {
                    term: "quantized",
                    searchIn: "both"
                  },
                ],
              },
              "intel-gpu": {
                // Keyword search - matches whole words only (with word boundaries)
                keywords: [
                  {
                    term: "B50",
                    searchIn: "both"
                  },
                  {
                    term: "B60",
                    searchIn: "both"
                  },
                  {
                    term: "B70",
                    searchIn: "both"
                  },
                  {
                    term: "intel gpu",
                    searchIn: "both"
                  },
                  {
                    term: "Arc GPU",
                    searchIn: "both"
                  },
                  {
                    term: "BMG",
                    searchIn: "both"
                  },
                ],
              },
              // Add more label configurations here as needed
              // example: {
              //   keywords: [...],
@@ -383,7 +323,7 @@ jobs:
            // {users} will be replaced with @mentions
            const ccConfig = {
              rocm: {
-                users: ['hongxiayang', 'tjtanaa', 'vllmellm', 'giuseppegrossi'],
+                users: ['hongxiayang', 'tjtanaa', 'vllmellm'],
                message: 'CC {users} for ROCm-related issue',
              },
              mistral: {
@@ -80,9 +80,9 @@ jobs:
                  '',
                  '\u{1f4ac} Join our developer Slack at https://slack.vllm.ai to discuss your PR in `#pr-reviews`, coordinate on features in `#feat-` channels, or join special interest groups in `#sig-` channels.',
                  '',
-                  'PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment `/ci run` whenever CI signals are needed.',
+                  'PRs do not trigger a full CI run by default. Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.',
                  '',
-                  'Once the PR is approved or has the `ready` label, the PR author can also use `/ci run` or `/ci retry`. New commits do not start CI automatically.',
+                  'To run CI, PR reviewers can either: Add `ready` label to the PR or enable auto-merge.',
                  '',
                  'If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.',
                  '',
@@ -41,7 +41,7 @@ jobs:
          if (hasReadyLabel || hasVerifiedLabel || mergedCount >= 4) {
            core.info(`Check passed: verified label=${hasVerifiedLabel}, ready label=${hasReadyLabel}, 4+ merged PRs=${mergedCount >= 4}`);
          } else {
-            core.setFailed(`PR must have the 'verified', 'ready', or 'ready-run-all-tests' label to run pre-commit, or the author must have at least 4 merged PRs (found ${mergedCount}).`);
+            core.setFailed(`PR must have the 'verified', 'ready', or 'ready-run-all-tests' label (the ready labels also trigger tests) or the author must have at least 4 merged PRs (found ${mergedCount}).`);
          }
  pre-commit:
@@ -1,40 +0,0 @@
 name: Run CI from PR comment
 on:
  issue_comment:
    types: [created]
 concurrency:
  group: run-ci-comment-${{ github.event.issue.number }}
  cancel-in-progress: false
 permissions:
  contents: read
  issues: write
  pull-requests: write
 jobs:
  run-ci-command:
    if: >-
      github.event.issue.pull_request &&
      (github.event.comment.body == '/ci run' ||
      github.event.comment.body == '/ci retry')
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
        with:
          ref: ${{ github.event.repository.default_branch }}
          persist-credentials: false
      - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0
        with:
          python-version: "3.12"
      - name: Authorize and run CI command
        run: >-
          uv run --no-project --python 3.12
          .github/workflows/scripts/run_ci_command.py
        env:
          BUILDKITE_API_TOKEN: ${{ secrets.BUILDKITE_API_TOKEN }}
          BUILDKITE_ORGANIZATION: vllm
          BUILDKITE_PIPELINE: ci
          CI_TRUSTED_USERS: ${{ vars.CI_TRUSTED_USERS }}
          GH_TOKEN: ${{ github.token }}
@@ -1,581 +0,0 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import json
 import os
 import sys
 import urllib.error
 import urllib.parse
 import urllib.request
 from collections.abc import Mapping, Sequence
 from typing import Any
 COMMAND_RUN_CI = "/ci run"
 COMMAND_RETRY_FAILED = "/ci retry"
 READY_LABELS = {"ready", "ready-run-all-tests"}
 TRUSTED_PERMISSIONS = {"admin", "maintain", "write"}
 ACTIVE_BUILD_STATES = {
    "blocked",
    "creating",
    "scheduled",
    "running",
    "failing",
    "canceling",
    "waiting",
    "waiting_failed",
 }
 RETRY_STATES = "failed,timed_out,expired"
 class ApiError(RuntimeError):
    def __init__(self, status: int | None, message: str) -> None:
        super().__init__(message)
        self.status = status
 class HttpTransport:
    def request(
        self,
        url: str,
        *,
        body: Mapping[str, Any] | None = None,
        headers: Mapping[str, str] | None = None,
        method: str = "GET",
    ) -> Any:
        data = None if body is None else json.dumps(body).encode()
        request = urllib.request.Request(
            url,
            data=data,
            headers=dict(headers or {}),
            method=method,
        )
        try:
            with urllib.request.urlopen(request, timeout=30) as response:
                response_body = response.read().decode()
        except urllib.error.HTTPError as error:
            response_body = error.read().decode()
            message = self._error_message(response_body, error.reason)
            raise ApiError(
                error.code,
                f"API returned {error.code}: {message}",
            ) from error
        except urllib.error.URLError as error:
            raise ApiError(None, f"API request failed: {error.reason}") from error
        if not response_body:
            return None
        try:
            return json.loads(response_body)
        except json.JSONDecodeError as error:
            raise ApiError(None, "API returned a non-JSON response.") from error
    @staticmethod
    def _error_message(response_body: str, fallback: str) -> str:
        try:
            parsed = json.loads(response_body)
        except json.JSONDecodeError:
            return fallback
        return str(parsed.get("message", fallback))
 class GitHubClient:
    def __init__(
        self,
        token: str,
        repository: str,
        transport: HttpTransport | None = None,
    ) -> None:
        if not token:
            raise RuntimeError("GH_TOKEN is not set.")
        self.owner, self.repo = repository.split("/", maxsplit=1)
        self.transport = transport or HttpTransport()
        self.headers = {
            "Accept": "application/vnd.github+json",
            "Authorization": f"Bearer {token}",
            "Content-Type": "application/json",
            "User-Agent": "vllm-ci-command",
            "X-GitHub-Api-Version": "2022-11-28",
        }
    def _request(
        self,
        path: str,
        *,
        body: Mapping[str, Any] | None = None,
        method: str = "GET",
    ) -> Any:
        return self.transport.request(
            f"https://api.github.com{path}",
            body=body,
            headers=self.headers,
            method=method,
        )
    def _repo_path(self, suffix: str) -> str:
        owner = urllib.parse.quote(self.owner, safe="")
        repo = urllib.parse.quote(self.repo, safe="")
        return f"/repos/{owner}/{repo}{suffix}"
    def _paginate(self, path: str) -> list[dict[str, Any]]:
        results: list[dict[str, Any]] = []
        separator = "&" if "?" in path else "?"
        for page in range(1, 101):
            response = self._request(f"{path}{separator}per_page=100&page={page}")
            if not isinstance(response, list):
                raise ApiError(None, "GitHub API returned an invalid list response.")
            results.extend(response)
            if len(response) < 100:
                return results
        raise ApiError(None, "GitHub API pagination exceeded 10,000 results.")
    def get_pr(self, number: int) -> dict[str, Any]:
        return self._request(self._repo_path(f"/pulls/{number}"))
    def get_permission(self, actor: str) -> str:
        username = urllib.parse.quote(actor, safe="")
        try:
            response = self._request(
                self._repo_path(f"/collaborators/{username}/permission")
            )
        except ApiError as error:
            if error.status == 404:
                return "none"
            raise
        return str(response["permission"])
    def get_review_decision(self, number: int) -> str | None:
        query = """
          query($owner: String!, $repo: String!, $number: Int!) {
            repository(owner: $owner, name: $repo) {
              pullRequest(number: $number) {
                reviewDecision
              }
            }
          }
        """
        response = self._request(
            "/graphql",
            body={
                "query": query,
                "variables": {
                    "number": number,
                    "owner": self.owner,
                    "repo": self.repo,
                },
            },
            method="POST",
        )
        return response["data"]["repository"]["pullRequest"]["reviewDecision"]
    def list_reviews(self, number: int) -> list[dict[str, Any]]:
        return self._paginate(self._repo_path(f"/pulls/{number}/reviews"))
    def list_reactions(self, comment_id: int) -> list[dict[str, Any]]:
        return self._paginate(
            self._repo_path(f"/issues/comments/{comment_id}/reactions")
        )
    def add_reaction(self, comment_id: int, content: str) -> None:
        self._request(
            self._repo_path(f"/issues/comments/{comment_id}/reactions"),
            body={"content": content},
            method="POST",
        )
    def add_comment(self, issue_number: int, body: str) -> None:
        self._request(
            self._repo_path(f"/issues/{issue_number}/comments"),
            body={"body": body},
            method="POST",
        )
 class BuildkiteClient:
    def __init__(
        self,
        token: str,
        organization: str,
        pipeline: str,
        transport: HttpTransport | None = None,
    ) -> None:
        self.token = token
        self.transport = transport or HttpTransport()
        organization = urllib.parse.quote(organization, safe="")
        pipeline = urllib.parse.quote(pipeline, safe="")
        self.base_url = (
            "https://api.buildkite.com/v2/organizations/"
            f"{organization}/pipelines/{pipeline}/builds"
        )
    def _request(
        self,
        *,
        body: Mapping[str, Any] | None = None,
        method: str = "GET",
        path: str = "",
        query: Sequence[tuple[str, str]] = (),
    ) -> Any:
        if not self.token:
            raise RuntimeError("The BUILDKITE_API_TOKEN repository secret is not set.")
        url = f"{self.base_url}{path}"
        if query:
            url = f"{url}?{urllib.parse.urlencode(query)}"
        return self.transport.request(
            url,
            body=body,
            headers={
                "Authorization": f"Bearer {self.token}",
                "Content-Type": "application/json",
                "User-Agent": "vllm-ci-command",
            },
            method=method,
        )
    def list_builds(
        self,
        commit: str,
        *,
        metadata: tuple[str, str] | None = None,
    ) -> list[dict[str, Any]]:
        query = [
            ("commit", commit),
            ("exclude_jobs", "true"),
            ("exclude_pipeline", "true"),
            ("per_page", "100"),
        ]
        if metadata:
            key, value = metadata
            query.append((f"meta_data[{key}]", value))
        response = self._request(query=query)
        if not isinstance(response, list):
            raise ApiError(None, "Buildkite API returned an invalid build list.")
        return response
    def create_build(self, body: Mapping[str, Any]) -> dict[str, Any]:
        return self._request(body=body, method="POST")
    def retry_failed_jobs(
        self,
        build_number: int,
        states: str,
    ) -> dict[str, Any]:
        number = urllib.parse.quote(str(build_number), safe="")
        return self._request(
            body={"states": states},
            method="PUT",
            path=f"/{number}/retry_failed_jobs",
        )
 def parse_command(body: str) -> str | None:
    if body in {COMMAND_RUN_CI, COMMAND_RETRY_FAILED}:
        return body
    return None
 def parse_trusted_users(value: str = "") -> set[str]:
    return {
        user.casefold() for item in value.split(",") for user in item.split() if user
    }
 def has_ready_label(pr: Mapping[str, Any]) -> bool:
    return any(label["name"] in READY_LABELS for label in pr["labels"])
 def is_trusted_permission(permission: str) -> bool:
    return permission in TRUSTED_PERMISSIONS
 def authorize(
    *,
    actor: str,
    permission: str,
    pr: Mapping[str, Any],
    trusted_approval: bool = False,
    trusted_users: set[str] | None = None,
 ) -> tuple[bool, str]:
    trusted_users = trusted_users or set()
    if is_trusted_permission(permission):
        return True, f"repository {permission} permission"
    if actor.casefold() in trusted_users:
        return True, "configured trusted contributor"
    if actor.casefold() != pr["user"]["login"].casefold():
        return (
            False,
            "Only reviewers with write access can run CI before it is "
            "delegated to the PR author.",
        )
    if pr["draft"]:
        return False, "PR authors cannot run CI while the PR is a draft."
    if has_ready_label(pr):
        return True, "ready label"
    if trusted_approval:
        return True, "approval from a trusted reviewer"
    return (
        False,
        "A reviewer with write access must run `/ci run`, approve the PR, "
        "or add the `ready` label first.",
    )
 def has_trusted_approval(
    github: GitHubClient,
    number: int,
    trusted_users: set[str],
 ) -> bool:
    if github.get_review_decision(number) != "APPROVED":
        return False
    latest_review_states: dict[str, tuple[str, str]] = {}
    for review in github.list_reviews(number):
        user = review.get("user") or {}
        login = user.get("login")
        state = review.get("state")
        if login and state in {"APPROVED", "CHANGES_REQUESTED", "DISMISSED"}:
            latest_review_states[login.casefold()] = (login, state)
    for login, state in latest_review_states.values():
        if state != "APPROVED":
            continue
        if login.casefold() in trusted_users:
            return True
        if is_trusted_permission(github.get_permission(login)):
            return True
    return False
 def is_build_for_pr(build: Mapping[str, Any], pr_number: int) -> bool:
    pull_request = build.get("pull_request")
    if isinstance(pull_request, Mapping):
        build_pr_number = pull_request.get("id", pull_request.get("number"))
        if build_pr_number is not None:
            return str(build_pr_number) == str(pr_number)
    metadata = build.get("meta_data") or {}
    return str(metadata.get("github-pr-number")) == str(pr_number)
 def is_active_build(build: Mapping[str, Any]) -> bool:
    return bool(build.get("blocked")) or build.get("state") in ACTIVE_BUILD_STATES
 def select_latest_build(
    builds: Sequence[dict[str, Any]],
    pr_number: int,
 ) -> dict[str, Any] | None:
    matching = [build for build in builds if is_build_for_pr(build, pr_number)]
    return max(matching, key=lambda build: build.get("created_at", ""), default=None)
 def create_build_payload(
    *,
    actor: str,
    comment_id: int,
    pr: Mapping[str, Any],
 ) -> dict[str, Any]:
    return {
        "commit": pr["head"]["sha"],
        "branch": pr["head"]["ref"],
        "message": f"PR #{pr['number']} {COMMAND_RUN_CI} by @{actor}",
        "pull_request_id": pr["number"],
        "pull_request_base_branch": pr["base"]["ref"],
        "pull_request_repository": pr["head"]["repo"]["clone_url"],
        "pull_request_labels": [label["name"] for label in pr["labels"]],
        "ignore_pipeline_branch_filters": True,
        "env": {
            "VLLM_CI_GITHUB_COMMENT_ID": str(comment_id),
            "VLLM_CI_TRIGGERED_BY": actor,
        },
        "meta_data": {
            "github-comment-id": str(comment_id),
            "github-pr-number": str(pr["number"]),
            "github-triggered-by": actor,
        },
    }
 def add_reaction_safely(
    github: GitHubClient,
    comment_id: int,
    content: str,
 ) -> None:
    try:
        github.add_reaction(comment_id, content)
    except Exception as error:
        print(f"Could not add {content} reaction: {error}", file=sys.stderr)
 def is_already_handled(github: GitHubClient, comment_id: int) -> bool:
    return any(
        reaction.get("content") in {"rocket", "-1"}
        and (reaction.get("user") or {}).get("login") == "github-actions[bot]"
        for reaction in github.list_reactions(comment_id)
    )
 def handle_run_ci(
    *,
    actor: str,
    buildkite: BuildkiteClient,
    comment_id: int,
    github: GitHubClient,
    pr: Mapping[str, Any],
 ) -> str:
    duplicate_builds = buildkite.list_builds(
        pr["head"]["sha"],
        metadata=("github-comment-id", str(comment_id)),
    )
    duplicate = select_latest_build(duplicate_builds, pr["number"])
    if duplicate:
        return f"CI was already requested by this comment: {duplicate['web_url']}"
    current_builds = buildkite.list_builds(pr["head"]["sha"])
    active_build = next(
        (
            build
            for build in current_builds
            if is_build_for_pr(build, pr["number"]) and is_active_build(build)
        ),
        None,
    )
    if active_build:
        return f"CI is already running for this commit: {active_build['web_url']}"
    current_pr = github.get_pr(pr["number"])
    if current_pr["state"] != "open" or current_pr["head"]["sha"] != pr["head"]["sha"]:
        return (
            "The PR head changed while processing the command. Comment `/ci run` again."
        )
    build = buildkite.create_build(
        create_build_payload(
            actor=actor,
            comment_id=comment_id,
            pr=current_pr,
        )
    )
    return (
        f"Triggered [Buildkite CI #{build['number']}]({build['web_url']}) "
        f"for commit `{current_pr['head']['sha'][:12]}`."
    )
 def handle_retry_failed(
    *,
    buildkite: BuildkiteClient,
    pr: Mapping[str, Any],
 ) -> str:
    builds = buildkite.list_builds(pr["head"]["sha"])
    build = select_latest_build(builds, pr["number"])
    if not build:
        return "No CI build exists for the current PR commit. Use `/ci run` first."
    if not build.get("finished_at") or is_active_build(build):
        return f"CI is still running for this commit: {build['web_url']}"
    retried = buildkite.retry_failed_jobs(build["number"], RETRY_STATES)
    if retried["retried_jobs_count"] == 0:
        return (
            f"No failed, timed-out, or expired jobs need retrying: {build['web_url']}"
        )
    return (
        f"Queued {retried['retried_jobs_count']} failed job(s) for retry in "
        f"[Buildkite CI #{build['number']}]({build['web_url']})."
    )
 def run(
    event: Mapping[str, Any],
    github: GitHubClient,
    buildkite: BuildkiteClient,
    trusted_users_value: str = "",
 ) -> None:
    command = parse_command(event["comment"]["body"])
    if not command or "pull_request" not in event["issue"]:
        return
    issue_number = event["issue"]["number"]
    comment_id = event["comment"]["id"]
    actor = event["comment"]["user"]["login"]
    if is_already_handled(github, comment_id):
        print(f"Comment {comment_id} was already handled.")
        return
    add_reaction_safely(github, comment_id, "eyes")
    try:
        pr = github.get_pr(issue_number)
        permission = github.get_permission(actor)
        if pr["state"] != "open":
            github.add_comment(issue_number, "CI commands require an open PR.")
            return
        trusted_users = parse_trusted_users(trusted_users_value)
        should_check_approval = (
            not is_trusted_permission(permission)
            and actor.casefold() not in trusted_users
            and actor.casefold() == pr["user"]["login"].casefold()
            and not pr["draft"]
            and not has_ready_label(pr)
        )
        trusted_approval = should_check_approval and has_trusted_approval(
            github,
            issue_number,
            trusted_users,
        )
        allowed, reason = authorize(
            actor=actor,
            permission=permission,
            pr=pr,
            trusted_approval=trusted_approval,
            trusted_users=trusted_users,
        )
        if not allowed:
            add_reaction_safely(github, comment_id, "-1")
            github.add_comment(issue_number, f"@{actor}, {reason}")
            return
        print(f"Authorized @{actor}: {reason}")
        if command == COMMAND_RUN_CI:
            message = handle_run_ci(
                actor=actor,
                buildkite=buildkite,
                comment_id=comment_id,
                github=github,
                pr=pr,
            )
        else:
            message = handle_retry_failed(buildkite=buildkite, pr=pr)
        add_reaction_safely(github, comment_id, "rocket")
        github.add_comment(issue_number, message)
    except Exception:
        add_reaction_safely(github, comment_id, "confused")
        raise
 def main() -> None:
    event_path = os.environ["GITHUB_EVENT_PATH"]
    with open(event_path, encoding="utf-8") as event_file:
        event = json.load(event_file)
    if not parse_command(event["comment"]["body"]):
        return
    github = GitHubClient(
        os.environ.get("GH_TOKEN", ""),
        os.environ["GITHUB_REPOSITORY"],
    )
    buildkite = BuildkiteClient(
        os.environ.get("BUILDKITE_API_TOKEN", ""),
        os.environ.get("BUILDKITE_ORGANIZATION", "vllm"),
        os.environ.get("BUILDKITE_PIPELINE", "ci"),
    )
    run(
        event,
        github,
        buildkite,
        os.environ.get("CI_TRUSTED_USERS", ""),
    )
 if __name__ == "__main__":
    main()
@@ -1,363 +0,0 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import unittest
 from typing import Any
 from run_ci_command import (
    COMMAND_RETRY_FAILED,
    COMMAND_RUN_CI,
    RETRY_STATES,
    BuildkiteClient,
    authorize,
    create_build_payload,
    has_trusted_approval,
    is_active_build,
    is_build_for_pr,
    parse_command,
    parse_trusted_users,
    run,
    select_latest_build,
 )
 def make_pr(**overrides: Any) -> dict[str, Any]:
    pr = {
        "base": {"ref": "main"},
        "draft": False,
        "head": {
            "ref": "feature",
            "repo": {"clone_url": "https://github.com/contributor/vllm.git"},
            "sha": "0123456789abcdef",
        },
        "labels": [],
        "number": 42,
        "state": "open",
        "user": {"login": "author"},
    }
    pr.update(overrides)
    return pr
 def make_event(command: str, actor: str = "reviewer") -> dict[str, Any]:
    return {
        "comment": {
            "body": command,
            "id": 99,
            "user": {"login": actor},
        },
        "issue": {
            "number": 42,
            "pull_request": {},
        },
    }
 class FakeGitHub:
    def __init__(
        self,
        *,
        permission: str = "write",
        permissions: dict[str, str] | None = None,
        pr: dict[str, Any] | None = None,
        review_decision: str = "REVIEW_REQUIRED",
        reviews: list[dict[str, Any]] | None = None,
    ) -> None:
        self.comments: list[str] = []
        self.permission = permission
        self.permissions = permissions or {}
        self.pr = pr or make_pr()
        self.reactions: list[str] = []
        self.review_decision = review_decision
        self.reviews = reviews or []
    def get_pr(self, number: int) -> dict[str, Any]:
        return self.pr
    def get_permission(self, actor: str) -> str:
        return self.permissions.get(actor, self.permission)
    def get_review_decision(self, number: int) -> str:
        return self.review_decision
    def list_reviews(self, number: int) -> list[dict[str, Any]]:
        return self.reviews
    def list_reactions(self, comment_id: int) -> list[dict[str, Any]]:
        return []
    def add_reaction(self, comment_id: int, content: str) -> None:
        self.reactions.append(content)
    def add_comment(self, issue_number: int, body: str) -> None:
        self.comments.append(body)
 class FakeBuildkite:
    def __init__(
        self,
        build_lists: list[list[dict[str, Any]]] | None = None,
    ) -> None:
        self.build_lists = build_lists or []
        self.created_builds: list[dict[str, Any]] = []
        self.list_calls: list[tuple[str, tuple[str, str] | None]] = []
        self.retry_calls: list[tuple[int, str]] = []
    def list_builds(
        self,
        commit: str,
        *,
        metadata: tuple[str, str] | None = None,
    ) -> list[dict[str, Any]]:
        self.list_calls.append((commit, metadata))
        return self.build_lists.pop(0)
    def create_build(self, body: dict[str, Any]) -> dict[str, Any]:
        self.created_builds.append(body)
        return {
            "number": 123,
            "web_url": "https://buildkite.example/builds/123",
        }
    def retry_failed_jobs(
        self,
        build_number: int,
        states: str,
    ) -> dict[str, Any]:
        self.retry_calls.append((build_number, states))
        return {"retried_jobs_count": 3}
 class FakeTransport:
    def __init__(self, response: Any) -> None:
        self.calls: list[dict[str, Any]] = []
        self.response = response
    def request(self, url: str, **kwargs: Any) -> Any:
        self.calls.append({"url": url, **kwargs})
        return self.response
 class RunCiCommandTest(unittest.TestCase):
    def test_only_exact_ci_commands_are_accepted(self) -> None:
        self.assertEqual(parse_command(COMMAND_RUN_CI), COMMAND_RUN_CI)
        self.assertEqual(
            parse_command(COMMAND_RETRY_FAILED),
            COMMAND_RETRY_FAILED,
        )
        self.assertIsNone(parse_command("/ci run please"))
        self.assertIsNone(parse_command(" /ci run"))
    def test_write_access_authorizes_reviewers_and_authors(self) -> None:
        allowed, _ = authorize(
            actor="reviewer",
            permission="write",
            pr=make_pr(),
        )
        self.assertTrue(allowed)
    def test_configured_trusted_contributors_can_run_ci(self) -> None:
        trusted_users = parse_trusted_users("trusted-one, TRUSTED-TWO")
        allowed, _ = authorize(
            actor="trusted-two",
            permission="read",
            pr=make_pr(),
            trusted_users=trusted_users,
        )
        self.assertTrue(allowed)
    def test_authors_need_an_approval_or_ready_label(self) -> None:
        pending, _ = authorize(
            actor="author",
            permission="read",
            pr=make_pr(),
        )
        approved, _ = authorize(
            actor="author",
            permission="read",
            pr=make_pr(),
            trusted_approval=True,
        )
        ready, _ = authorize(
            actor="author",
            permission="read",
            pr=make_pr(labels=[{"name": "ready"}]),
        )
        self.assertFalse(pending)
        self.assertTrue(approved)
        self.assertTrue(ready)
    def test_non_author_contributors_without_write_are_denied(self) -> None:
        allowed, _ = authorize(
            actor="contributor",
            permission="read",
            pr=make_pr(),
            trusted_approval=True,
        )
        self.assertFalse(allowed)
    def test_authors_cannot_use_ready_state_on_draft_prs(self) -> None:
        allowed, _ = authorize(
            actor="author",
            permission="read",
            pr=make_pr(draft=True, labels=[{"name": "ready"}]),
            trusted_approval=True,
        )
        self.assertFalse(allowed)
    def test_only_trusted_reviewers_can_delegate_through_approval(self) -> None:
        approved_review = {
            "state": "APPROVED",
            "user": {"login": "reviewer"},
        }
        trusted = FakeGitHub(
            permission="read",
            permissions={"reviewer": "write"},
            review_decision="APPROVED",
            reviews=[approved_review],
        )
        untrusted = FakeGitHub(
            permission="read",
            review_decision="APPROVED",
            reviews=[approved_review],
        )
        self.assertTrue(has_trusted_approval(trusted, 42, set()))
        self.assertFalse(has_trusted_approval(untrusted, 42, set()))
    def test_build_matching_is_scoped_to_the_pr(self) -> None:
        self.assertTrue(is_build_for_pr({"pull_request": {"id": 42}}, 42))
        self.assertFalse(is_build_for_pr({"pull_request": {"id": 43}}, 42))
        self.assertTrue(
            is_build_for_pr(
                {"meta_data": {"github-pr-number": "42"}},
                42,
            )
        )
    def test_latest_build_selection_ignores_other_prs(self) -> None:
        latest = select_latest_build(
            [
                {
                    "created_at": "2026-07-28T02:00:00Z",
                    "number": 3,
                    "pull_request": {"id": 43},
                },
                {
                    "created_at": "2026-07-28T01:00:00Z",
                    "number": 2,
                    "pull_request": {"id": 42},
                },
                {
                    "created_at": "2026-07-28T00:00:00Z",
                    "number": 1,
                    "pull_request": {"id": 42},
                },
            ],
            42,
        )
        self.assertEqual(latest["number"], 2)
    def test_active_build_states_prevent_duplicate_runs(self) -> None:
        self.assertTrue(is_active_build({"state": "scheduled"}))
        self.assertTrue(is_active_build({"state": "running"}))
        self.assertTrue(is_active_build({"state": "waiting"}))
        self.assertTrue(is_active_build({"blocked": True, "state": "passed"}))
        self.assertFalse(is_active_build({"state": "failed"}))
    def test_build_payload_preserves_pr_context(self) -> None:
        payload = create_build_payload(
            actor="reviewer",
            comment_id=99,
            pr=make_pr(labels=[{"name": "ready"}, {"name": "v1"}]),
        )
        self.assertEqual(
            payload,
            {
                "commit": "0123456789abcdef",
                "branch": "feature",
                "message": "PR #42 /ci run by @reviewer",
                "pull_request_id": 42,
                "pull_request_base_branch": "main",
                "pull_request_repository": ("https://github.com/contributor/vllm.git"),
                "pull_request_labels": ["ready", "v1"],
                "ignore_pipeline_branch_filters": True,
                "env": {
                    "VLLM_CI_GITHUB_COMMENT_ID": "99",
                    "VLLM_CI_TRIGGERED_BY": "reviewer",
                },
                "meta_data": {
                    "github-comment-id": "99",
                    "github-pr-number": "42",
                    "github-triggered-by": "reviewer",
                },
            },
        )
    def test_ci_run_dispatches_build_with_current_pr_metadata(self) -> None:
        github = FakeGitHub()
        buildkite = FakeBuildkite([[], []])
        run(make_event(COMMAND_RUN_CI), github, buildkite)
        self.assertEqual(len(buildkite.created_builds), 1)
        self.assertEqual(
            buildkite.created_builds[0]["message"],
            "PR #42 /ci run by @reviewer",
        )
        self.assertEqual(github.reactions, ["eyes", "rocket"])
        self.assertIn("Buildkite CI #123", github.comments[0])
    def test_unapproved_authors_are_denied_without_buildkite(self) -> None:
        github = FakeGitHub(
            permission="read",
            pr=make_pr(),
            review_decision="REVIEW_REQUIRED",
        )
        buildkite = FakeBuildkite()
        run(make_event(COMMAND_RUN_CI, "author"), github, buildkite)
        self.assertEqual(buildkite.list_calls, [])
        self.assertEqual(github.reactions, ["eyes", "-1"])
        self.assertIn("approve the PR", github.comments[0])
    def test_ci_retry_uses_latest_current_sha_build(self) -> None:
        github = FakeGitHub(
            permission="read",
            pr=make_pr(labels=[{"name": "ready"}]),
        )
        buildkite = FakeBuildkite(
            [
                [
                    {
                        "created_at": "2026-07-28T01:00:00Z",
                        "finished_at": "2026-07-28T02:00:00Z",
                        "number": 123,
                        "pull_request": {"id": 42},
                        "state": "failed",
                        "web_url": "https://buildkite.example/builds/123",
                    }
                ]
            ]
        )
        run(make_event(COMMAND_RETRY_FAILED, "author"), github, buildkite)
        self.assertEqual(buildkite.retry_calls, [(123, RETRY_STATES)])
        self.assertIn("Queued 3 failed job", github.comments[0])
    def test_buildkite_retry_uses_retry_failed_jobs_endpoint(self) -> None:
        transport = FakeTransport({"retried_jobs_count": 2})
        client = BuildkiteClient(
            "secret",
            "vllm",
            "ci",
            transport=transport,
        )
        client.retry_failed_jobs(123, RETRY_STATES)
        call = transport.calls[0]
        self.assertEqual(call["method"], "PUT")
        self.assertTrue(call["url"].endswith("/123/retry_failed_jobs"))
        self.assertEqual(call["body"], {"states": RETRY_STATES})
 if __name__ == "__main__":
    unittest.main()
@@ -18,9 +18,6 @@ vllm/third_party/deep_gemm/
 # fmha_sm100 vendored package built from source
 vllm/third_party/fmha_sm100/
 # tml-fa4 vendored package built from source
 vllm/third_party/tml_fa4/
 # triton jit
 .triton
@@ -173,6 +170,9 @@ venv.bak/
 # mkdocs documentation
 /site
 docs/argparse
 docs/examples/*
 !docs/examples/README.md
 # mypy
 .mypy_cache/
@@ -3,9 +3,6 @@ MD007:
 MD013: false
 MD024:
  siblings_only: true
 MD025:
  # Allow front matter title to be different from the first heading in the document.
  front_matter_title: ""
 MD031:
  list_items: false
 MD033: false
@@ -4,7 +4,7 @@ default_install_hook_types:
 default_stages:
  - pre-commit # Run locally
  - manual # Run in CI
-exclude: 'vllm/third_party/.*|vllm/models/kimi_k3/nvidia/ops/third_party/.*|vllm/models/kimi_k3/amd/ops/third_party/.*'
+exclude: 'vllm/third_party/.*'
 repos:
 - repo: https://github.com/astral-sh/ruff-pre-commit
  rev: v0.14.0
@@ -30,7 +30,7 @@ repos:
  - id: markdownlint-cli2
    language_version: lts
    args: [--fix]
-    exclude: (^|/)CLAUDE\.md$
+    exclude: ^CLAUDE\.md$
 - repo: https://github.com/rhysd/actionlint
  rev: v1.7.7
  hooks:
@@ -210,7 +210,7 @@ repos:
    name: Check SPDX headers
    entry: python tools/pre_commit/check_spdx_header.py
    language: python
-    types_or: [python, rust, proto]
+    types: [python]
  - id: check-root-lazy-imports
    name: Check root lazy imports
    entry: python tools/pre_commit/check_init_lazy_imports.py
@@ -260,6 +260,10 @@ repos:
    files: ^docker/(Dockerfile|versions\.json)$
    pass_filenames: false
    additional_dependencies: [dockerfile-parse]
  - id: attention-backend-docs
    name: Check attention backend documentation is up to date
    entry: python tools/pre_commit/generate_attention_backend_docs.py --check
    language: python
  - id: check-boolean-context-manager
    name: Check for boolean ops in with-statements
    entry: python tools/pre_commit/check_boolean_context_manager.py
@@ -0,0 +1,2 @@
 collect_env.py
 vllm/model_executor/layers/fla/ops/*.py
@@ -68,8 +68,8 @@ endif()
 # requirements.txt files and should be kept consistent.  The ROCm torch
 # versions are derived from docker/Dockerfile.rocm
 #
-set(TORCH_SUPPORTED_VERSION_CUDA "2.13.0")
+set(TORCH_SUPPORTED_VERSION_CUDA "2.11.0")
-set(TORCH_SUPPORTED_VERSION_ROCM "2.13.0")
+set(TORCH_SUPPORTED_VERSION_ROCM "2.11.0")
 # TORCH_NIGHTLY=1 builds run against unpinned nightly wheels, so the supported-
 # version check would always warn. Only treat it as a nightly build when the
 # value is exactly "1" (the bootstrap exports TORCH_NIGHTLY=0 by default, which
@@ -114,11 +114,6 @@ find_package(Torch REQUIRED)
 # Supported NVIDIA architectures.
 # This check must happen after find_package(Torch) because that's when CMAKE_CUDA_COMPILER_VERSION gets defined
 if(DEFINED CMAKE_CUDA_COMPILER_VERSION AND
   CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.4)
  # Rubin (10.7) can run SM100 family code, but CUDA 13.4 also supports
  # targeting it directly.
  set(CUDA_SUPPORTED_ARCHS "7.5;8.0;8.6;8.7;8.9;9.0;10.0;10.7;11.0;12.0")
 elseif(DEFINED CMAKE_CUDA_COMPILER_VERSION AND
   CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.0)
   # starting from CUDA 12.9 and Blackwell (10.0), we use family-specific targets (10.0f, 12.0f, etc)
   # to support the whole generation without specifying all sub-architectures
@@ -219,8 +214,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
  # the set of architectures we want to compile for and remove the from the
  # CMAKE_CUDA_FLAGS so that they are not applied globally.
  #
-  # `+PTX` in TORCH_CUDA_ARCH_LIST is not preserved here. If a kernel really
+  # `+PTX` in TORCH_CUDA_ARCH_LIST is not preserved here. It is emitted by torch
-  # needs PTX, add `+PTX` to that kernel's component-specific arch list below.
+  # as `code=compute_*`, while extract_unique_cuda_archs_ascending() records only
  # `arch=compute_*`. If a kernel really needs PTX, add `+PTX` to that kernel's
  # component-specific arch list below.
  #
  clear_cuda_arches(CUDA_ARCH_FLAGS)
  extract_unique_cuda_archs_ascending(CUDA_ARCHS "${CUDA_ARCH_FLAGS}")
@@ -230,13 +227,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
  cuda_archs_loose_intersection(CUDA_ARCHS
    "${CUDA_SUPPORTED_ARCHS}" "${CUDA_ARCHS}")
  message(STATUS "CUDA supported target architectures: ${CUDA_ARCHS}")
  if(NOT CUDA_ARCHS)
    message(FATAL_ERROR
      "No supported CUDA architectures; the build would produce a binary "
      "with no usable kernels. Detected gencode flags: ${CUDA_ARCH_FLAGS}; "
      "supported: ${CUDA_SUPPORTED_ARCHS}. "
      "Set TORCH_CUDA_ARCH_LIST for your GPU (e.g. 12.0).")
  endif()
 else()
  #
  # For other GPU targets override the GPU architectures detected by cmake/torch
@@ -421,11 +411,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
    "csrc/libtorch_stable/mamba/selective_scan_fwd.cu"
    "csrc/libtorch_stable/cache_kernels.cu"
    "csrc/libtorch_stable/cache_kernels_fused.cu"
    "csrc/libtorch_stable/custom_all_gather_reduce_scatter.cu"
    "csrc/libtorch_stable/custom_all_gather_reduce_scatter_ops.cpp"
    "csrc/libtorch_stable/custom_all_reduce.cu"
-    "csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu"
+    "csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")
    "csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu")
  if(VLLM_GPU_LANG STREQUAL "CUDA" AND
     DEFINED CMAKE_CUDA_COMPILER_VERSION AND
@@ -433,7 +420,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
    if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
      cuda_archs_loose_intersection(COOPERATIVE_TOPK_ARCHS
-        "9.0a;10.0f;10.1f;10.3f;10.7f;11.0f;12.0f;12.1f" "${CUDA_ARCHS}")
+        "9.0a;10.0f;10.1f;10.3f;11.0f;12.0f;12.1f" "${CUDA_ARCHS}")
    else()
      cuda_archs_loose_intersection(COOPERATIVE_TOPK_ARCHS
        "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
@@ -708,7 +695,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
  # DeepSeek V3 fused A GEMM kernel (requires SM 9.0+, Hopper and later)
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0f;10.7f;11.0f;12.0f" "${CUDA_ARCHS}")
+    cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0f;11.0f;12.0f" "${CUDA_ARCHS}")
  else()
    cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
  endif()
@@ -828,7 +815,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
  # The cutlass_scaled_mm kernels for Blackwell SM100 (c3x, i.e. CUTLASS 3.x)
  # require CUDA 12.8 or later
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;10.7f;11.0f" "${CUDA_ARCHS}")
+    cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
  else()
    cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
  endif()
@@ -912,7 +899,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
  endif()
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;10.7f;11.0f" "${CUDA_ARCHS}")
+    cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
  else()
    cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
  endif()
@@ -937,7 +924,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
  # moe_data.cu is used by all CUTLASS MoE kernels.
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0f;10.7f;11.0f;12.0f" "${CUDA_ARCHS}")
+    cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0f;11.0f;12.0f" "${CUDA_ARCHS}")
  else()
    cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
  endif()
@@ -994,7 +981,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
  # SM10x/11x FP4 kernels. MXFP4 experts quantization is currently compiled
  # only in this block; SM12x has separate NVFP4 matmul/MoE kernels above.
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(FP4_SM100_ARCHS "10.0f;10.7f;11.0f" "${CUDA_ARCHS}")
+    cuda_archs_loose_intersection(FP4_SM100_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
  else()
    cuda_archs_loose_intersection(FP4_SM100_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
  endif()
@@ -1060,7 +1047,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
  # Runtime dispatch is gated in
  # vllm/v1/attention/backends/mla/cutlass_mla.py.
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(MLA_ARCHS "10.0f;10.7f;11.0f" "${CUDA_ARCHS}")
+    cuda_archs_loose_intersection(MLA_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
  else()
    cuda_archs_loose_intersection(MLA_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
  endif()
@@ -1082,41 +1069,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
    set(MLA_ARCHS)
  endif()
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
    cuda_archs_loose_intersection(FUSED_KDA_DECODE_ARCHS
      "9.0a;10.0f;12.0f" "${CUDA_ARCHS}")
  endif()
  if(FUSED_KDA_DECODE_ARCHS)
    set(FUSED_KDA_DECODE_SRC
      "csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel.cu")
    set_gencode_flags_for_srcs(
      SRCS "${FUSED_KDA_DECODE_SRC}"
      CUDA_ARCHS "${FUSED_KDA_DECODE_ARCHS}")
    set_property(SOURCE ${FUSED_KDA_DECODE_SRC} APPEND PROPERTY
      COMPILE_OPTIONS "$<$<COMPILE_LANGUAGE:CUDA>:--use_fast_math>")
    list(APPEND VLLM_STABLE_EXT_SRC "${FUSED_KDA_DECODE_SRC}")
    message(STATUS
      "Building fused KDA decode for archs: ${FUSED_KDA_DECODE_ARCHS}")
  endif()
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
    cuda_archs_loose_intersection(KIMI_K3_ATTN_RES_ARCHS
      "10.0f" "${CUDA_ARCHS}")
  endif()
  if(KIMI_K3_ATTN_RES_ARCHS)
    set(KIMI_K3_ATTN_RES_SRC
      "csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu")
    set_gencode_flags_for_srcs(
      SRCS "${KIMI_K3_ATTN_RES_SRC}"
      CUDA_ARCHS "${KIMI_K3_ATTN_RES_ARCHS}")
    set_property(SOURCE ${KIMI_K3_ATTN_RES_SRC} APPEND PROPERTY
      COMPILE_OPTIONS
      "$<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr;--expt-extended-lambda;--use_fast_math>")
    list(APPEND VLLM_STABLE_EXT_SRC "${KIMI_K3_ATTN_RES_SRC}")
    message(STATUS
      "Building Kimi K3 AttnRes for archs: ${KIMI_K3_ATTN_RES_ARCHS}")
  endif()
  # Hadacore kernels
  cuda_archs_loose_intersection(HADACORE_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}")
  if(HADACORE_ARCHS)
@@ -1158,14 +1110,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
      target_compile_definitions(_C_stable_libtorch PRIVATE
        VLLM_ENABLE_COOPERATIVE_TOPK=1)
    endif()
    if(FUSED_KDA_DECODE_ARCHS)
      target_compile_definitions(_C_stable_libtorch PRIVATE
        VLLM_ENABLE_FUSED_KDA_DECODE=1)
    endif()
    if(KIMI_K3_ATTN_RES_ARCHS)
      target_compile_definitions(_C_stable_libtorch PRIVATE
        VLLM_ENABLE_KIMI_K3_ATTN_RES=1)
    endif()
    # Needed by CUTLASS kernels
    target_compile_definitions(_C_stable_libtorch PRIVATE
      CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
@@ -1420,7 +1364,6 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
  set(VLLM_ROCM_EXT_SRC
    "csrc/rocm/torch_bindings.cpp"
    "csrc/rocm/skinny_gemms.cu"
    "csrc/rocm/skinny_gemms_int4.cu"
    "csrc/rocm/attention.cu")
  set(VLLM_ROCM_HAS_GFX1100 OFF)
@@ -1463,9 +1406,7 @@ if (VLLM_GPU_LANG STREQUAL "CUDA")
    include(cmake/external_projects/deepgemm.cmake)
    include(cmake/external_projects/fmha_sm100.cmake)
    include(cmake/external_projects/flashmla.cmake)
    include(cmake/external_projects/flashkda.cmake)
    include(cmake/external_projects/qutlass.cmake)
    include(cmake/external_projects/tml_fa4.cmake)
    # vllm-flash-attn should be last as it overwrites some CMake functions
    include(cmake/external_projects/vllm_flash_attn.cmake)
@@ -48,7 +48,7 @@ vLLM is flexible and easy to use with:
 - Tool calling and reasoning parsers
 - OpenAI-compatible API server, plus Anthropic Messages API and gRPC support
 - Efficient multi-LoRA support for dense and MoE layers
- Support for NVIDIA GPUs, AMD GPUs, Intel GPUs, and x86/ARM/PowerPC CPUs. Additionally, diverse hardware plugins such as Google TPUs, Intel Gaudi, IBM Spyre, Huawei Ascend, Rebellions NPU, Apple Silicon, MetaX GPU, and more.
+- Support for NVIDIA GPUs, AMD GPUs, and x86/ARM/PowerPC CPUs. Additionally, diverse hardware plugins such as Google TPUs, Intel Gaudi, IBM Spyre, Huawei Ascend, Rebellions NPU, Apple Silicon, MetaX GPU, and more.
 vLLM seamlessly supports 200+ model architectures on Hugging Face, including:
@@ -75,11 +75,7 @@ def run_mla_benchmark(config: BenchmarkConfig, **kwargs) -> BenchmarkResult:
    from mla_runner import run_mla_benchmark as run_mla
    return run_mla(
-        config.backend,
+        config.backend, config, prefill_backend=config.prefill_backend, **kwargs
        config,
        prefill_backend=config.prefill_backend,
        sparse_mla_force_mqa=config.sparse_mla_force_mqa,
        **kwargs,
    )
@@ -596,30 +592,6 @@ def main():
        default="profile",
        help="Output file name for ncu profile (default: 'profile').",
    )
    parser.add_argument(
        "--torch-profile",
        action="store_true",
        default=False,
        help="Collect a PyTorch profiler Chrome trace for each benchmark run.",
    )
    parser.add_argument(
        "--torch-profile-dir",
        default=None,
        help="Directory for PyTorch profiler traces.",
    )
    parser.add_argument(
        "--torch-profile-iters",
        type=int,
        default=3,
        help="Number of forward passes to record per PyTorch profiler trace.",
    )
    parser.add_argument(
        "--sparse-mla-mha-variants",
        nargs="+",
        default=None,
        choices=["dense_mha", "mqa"],
        help="Sparse MLA variants to run in mha_vs_mqa mode. Defaults to both.",
    )
    # Parameter sweep (use YAML config for advanced sweeps)
    parser.add_argument(
@@ -669,7 +641,6 @@ def main():
        # Prefill backends (e.g., ["fa3", "fa4"])
        args.prefill_backends = yaml_config.get("prefill_backends", None)
        args.prefill_backend = yaml_config.get("prefill_backend", None)
        # FP8 output benchmark knobs; CLI wins.
        if args.fp8_output_scale is None:
@@ -712,9 +683,6 @@ def main():
            args.num_q_heads = model.get("num_q_heads", args.num_q_heads)
            args.num_kv_heads = model.get("num_kv_heads", args.num_kv_heads)
            args.block_size = model.get("block_size", args.block_size)
            args.max_model_len = model.get(
                "max_model_len", getattr(args, "max_model_len", None)
            )
            # MLA-specific dimensions
            args.kv_lora_rank = model.get("kv_lora_rank", args.kv_lora_rank)
            args.qk_nope_head_dim = model.get("qk_nope_head_dim", args.qk_nope_head_dim)
@@ -733,21 +701,6 @@ def main():
            args.cuda_graphs = yaml_config["cuda_graphs"]
        if "ncu_profile" in yaml_config:
            args.ncu_profile = yaml_config["ncu_profile"]
        if "torch_profile" in yaml_config:
            args.torch_profile = yaml_config["torch_profile"]
        if "torch_profile_dir" in yaml_config:
            args.torch_profile_dir = yaml_config["torch_profile_dir"]
        if "torch_profile_iters" in yaml_config:
            args.torch_profile_iters = yaml_config["torch_profile_iters"]
        args.sparse_mla_topk_pattern = yaml_config.get(
            "sparse_mla_topk_pattern", "random"
        )
        args.sparse_mla_dense_mha_max_seq_len = yaml_config.get(
            "sparse_mla_dense_mha_max_seq_len", None
        )
        args.sparse_mla_mha_variants = yaml_config.get(
            "sparse_mla_mha_variants", args.sparse_mla_mha_variants
        )
        # Parameter sweep configuration
        if "parameter_sweep" in yaml_config:
@@ -889,6 +842,8 @@ def main():
                        num_kv_heads=args.num_kv_heads,
                        block_size=args.block_size,
                        device=args.device,
                        repeats=args.repeats,
                        warmup_iters=args.warmup_iters,
                        profile_memory=args.profile_memory,
                        kv_cache_dtype=args.kv_cache_dtype,
                        use_cuda_graphs=args.cuda_graphs,
@@ -1108,133 +1063,6 @@ def main():
                    f"\n  [yellow]Prefill always faster for batch_size={bs}[/]"
                )
    # Handle MHA vs MQA comparison mode for sparse MLA
    elif hasattr(args, "mode") and args.mode == "mha_vs_mqa":
        console.print("[yellow]Mode: MHA vs MQA comparison for sparse MLA[/]")
        sparse_mla_topk_pattern = getattr(args, "sparse_mla_topk_pattern", "random")
        dense_mha_max_seq_len = getattr(args, "sparse_mla_dense_mha_max_seq_len", None)
        prefill_backend = getattr(args, "prefill_backend", None)
        if prefill_backend:
            console.print(f"Prefill backend: {prefill_backend}")
        available_variants = [
            ("dense_mha", False, "dense"),
            ("mqa", True, "auto"),
        ]
        requested_variants = getattr(args, "sparse_mla_mha_variants", None)
        if requested_variants is not None:
            valid_variants = {label for label, _, _ in available_variants}
            invalid_variants = sorted(set(requested_variants) - valid_variants)
            if invalid_variants:
                raise ValueError(
                    "Invalid sparse_mla_mha_variants entries: "
                    f"{invalid_variants}. Valid variants are: "
                    f"{sorted(valid_variants)}"
                )
            requested_variant_set = set(requested_variants)
            variants = [
                variant
                for variant in available_variants
                if variant[0] in requested_variant_set
            ]
        else:
            variants = available_variants
        formatter = ResultsFormatter(console)
        total = 0
        for spec in args.batch_specs:
            q_len = max(request.q_len for request in parse_batch_spec(spec))
            for variant_label, _, _ in variants:
                if (
                    variant_label == "dense_mha"
                    and dense_mha_max_seq_len is not None
                    and q_len > dense_mha_max_seq_len
                ):
                    continue
                total += len(backends)
        with tqdm(total=total, desc="Benchmarking") as pbar:
            for spec in args.batch_specs:
                q_len = max(request.q_len for request in parse_batch_spec(spec))
                for backend in backends:
                    for variant_label, force_mqa, mha_mode in variants:
                        if (
                            variant_label == "dense_mha"
                            and dense_mha_max_seq_len is not None
                            and q_len > dense_mha_max_seq_len
                        ):
                            continue
                        config = BenchmarkConfig(
                            backend=f"{backend}_{variant_label}",
                            batch_spec=spec,
                            num_layers=args.num_layers,
                            head_dim=args.head_dim,
                            num_q_heads=args.num_q_heads,
                            num_kv_heads=args.num_kv_heads,
                            block_size=args.block_size,
                            device=args.device,
                            max_model_len=getattr(args, "max_model_len", None),
                            kv_cache_dtype=args.kv_cache_dtype,
                            profile_memory=args.profile_memory,
                            use_cuda_graphs=args.cuda_graphs,
                            ncu_profile=args.ncu_profile,
                            torch_profile=args.torch_profile,
                            torch_profile_dir=args.torch_profile_dir,
                            torch_profile_iters=args.torch_profile_iters,
                            warmup_ms=args.warmup_ms,
                            kv_lora_rank=getattr(args, "kv_lora_rank", None),
                            qk_nope_head_dim=getattr(args, "qk_nope_head_dim", None),
                            qk_rope_head_dim=getattr(args, "qk_rope_head_dim", None),
                            v_head_dim=getattr(args, "v_head_dim", None),
                            sparse_mla_force_mqa=force_mqa,
                            sparse_mla_mha_mode=mha_mode,
                            sparse_mla_dense_mha_max_seq_len=dense_mha_max_seq_len,
                            sparse_mla_topk_pattern=sparse_mla_topk_pattern,
                            prefill_backend=prefill_backend,
                        )
                        # run_mla_benchmark needs the real backend name
                        from mla_runner import run_mla_benchmark as run_mla
                        run_label = f"{backend}_{variant_label} {spec}"
                        pbar.set_postfix_str(run_label)
                        try:
                            result = run_mla(
                                backend,
                                config,
                                prefill_backend=prefill_backend,
                                sparse_mla_force_mqa=force_mqa,
                            )
                        except Exception as e:
                            result = BenchmarkResult(
                                config=config,
                                mean_time=float("inf"),
                                median_time=float("inf"),
                                std_time=0,
                                min_time=float("inf"),
                                max_time=float("inf"),
                                error=str(e),
                            )
                        all_results.append(result)
                        if args.output_csv:
                            formatter.save_csv(all_results, args.output_csv)
                        if args.output_json:
                            formatter.save_json(all_results, args.output_json)
                        if not result.success:
                            console.print(
                                f"[red]Error {backend}_{variant_label} "
                                f"{spec}: {result.error}[/]"
                            )
                        pbar.update(1)
        # Display results with variant labels as separate "backends"
        console.print("\n[bold green]MHA vs MQA Results:[/]")
        variant_backends = [f"{b}_{v}" for b in backends for v, _, _ in variants]
        formatter.print_table(all_results, variant_backends)
    # Handle model parameter sweep mode
    elif hasattr(args, "model_parameter_sweep") and args.model_parameter_sweep:
        # Model parameter sweep
@@ -1358,10 +1186,6 @@ def main():
                            profile_memory=args.profile_memory,
                            warmup_ms=args.warmup_ms,
                            prefill_backend=pb,
                            kv_lora_rank=args.kv_lora_rank,
                            qk_nope_head_dim=args.qk_nope_head_dim,
                            qk_rope_head_dim=args.qk_rope_head_dim,
                            v_head_dim=args.v_head_dim,
                        )
                        result = run_benchmark(config)
@@ -4,10 +4,8 @@
 """Common utilities for attention benchmarking."""
 import csv
 import gc
 import json
 import math
 from collections.abc import Sequence
 from dataclasses import asdict, dataclass
 from pathlib import Path
 from typing import Any
@@ -46,13 +44,10 @@ def run_do_bench(
    kwargs: dict[str, Any] = {"return_mode": "all"}
    if use_cuda_graphs:
        result = triton.testing.do_bench_cudagraph(benchmark_fn, **kwargs)
        gc.collect()
        torch.accelerator.empty_cache()
    else:
        if warmup_ms is not None:
            kwargs["warmup"] = warmup_ms
        result = triton.testing.do_bench(benchmark_fn, **kwargs)
    torch.accelerator.synchronize()
    return result
@@ -96,6 +91,42 @@ except ImportError:
    AttentionLayerBase = object  # Fallback
 class MockKVBProj:
    """Mock KV projection layer for MLA prefill mode.
    Mimics ColumnParallelLinear behavior for kv_b_proj in MLA backends.
    Projects kv_c_normed to [qk_nope_head_dim + v_head_dim] per head.
    """
    def __init__(self, num_heads: int, qk_nope_head_dim: int, v_head_dim: int):
        self.num_heads = num_heads
        self.qk_nope_head_dim = qk_nope_head_dim
        self.v_head_dim = v_head_dim
        self.out_dim = qk_nope_head_dim + v_head_dim
        self.weight = torch.empty(0, dtype=torch.bfloat16)
    def __call__(self, x: torch.Tensor) -> tuple[torch.Tensor]:
        """
        Project kv_c_normed to output space.
        Args:
            x: Input tensor [num_tokens, kv_lora_rank]
        Returns:
            Tuple containing output tensor
                [num_tokens, num_heads, qk_nope_head_dim + v_head_dim]
        """
        num_tokens = x.shape[0]
        result = torch.randn(
            num_tokens,
            self.num_heads,
            self.out_dim,
            device=x.device,
            dtype=x.dtype,
        )
        return (result,)  # Return as tuple to match ColumnParallelLinear API
 class MockIndexer:
    """Mock Indexer for sparse MLA backends.
@@ -127,60 +158,6 @@ class MockIndexer:
        )
        self.topk_indices_buffer[:num_tokens] = indices
    def fill_indices(
        self,
        num_tokens: int,
        max_kv_len: int,
        pattern: str = "random",
        requests: Sequence[Any] | None = None,
    ):
        if pattern == "random":
            self.fill_random_indices(num_tokens, max_kv_len)
            return
        if pattern == "prefix":
            indices = torch.arange(
                self.topk_tokens,
                dtype=torch.int32,
                device=self.topk_indices_buffer.device,
            )
            indices = (indices % max_kv_len).expand(num_tokens, -1)
            self.topk_indices_buffer[:num_tokens] = indices
            return
        if pattern == "sliding_window":
            if requests is None:
                start = max(max_kv_len - self.topk_tokens, 0)
                indices = torch.arange(
                    start,
                    start + self.topk_tokens,
                    dtype=torch.int32,
                    device=self.topk_indices_buffer.device,
                )
                indices = indices.clamp(max=max_kv_len - 1).expand(num_tokens, -1)
                self.topk_indices_buffer[:num_tokens] = indices
                return
            rows = []
            offsets = torch.arange(
                self.topk_tokens,
                dtype=torch.int32,
                device=self.topk_indices_buffer.device,
            ) - (self.topk_tokens - 1)
            for request in requests:
                q_len = request.q_len
                kv_len = request.kv_len
                context_len = kv_len - q_len
                positions = torch.arange(
                    context_len,
                    kv_len,
                    dtype=torch.int32,
                    device=self.topk_indices_buffer.device,
                )
                row_indices = positions[:, None] + offsets[None, :]
                rows.append(row_indices.clamp(min=0, max=kv_len - 1))
            self.topk_indices_buffer[:num_tokens] = torch.cat(rows, dim=0)
            return
        raise ValueError(f"Unknown sparse MLA topk pattern: {pattern}")
 class MockLayer(AttentionLayerBase):
    """Mock attention layer with scale parameters and impl.
@@ -275,14 +252,10 @@ class BenchmarkConfig:
    num_kv_heads: int
    block_size: int
    device: str
    max_model_len: int | None = None
    dtype: torch.dtype = torch.float16
    profile_memory: bool = False
-    use_cuda_graphs: bool = True
+    use_cuda_graphs: bool = False
    ncu_profile: bool = False
    torch_profile: bool = False
    torch_profile_dir: str | None = None
    torch_profile_iters: int = 3
    warmup_ms: int | None = None
    # "auto" or "fp8"
@@ -298,10 +271,6 @@ class BenchmarkConfig:
    # Backend-specific tuning
    num_kv_splits: int | None = None  # CUTLASS MLA
    reorder_batch_threshold: int | None = None  # FlashAttn MLA, FlashMLA
    sparse_mla_force_mqa: bool = False  # Force MQA path for sparse MLA
    sparse_mla_mha_mode: str = "auto"  # "auto" or "dense"
    sparse_mla_dense_mha_max_seq_len: int | None = None
    sparse_mla_topk_pattern: str = "random"  # "random", "prefix", "sliding_window"
    num_splits: int | None = None  # FlashAttention split-K (0=auto, 1=disabled)
@@ -1,474 +0,0 @@
 # Sparse MLA benchmark: forward_mha vs forward_mqa
 #
 # Usage:
 #   python benchmark.py --config configs/mla_sparse_mha_vs_mqa.yaml
 #
 # Heatmap grid:
 #   - batch_size: 1, 2, 4, 8, 16, 32
 #   - seq_len: 32, 64, 128, 256, 512, 1024, 2048
 #   - q_len: powers of two through seq_len
 #
 # Specs with q_len < seq_len include context; the q_len == seq_len diagonal
 # covers pure prefill.
 # The model shape below is the DP case. For the TP8 run, manually change
 # model.num_q_heads from 128 to 16 before rerunning this benchmark.
 mode: mha_vs_mqa
 model:
  name: "deepseek-v3"
  num_layers: 60
  num_q_heads: 128
  num_kv_heads: 1
  head_dim: 576
  kv_lora_rank: 512
  qk_nope_head_dim: 128
  qk_rope_head_dim: 64
  v_head_dim: 128
  block_size: 128
  max_model_len: 2048
 batch_specs:
  # Batch size 1
  # seq_len = 32
  - "1q1s32"
  - "1q2s32"
  - "1q4s32"
  - "1q8s32"
  - "1q16s32"
  - "1q32"
  # seq_len = 64
  - "1q1s64"
  - "1q2s64"
  - "1q4s64"
  - "1q8s64"
  - "1q16s64"
  - "1q32s64"
  - "1q64"
  # seq_len = 128
  - "1q1s128"
  - "1q2s128"
  - "1q4s128"
  - "1q8s128"
  - "1q16s128"
  - "1q32s128"
  - "1q64s128"
  - "1q128"
  # seq_len = 256
  - "1q1s256"
  - "1q2s256"
  - "1q4s256"
  - "1q8s256"
  - "1q16s256"
  - "1q32s256"
  - "1q64s256"
  - "1q128s256"
  - "1q256"
  # seq_len = 512
  - "1q1s512"
  - "1q2s512"
  - "1q4s512"
  - "1q8s512"
  - "1q16s512"
  - "1q32s512"
  - "1q64s512"
  - "1q128s512"
  - "1q256s512"
  - "1q512"
  # seq_len = 1024
  - "1q1s1024"
  - "1q2s1024"
  - "1q4s1024"
  - "1q8s1024"
  - "1q16s1024"
  - "1q32s1024"
  - "1q64s1024"
  - "1q128s1024"
  - "1q256s1024"
  - "1q512s1024"
  - "1q1024"
  # seq_len = 2048
  - "1q1s2048"
  - "1q2s2048"
  - "1q4s2048"
  - "1q8s2048"
  - "1q16s2048"
  - "1q32s2048"
  - "1q64s2048"
  - "1q128s2048"
  - "1q256s2048"
  - "1q512s2048"
  - "1q1024s2048"
  - "1q2048"
  # Batch size 2
  # seq_len = 32
  - "2q1s32"
  - "2q2s32"
  - "2q4s32"
  - "2q8s32"
  - "2q16s32"
  - "2q32"
  # seq_len = 64
  - "2q1s64"
  - "2q2s64"
  - "2q4s64"
  - "2q8s64"
  - "2q16s64"
  - "2q32s64"
  - "2q64"
  # seq_len = 128
  - "2q1s128"
  - "2q2s128"
  - "2q4s128"
  - "2q8s128"
  - "2q16s128"
  - "2q32s128"
  - "2q64s128"
  - "2q128"
  # seq_len = 256
  - "2q1s256"
  - "2q2s256"
  - "2q4s256"
  - "2q8s256"
  - "2q16s256"
  - "2q32s256"
  - "2q64s256"
  - "2q128s256"
  - "2q256"
  # seq_len = 512
  - "2q1s512"
  - "2q2s512"
  - "2q4s512"
  - "2q8s512"
  - "2q16s512"
  - "2q32s512"
  - "2q64s512"
  - "2q128s512"
  - "2q256s512"
  - "2q512"
  # seq_len = 1024
  - "2q1s1024"
  - "2q2s1024"
  - "2q4s1024"
  - "2q8s1024"
  - "2q16s1024"
  - "2q32s1024"
  - "2q64s1024"
  - "2q128s1024"
  - "2q256s1024"
  - "2q512s1024"
  - "2q1024"
  # seq_len = 2048
  - "2q1s2048"
  - "2q2s2048"
  - "2q4s2048"
  - "2q8s2048"
  - "2q16s2048"
  - "2q32s2048"
  - "2q64s2048"
  - "2q128s2048"
  - "2q256s2048"
  - "2q512s2048"
  - "2q1024s2048"
  - "2q2048"
  # Batch size 4
  # seq_len = 32
  - "4q1s32"
  - "4q2s32"
  - "4q4s32"
  - "4q8s32"
  - "4q16s32"
  - "4q32"
  # seq_len = 64
  - "4q1s64"
  - "4q2s64"
  - "4q4s64"
  - "4q8s64"
  - "4q16s64"
  - "4q32s64"
  - "4q64"
  # seq_len = 128
  - "4q1s128"
  - "4q2s128"
  - "4q4s128"
  - "4q8s128"
  - "4q16s128"
  - "4q32s128"
  - "4q64s128"
  - "4q128"
  # seq_len = 256
  - "4q1s256"
  - "4q2s256"
  - "4q4s256"
  - "4q8s256"
  - "4q16s256"
  - "4q32s256"
  - "4q64s256"
  - "4q128s256"
  - "4q256"
  # seq_len = 512
  - "4q1s512"
  - "4q2s512"
  - "4q4s512"
  - "4q8s512"
  - "4q16s512"
  - "4q32s512"
  - "4q64s512"
  - "4q128s512"
  - "4q256s512"
  - "4q512"
  # seq_len = 1024
  - "4q1s1024"
  - "4q2s1024"
  - "4q4s1024"
  - "4q8s1024"
  - "4q16s1024"
  - "4q32s1024"
  - "4q64s1024"
  - "4q128s1024"
  - "4q256s1024"
  - "4q512s1024"
  - "4q1024"
  # seq_len = 2048
  - "4q1s2048"
  - "4q2s2048"
  - "4q4s2048"
  - "4q8s2048"
  - "4q16s2048"
  - "4q32s2048"
  - "4q64s2048"
  - "4q128s2048"
  - "4q256s2048"
  - "4q512s2048"
  - "4q1024s2048"
  - "4q2048"
  # Batch size 8
  # seq_len = 32
  - "8q1s32"
  - "8q2s32"
  - "8q4s32"
  - "8q8s32"
  - "8q16s32"
  - "8q32"
  # seq_len = 64
  - "8q1s64"
  - "8q2s64"
  - "8q4s64"
  - "8q8s64"
  - "8q16s64"
  - "8q32s64"
  - "8q64"
  # seq_len = 128
  - "8q1s128"
  - "8q2s128"
  - "8q4s128"
  - "8q8s128"
  - "8q16s128"
  - "8q32s128"
  - "8q64s128"
  - "8q128"
  # seq_len = 256
  - "8q1s256"
  - "8q2s256"
  - "8q4s256"
  - "8q8s256"
  - "8q16s256"
  - "8q32s256"
  - "8q64s256"
  - "8q128s256"
  - "8q256"
  # seq_len = 512
  - "8q1s512"
  - "8q2s512"
  - "8q4s512"
  - "8q8s512"
  - "8q16s512"
  - "8q32s512"
  - "8q64s512"
  - "8q128s512"
  - "8q256s512"
  - "8q512"
  # seq_len = 1024
  - "8q1s1024"
  - "8q2s1024"
  - "8q4s1024"
  - "8q8s1024"
  - "8q16s1024"
  - "8q32s1024"
  - "8q64s1024"
  - "8q128s1024"
  - "8q256s1024"
  - "8q512s1024"
  - "8q1024"
  # seq_len = 2048
  - "8q1s2048"
  - "8q2s2048"
  - "8q4s2048"
  - "8q8s2048"
  - "8q16s2048"
  - "8q32s2048"
  - "8q64s2048"
  - "8q128s2048"
  - "8q256s2048"
  - "8q512s2048"
  - "8q1024s2048"
  - "8q2048"
  # Batch size 16
  # seq_len = 32
  - "16q1s32"
  - "16q2s32"
  - "16q4s32"
  - "16q8s32"
  - "16q16s32"
  - "16q32"
  # seq_len = 64
  - "16q1s64"
  - "16q2s64"
  - "16q4s64"
  - "16q8s64"
  - "16q16s64"
  - "16q32s64"
  - "16q64"
  # seq_len = 128
  - "16q1s128"
  - "16q2s128"
  - "16q4s128"
  - "16q8s128"
  - "16q16s128"
  - "16q32s128"
  - "16q64s128"
  - "16q128"
  # seq_len = 256
  - "16q1s256"
  - "16q2s256"
  - "16q4s256"
  - "16q8s256"
  - "16q16s256"
  - "16q32s256"
  - "16q64s256"
  - "16q128s256"
  - "16q256"
  # seq_len = 512
  - "16q1s512"
  - "16q2s512"
  - "16q4s512"
  - "16q8s512"
  - "16q16s512"
  - "16q32s512"
  - "16q64s512"
  - "16q128s512"
  - "16q256s512"
  - "16q512"
  # seq_len = 1024
  - "16q1s1024"
  - "16q2s1024"
  - "16q4s1024"
  - "16q8s1024"
  - "16q16s1024"
  - "16q32s1024"
  - "16q64s1024"
  - "16q128s1024"
  - "16q256s1024"
  - "16q512s1024"
  - "16q1024"
  # seq_len = 2048
  - "16q1s2048"
  - "16q2s2048"
  - "16q4s2048"
  - "16q8s2048"
  - "16q16s2048"
  - "16q32s2048"
  - "16q64s2048"
  - "16q128s2048"
  - "16q256s2048"
  - "16q512s2048"
  - "16q1024s2048"
  - "16q2048"
  # Batch size 32
  # seq_len = 32
  - "32q1s32"
  - "32q2s32"
  - "32q4s32"
  - "32q8s32"
  - "32q16s32"
  - "32q32"
  # seq_len = 64
  - "32q1s64"
  - "32q2s64"
  - "32q4s64"
  - "32q8s64"
  - "32q16s64"
  - "32q32s64"
  - "32q64"
  # seq_len = 128
  - "32q1s128"
  - "32q2s128"
  - "32q4s128"
  - "32q8s128"
  - "32q16s128"
  - "32q32s128"
  - "32q64s128"
  - "32q128"
  # seq_len = 256
  - "32q1s256"
  - "32q2s256"
  - "32q4s256"
  - "32q8s256"
  - "32q16s256"
  - "32q32s256"
  - "32q64s256"
  - "32q128s256"
  - "32q256"
  # seq_len = 512
  - "32q1s512"
  - "32q2s512"
  - "32q4s512"
  - "32q8s512"
  - "32q16s512"
  - "32q32s512"
  - "32q64s512"
  - "32q128s512"
  - "32q256s512"
  - "32q512"
  # seq_len = 1024
  - "32q1s1024"
  - "32q2s1024"
  - "32q4s1024"
  - "32q8s1024"
  - "32q16s1024"
  - "32q32s1024"
  - "32q64s1024"
  - "32q128s1024"
  - "32q256s1024"
  - "32q512s1024"
  - "32q1024"
  # seq_len = 2048
  - "32q1s2048"
  - "32q2s2048"
  - "32q4s2048"
  - "32q8s2048"
  - "32q16s2048"
  - "32q32s2048"
  - "32q64s2048"
  - "32q128s2048"
  - "32q256s2048"
  - "32q512s2048"
  - "32q1024s2048"
  - "32q2048"
 backends:
  - FLASHMLA_SPARSE
 device: "cuda:0"
 profile_memory: false
 sparse_mla_dense_mha_max_seq_len: 2048
 sparse_mla_topk_pattern: "random"
 output:
  csv: "benchmark_output/mla_sparse_mha_vs_mqa.csv"
  json: "benchmark_output/mla_sparse_mha_vs_mqa.json"
@@ -9,8 +9,6 @@ needing full VllmConfig integration.
 """
 import statistics
 import tempfile
 from pathlib import Path
 import numpy as np
 import torch
@@ -19,6 +17,7 @@ from common import (
    BenchmarkResult,
    MockHfConfig,
    MockIndexer,
    MockKVBProj,
    MockLayer,
    run_do_bench,
    run_ncu_profile,
@@ -34,59 +33,8 @@ from vllm.config import (
    VllmConfig,
    set_current_vllm_config,
 )
 from vllm.model_executor.layers.linear import ColumnParallelLinear
 from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum
 def _safe_profile_name(value: str) -> str:
    return "".join(c if c.isalnum() or c in "._-" else "_" for c in value)
 def _create_kv_b_proj(
    mla_dims: dict,
    device: torch.device,
 ):
    kv_b_proj = ColumnParallelLinear(
        mla_dims["kv_lora_rank"],
        mla_dims["num_q_heads"]
        * (mla_dims["qk_nope_head_dim"] + mla_dims["v_head_dim"]),
        bias=False,
        params_dtype=torch.bfloat16,
        quant_config=None,
        prefix="benchmark.kv_b_proj",
    ).to(device)
    with torch.no_grad():
        kv_b_proj.weight.copy_(torch.randn_like(kv_b_proj.weight))
    return kv_b_proj
 def _ensure_single_rank_model_parallel() -> None:
    import torch.distributed as dist
    from vllm.distributed import (
        ensure_model_parallel_initialized,
        init_distributed_environment,
        model_parallel_is_initialized,
    )
    if not dist.is_available():
        return
    if not dist.is_initialized():
        with tempfile.NamedTemporaryFile(
            prefix="vllm_bench_dist_", delete=False
        ) as init_file:
            distributed_init_method = f"file://{init_file.name}"
        init_distributed_environment(
            world_size=1,
            rank=0,
            distributed_init_method=distributed_init_method,
            local_rank=0,
            backend="nccl",
        )
    if not model_parallel_is_initialized():
        ensure_model_parallel_initialized(1, 1)
 # ============================================================================
 # VllmConfig Creation
 # ============================================================================
@@ -118,12 +66,10 @@ def create_minimal_vllm_config(
    block_size: int = 128,
    max_num_seqs: int = 256,
    max_num_batched_tokens: int = 8192,
    max_model_len: int = 32768,
    mla_dims: dict | None = None,
    index_topk: int | None = None,
    prefill_backend: str | None = None,
    kv_cache_dtype: str = "auto",
    sparse_mla_force_mqa: bool = False,
 ) -> VllmConfig:
    """
    Create minimal VllmConfig for MLA benchmarks.
@@ -140,8 +86,6 @@ def create_minimal_vllm_config(
        prefill_backend: Prefill backend name (e.g., "fa3", "fa4", "flashinfer",
                        "trtllm"). Configures the attention config to force
                        the specified prefill backend.
        sparse_mla_force_mqa: If True, forces all sparse MLA tokens through
                    forward_mqa (even prefill tokens).
    Returns:
        VllmConfig for benchmarking
@@ -187,7 +131,7 @@ def create_minimal_vllm_config(
            trust_remote_code=True,
            dtype="bfloat16",
            seed=0,
-            max_model_len=max_model_len,
+            max_model_len=32768,
            quantization=None,
            enforce_eager=False,
            max_logprobs=20,
@@ -219,7 +163,7 @@ def create_minimal_vllm_config(
    scheduler_config = SchedulerConfig(
        max_num_seqs=max_num_seqs,
        max_num_batched_tokens=max(max_num_batched_tokens, max_num_seqs),
-        max_model_len=max_model_len,
+        max_model_len=32768,
        is_encoder_decoder=False,
        enable_chunked_prefill=True,
    )
@@ -248,9 +192,6 @@ def create_minimal_vllm_config(
                "flash_attn_version"
            ]
    if sparse_mla_force_mqa:
        vllm_config.attention_config.sparse_mla_force_mqa = True
    return vllm_config
@@ -607,7 +548,12 @@ def _create_backend_impl(
    # Calculate scale
    scale = 1.0 / np.sqrt(mla_dims["qk_nope_head_dim"] + mla_dims["qk_rope_head_dim"])
-    kv_b_proj = _create_kv_b_proj(mla_dims, device)
+    # Create mock kv_b_proj layer for prefill mode
    mock_kv_b_proj = MockKVBProj(
        num_heads=mla_dims["num_q_heads"],
        qk_nope_head_dim=mla_dims["qk_nope_head_dim"],
        v_head_dim=mla_dims["v_head_dim"],
    )
    # Create indexer for sparse backends
    indexer = None
@@ -638,7 +584,7 @@ def _create_backend_impl(
        "qk_rope_head_dim": mla_dims["qk_rope_head_dim"],
        "qk_head_dim": mla_dims["qk_nope_head_dim"] + mla_dims["qk_rope_head_dim"],
        "v_head_dim": mla_dims["v_head_dim"],
-        "kv_b_proj": kv_b_proj,
+        "kv_b_proj": mock_kv_b_proj,
    }
    # Add indexer for sparse backends
@@ -839,35 +785,14 @@ def _run_single_benchmark(
    # Fill indexer with random indices for sparse backends
    is_sparse = backend_cfg.get("is_sparse", False)
    if is_sparse and indexer is not None:
-        indexer.fill_indices(
+        indexer.fill_random_indices(total_q, max_kv_len)
            total_q,
            max_kv_len,
            getattr(config, "sparse_mla_topk_pattern", "random"),
        )
    # Determine which forward methods to use based on metadata.
-    # Non-sparse backends use .decode/.prefill sub-objects.
+    # Sparse MLA backends always use forward_mqa
-    # Sparse backends use num_decode_tokens/num_prefills directly.
+    has_decode = is_sparse or getattr(metadata, "decode", None) is not None
-    #
+    has_prefill = not is_sparse and getattr(metadata, "prefill", None) is not None
    # sparse_mla_force_mqa overrides: even for prefill metadata, use MQA.
    force_mqa = getattr(config, "sparse_mla_force_mqa", False)
    force_dense_mha = getattr(config, "sparse_mla_mha_mode", "auto") == "dense"
    if force_mqa:
        has_decode = True
        has_prefill = False
    elif is_sparse:
        has_decode = metadata.num_decode_tokens > 0
        has_prefill = metadata.num_prefills > 0
    else:
        has_decode = metadata.decode is not None
        has_prefill = metadata.prefill is not None
    if not has_decode and not has_prefill:
        raise RuntimeError("Metadata has neither decode nor prefill metadata")
    if is_sparse and force_dense_mha and not has_prefill:
        raise RuntimeError(
            "Sparse MLA dense_mha benchmark did not produce prefill metadata. "
            "Check reorder_batch_threshold/path forcing."
        )
    num_decode = (
        metadata.num_decode_tokens
@@ -946,6 +871,7 @@ def _run_single_benchmark(
                metadata,
                prefill_inputs["k_scale"],
                prefill_fp8_output if fused_output else prefill_inputs["output"],
                prefill_output_scale if fused_output else None,
            )
            if fused_output:
                out = prefill_fp8_output
@@ -972,48 +898,6 @@ def _run_single_benchmark(
            throughput_tokens_per_sec=0.0,
        )
    if config.torch_profile:
        profile_dir = Path(
            config.torch_profile_dir or "benchmark_outputs/torch_profiles"
        )
        profile_dir.mkdir(parents=True, exist_ok=True)
        trace_name = _safe_profile_name(f"{config.backend}_{config.batch_spec}")
        trace_path = profile_dir / f"{trace_name}.json"
        iters = max(config.torch_profile_iters, 1)
        forward_fn()
        torch.accelerator.synchronize()
        with torch.profiler.profile(
            activities=[
                torch.profiler.ProfilerActivity.CPU,
                torch.profiler.ProfilerActivity.CUDA,
            ],
            record_shapes=True,
            profile_memory=True,
            with_stack=False,
        ) as prof:
            for _ in range(iters):
                forward_fn()
                torch.accelerator.synchronize()
                prof.step()
        prof.export_chrome_trace(str(trace_path))
        print(f"Saved PyTorch profiler trace to {trace_path}")
        print(
            prof.key_averages().table(
                sort_by="cuda_time_total",
                row_limit=25,
            )
        )
        return BenchmarkResult(
            config=config,
            mean_time=0.0,
            median_time=0.0,
            std_time=0.0,
            min_time=0.0,
            max_time=0.0,
            throughput_tokens_per_sec=0.0,
        )
    all_ms = run_do_bench(benchmark_fn, config.use_cuda_graphs, config.warmup_ms)
    # Convert ms to seconds per layer
@@ -1036,7 +920,6 @@ def _run_mla_benchmark_batched(
    configs_with_params: list[tuple],  # [(config, threshold, num_splits), ...]
    index_topk: int = 2048,
    prefill_backend: str | None = None,
    sparse_mla_force_mqa: bool = False,
    output_scale: float | None = None,
    fuse_quant_op: bool = False,
 ) -> list[BenchmarkResult]:
@@ -1057,8 +940,6 @@ def _run_mla_benchmark_batched(
        index_topk: Topk value for sparse MLA backends (default 2048)
        prefill_backend: Prefill backend name (e.g., "fa3", "fa4").
            When set, forces the specified FlashAttention version for prefill.
        sparse_mla_force_mqa: If True, forces all sparse MLA tokens through
            forward_mqa (even prefill tokens).
    Returns:
        List of BenchmarkResult objects
@@ -1099,41 +980,21 @@ def _run_mla_benchmark_batched(
        sum(r.q_len for r in parse_batch_spec(cfg.batch_spec))
        for cfg, *_ in configs_with_params
    )
    max_model_len = max(
        max_total_q,
        max(
            getattr(cfg, "max_model_len", None) or 32768
            for cfg, *_ in configs_with_params
        ),
    )
    # Create and set vLLM config for MLA (reused across all benchmarks)
    vllm_config = create_minimal_vllm_config(
        model_name="deepseek-v3",  # Used only for model path
        block_size=block_size,
        max_num_batched_tokens=max_total_q,
        max_model_len=max_model_len,
        mla_dims=mla_dims,  # Use custom dims from config or default
        index_topk=index_topk if is_sparse else None,
        prefill_backend=prefill_backend,
        kv_cache_dtype=kv_cache_dtype,
        sparse_mla_force_mqa=sparse_mla_force_mqa,
    )
    results = []
    # Initialize workspace manager (needed by metadata builders)
    from vllm.v1.worker.workspace import (
        init_workspace_manager,
        is_workspace_manager_initialized,
    )
    if not is_workspace_manager_initialized():
        init_workspace_manager(device)
    with set_current_vllm_config(vllm_config):
        _ensure_single_rank_model_parallel()
        # Create backend impl, layer, builder, and indexer (reused across benchmarks)
        impl, layer, builder_instance, indexer = _create_backend_impl(
            backend_cfg,
@@ -1179,20 +1040,9 @@ def _run_mla_benchmark_batched(
        for config, threshold, num_splits in configs_with_params:
            # Set threshold for this benchmark (FlashAttn/FlashMLA only)
            original_threshold = None
-            effective_threshold = threshold
+            if threshold is not None and builder_instance:
            force_dense_mha = (
                is_sparse
                and getattr(config, "sparse_mla_mha_mode", "auto") == "dense"
                and not getattr(config, "sparse_mla_force_mqa", False)
            )
            if force_dense_mha:
                # Sparse MLA normally treats q_len <= 1 as decode. Use an
                # impossible threshold so dense_mha benchmarks actually run
                # the prefill/MHA path, including q_len=1 short extends.
                effective_threshold = -1
            if effective_threshold is not None and builder_instance:
                original_threshold = builder_instance.reorder_batch_threshold
-                builder_instance.reorder_batch_threshold = effective_threshold
+                builder_instance.reorder_batch_threshold = threshold
            # Set num_splits for CUTLASS
            original_num_splits = None
@@ -1240,7 +1090,6 @@ def run_mla_benchmark(
    num_kv_splits: int | None = None,
    index_topk: int = 2048,
    prefill_backend: str | None = None,
    sparse_mla_force_mqa: bool = False,
    output_scale: float | None = None,
    fuse_quant_op: bool = False,
 ) -> BenchmarkResult | list[BenchmarkResult]:
@@ -1262,8 +1111,6 @@ def run_mla_benchmark(
        index_topk: Topk value for sparse MLA backends (default 2048)
        prefill_backend: Prefill backend name (e.g., "fa3", "fa4").
            When set, forces the specified FlashAttention version for prefill.
        sparse_mla_force_mqa: If True, forces all sparse MLA tokens through
            forward_mqa (even prefill tokens).
        output_scale: Static per-tensor FP8 scale for prefill output (None = bf16).
        fuse_quant_op: With output_scale set, fuse the FP8 write into the prefill
            kernel vs a standalone post-quant kernel. See _run_single_benchmark.
@@ -1295,7 +1142,6 @@ def run_mla_benchmark(
        configs_with_params,
        index_topk,
        prefill_backend=prefill_backend,
        sparse_mla_force_mqa=sparse_mla_force_mqa,
        output_scale=output_scale,
        fuse_quant_op=fuse_quant_op,
    )
@@ -69,11 +69,12 @@ def make_inputs(total_tokens, num_reqs, block_size):
    # Output workspace
    dst = torch.zeros(total_tokens, HEAD_DIM, dtype=torch.bfloat16, device="cuda")
    seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device="cuda")
    workspace_starts_t = torch.tensor(
        workspace_starts, dtype=torch.int32, device="cuda"
    )
-    return cache, dst, block_table, workspace_starts_t
+    return cache, dst, block_table, seq_lens_t, workspace_starts_t
 def bench_scenario(label, num_reqs, total_tokens_list, save_path):
@@ -93,7 +94,7 @@ def bench_scenario(label, num_reqs, total_tokens_list, save_path):
        )
    )
    def bench_fn(total_tokens, provider, num_reqs):
-        cache, dst, block_table, ws_starts = make_inputs(
+        cache, dst, block_table, seq_lens_t, ws_starts = make_inputs(
            total_tokens, num_reqs, BLOCK_SIZE
        )
@@ -101,7 +102,7 @@ def bench_scenario(label, num_reqs, total_tokens_list, save_path):
        ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
            lambda: ops.cp_gather_and_upconvert_fp8_kv_cache(
-                cache, dst, block_table, ws_starts, num_reqs
+                cache, dst, block_table, seq_lens_t, ws_starts, num_reqs
            ),
            quantiles=quantiles,
            rep=500,
@@ -1,176 +0,0 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import statistics
 import torch
 from tabulate import tabulate
 from vllm.models.inkling.nvidia.ops import qkvr_prep
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 def make_inputs(tokens: int, tp_size: int, is_local: bool):
    torch.manual_seed(0)
    num_q_heads = 64 // tp_size
    num_kv_heads = (16 if is_local else 8) // tp_size
    head_dim = 128
    d_rel = 16
    rel_extent = 512 if is_local else 1024
    page_size = 16
    num_blocks = (tokens + page_size - 1) // page_size
    q_width = num_q_heads * head_dim
    kv_width = num_kv_heads * head_dim
    r_width = num_q_heads * d_rel
    device = "cuda"
    qkvr = torch.randn(
        tokens,
        q_width + 2 * kv_width + r_width,
        device=device,
        dtype=torch.bfloat16,
    )
    k_weight = torch.randn(kv_width, 4, device=device, dtype=torch.bfloat16)
    v_weight = torch.randn_like(k_weight)
    q_norm_weight = torch.randn(head_dim, device=device, dtype=torch.bfloat16)
    k_norm_weight = torch.randn_like(q_norm_weight)
    rel_proj = torch.randn(d_rel, rel_extent, device=device, dtype=torch.bfloat16)
    conv_cache = torch.zeros(
        num_blocks,
        num_kv_heads,
        page_size,
        2 * head_dim,
        device=device,
        dtype=torch.bfloat16,
    )
    key_cache = torch.empty(
        num_blocks,
        page_size,
        num_kv_heads,
        head_dim,
        device=device,
        dtype=torch.bfloat16,
    )
    value_cache = torch.empty_like(key_cache)
    positions = torch.arange(tokens, device=device, dtype=torch.int64)
    block_table = torch.arange(num_blocks, device=device, dtype=torch.int32)[None]
    seq_idx = torch.zeros(tokens, device=device, dtype=torch.int32)
    slots = torch.arange(tokens, device=device, dtype=torch.int64)
    query_start = torch.zeros(tokens, device=device, dtype=torch.int32)
    log_scaling = None
    if not is_local:
        effective_n = (positions + 1).to(torch.float32)
        log_scaling = 1.0 + 0.1 * torch.log(torch.clamp(effective_n / 128000, min=1.0))
    return (
        qkvr,
        k_weight,
        v_weight,
        q_norm_weight,
        k_norm_weight,
        rel_proj,
        1e-6,
        num_q_heads,
        num_kv_heads,
        head_dim,
        d_rel,
        conv_cache,
        key_cache,
        value_cache,
        positions,
        block_table,
        seq_idx,
        slots,
        query_start,
        slots,
        0,
        head_dim,
        page_size,
        log_scaling,
    )
 def capture(implementation, inputs):
    outputs = []
    def run():
        outputs[:] = implementation.fused_qkvr_prep(*inputs)
    stream = torch.cuda.Stream()
    stream.wait_stream(torch.cuda.current_stream())
    with torch.cuda.stream(stream):
        for _ in range(3):
            run()
    torch.cuda.current_stream().wait_stream(stream)
    torch.accelerator.synchronize()
    graph = torch.cuda.CUDAGraph()
    with torch.cuda.graph(graph):
        run()
    torch.accelerator.synchronize()
    return graph, outputs
 def time_graph(graph: torch.cuda.CUDAGraph, warmup: int, repeats: int) -> float:
    for _ in range(warmup):
        graph.replay()
    torch.accelerator.synchronize()
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)
    start.record()
    for _ in range(repeats):
        graph.replay()
    end.record()
    end.synchronize()
    return start.elapsed_time(end) * 1000 / repeats
 def benchmark(inputs, args) -> float:
    graph, _ = capture(qkvr_prep, inputs)
    return statistics.median(
        time_graph(graph, args.warmup, args.repeats) for _ in range(args.trials)
    )
@torch.inference_mode()
 def main(args):
    rows = []
    for tp_size in args.tp_sizes:
        for tokens in args.tokens:
            for is_local in (True, False):
                triton_us = benchmark(make_inputs(tokens, tp_size, is_local), args)
                rows.append(
                    [
                        tp_size,
                        tokens,
                        "local" if is_local else "global",
                        triton_us,
                    ]
                )
    print("Inkling QKVR prep (CUDA graph, median latency)")
    print(
        tabulate(
            rows,
            headers=[
                "TP",
                "tokens",
                "scope",
                "Triton (us)",
            ],
            floatfmt=("d", "d", "", ".2f"),
        )
    )
 if __name__ == "__main__":
    parser = FlexibleArgumentParser()
    parser.add_argument(
        "--tokens",
        type=int,
        nargs="+",
        default=[1 << power for power in range(15)],
    )
    parser.add_argument("--tp-sizes", type=int, nargs="+", default=[4, 8])
    parser.add_argument("--warmup", type=int, default=20)
    parser.add_argument("--repeats", type=int, default=200)
    parser.add_argument("--trials", type=int, default=5)
    main(parser.parse_args())
@@ -1,367 +0,0 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Benchmark the Kimi-K3 latent MoE addmm against CuTe residual GEMM.
 The benchmark covers ``BF16[M, 3584] @ BF16[7168, 3584].T + BF16[M, 7168]``
 with FP32 accumulation and BF16 output. Both backends execute through CUDA
 Graph replay. Weights and residuals rotate across buffers exceeding L2 so the
 comparison models the full latent MoE projection-and-add path.
 """
 from __future__ import annotations
 import argparse
 import dataclasses
 import importlib.util
 import json
 import math
 import statistics
 from collections.abc import Callable, Sequence
 from pathlib import Path
 from typing import Any
 import cutlass
 import cutlass.cute as cute
 import torch
 from cuda.bindings import driver as cuda
 from cuda.bindings.driver import CUstream
 from quack.compile_utils import make_fake_tensor
 N = 7168
 K = 3584
@dataclasses.dataclass(frozen=True, slots=True)
 class Config:
    block_size: int
    outputs_per_block: int
    k_unroll: int
    vector_width: int = 8
 def parse_config(value: str) -> Config:
    try:
        parts = [int(part) for part in value.split(",")]
    except ValueError as error:
        raise argparse.ArgumentTypeError(
            "config must be BLOCK,OUTPUTS,K_UNROLL[,VECTOR_WIDTH]"
        ) from error
    if len(parts) == 3:
        return Config(*parts)
    if len(parts) == 4:
        return Config(*parts)
    raise argparse.ArgumentTypeError(
        "config must be BLOCK,OUTPUTS,K_UNROLL[,VECTOR_WIDTH]"
    )
 def production_residual_config(m: int) -> Config | None:
    """The measured Latent-MoE residual config for M, from the K3 table."""
    from vllm.models.kimi_k3.nvidia.low_latency_gemm import KIMI_K3_PROJECTIONS
    spec = KIMI_K3_PROJECTIONS.get((N, K))
    config = spec.residual_config(m) if spec is not None else None
    if config is None:
        return None
    return Config(
        config.block_size,
        config.outputs_per_block,
        config.k_unroll,
        config.vector_width,
    )
 def candidate_configs(mode: str, selected: Config | None, m: int) -> list[Config]:
    if mode == "selected":
        if selected is not None:
            return [selected]
        # No explicit --config: fall back to the production table for this M.
        config = production_residual_config(m)
        return [config] if config is not None else []
    if mode == "baseline":
        return [Config(224, 4, 2)]
    return [
        Config(block_size, outputs_per_block, k_unroll, vector_width)
        for vector_width in (4, 8)
        for block_size in (32, 64, 128, 224, 448)
        if block_size % 32 == 0 and K % (block_size * vector_width) == 0
        for outputs_per_block in (1, 2, 4, 7, 8)
        if N % outputs_per_block == 0
        for k_unroll in (1, 2, 4)
    ]
 def load_kernel_class(path: Path):
    spec = importlib.util.spec_from_file_location("cute_skinny_device", path)
    if spec is None or spec.loader is None:
        raise RuntimeError(f"cannot load CuTe kernel from {path}")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module.CuteSkinnyGemm
 def stream() -> CUstream:
    return CUstream(torch.cuda.current_stream().cuda_stream)
 def compile_kernel(kernel_class, m: int, config: Config, max_registers: int):
    element_type = cutlass.BFloat16
    n = cute.sym_int(divisibility=config.outputs_per_block)
    k = cute.sym_int(divisibility=config.block_size * config.vector_width)
    a = make_fake_tensor(element_type, (m, k), divisibility=config.vector_width)
    b = make_fake_tensor(element_type, (n, k), divisibility=config.vector_width)
    residual = make_fake_tensor(element_type, (m, n), divisibility=1)
    c = make_fake_tensor(element_type, (m, n), divisibility=1)
    kernel = kernel_class(
        element_type=element_type,
        num_rows=m,
        block_size=config.block_size,
        outputs_per_block=config.outputs_per_block,
        vector_width=config.vector_width,
        k_unroll=config.k_unroll,
        has_residual=True,
        use_pdl=True,
    )
    return cute.compile(
        kernel,
        a,
        b,
        residual,
        c,
        stream(),
        options=(
            "--enable-tvm-ffi --keep-cubin "
            f"--ptxas-options -maxrregcount={max_registers} "
            "--ptxas-options -lineinfo"
        ),
    )
 def resource_usage(compiled) -> dict[str, Any]:
    executor = getattr(compiled, "_default_executor", None)
    context = getattr(executor, "exec_context", None)
    functions = getattr(context, "kernel_functions", None)
    if not functions:
        return {"resource_metrics_available": False}
    def attribute(name, function) -> int:
        error, value = cuda.cuFuncGetAttribute(name, function)
        if error != cuda.CUresult.CUDA_SUCCESS:
            raise RuntimeError(f"cuFuncGetAttribute failed with {error}")
        return int(value)
    registers = [
        attribute(cuda.CUfunction_attribute.CU_FUNC_ATTRIBUTE_NUM_REGS, function)
        for function in functions
    ]
    local_bytes = [
        attribute(
            cuda.CUfunction_attribute.CU_FUNC_ATTRIBUTE_LOCAL_SIZE_BYTES,
            function,
        )
        for function in functions
    ]
    return {
        "resource_metrics_available": True,
        "registers_per_thread": max(registers, default=0),
        "spill_bytes": max(local_bytes, default=0),
    }
 def rotating_buffer_count(m: int, multiplier: float, limit: int) -> int:
    properties = torch.cuda.get_device_properties(0)
    bytes_per_pair = (N * K + m * N) * 2
    target = math.ceil(multiplier * properties.L2_cache_size)
    return max(2, min(limit, math.ceil(target / bytes_per_pair)))
 def graph_samples(
    launch: Callable[[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], None],
    activation: torch.Tensor,
    weights: Sequence[torch.Tensor],
    residuals: Sequence[torch.Tensor],
    repeats: int,
    replays: int,
 ) -> tuple[list[float], list[torch.Tensor]]:
    outputs = [torch.empty_like(residual) for residual in residuals]
    for weight, residual, output in zip(weights, residuals, outputs):
        launch(activation, weight, residual, output)
    torch.accelerator.synchronize()
    graph = torch.cuda.CUDAGraph()
    with torch.cuda.graph(graph):
        for weight, residual, output in zip(weights, residuals, outputs):
            launch(activation, weight, residual, output)
    for _ in range(20):
        graph.replay()
    torch.accelerator.synchronize()
    samples = []
    for _ in range(repeats):
        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)
        start.record()
        for _ in range(replays):
            graph.replay()
        end.record()
        end.synchronize()
        samples.append(start.elapsed_time(end) * 1000.0 / (replays * len(weights)))
    return samples, outputs
 def summarize(samples: Sequence[float]) -> dict[str, Any]:
    ordered = sorted(samples)
    def percentile(fraction: float) -> float:
        position = fraction * (len(ordered) - 1)
        lower = math.floor(position)
        upper = math.ceil(position)
        if lower == upper:
            return ordered[lower]
        weight = position - lower
        return ordered[lower] * (1.0 - weight) + ordered[upper] * weight
    mean = statistics.mean(samples)
    return {
        "median_us": statistics.median(samples),
        "p10_us": percentile(0.1),
        "p90_us": percentile(0.9),
        "mean_us": mean,
        "cv_pct": statistics.pstdev(samples) / mean * 100.0,
        "samples_us": list(samples),
    }
 def correctness(
    output: torch.Tensor,
    activation: torch.Tensor,
    weight: torch.Tensor,
    residual: torch.Tensor,
 ) -> dict[str, Any]:
    actual = output.float()
    reference = activation.float() @ weight.float().t() + residual.float()
    error = (actual - reference).abs()
    scaled_error = error / (reference.abs() + 1.0)
    cosine = torch.nn.functional.cosine_similarity(
        actual.flatten(), reference.flatten(), dim=0
    ).item()
    return {
        "valid": cosine > 0.999,
        "cosine": cosine,
        "max_abs_error": error.max().item(),
        "max_scaled_error": scaled_error.max().item(),
    }
 def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--kernel", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument(
        "--mode", choices=("baseline", "sweep", "selected"), default="baseline"
    )
    parser.add_argument("--config", type=parse_config)
    parser.add_argument("--m", type=int, action="append")
    parser.add_argument("--config-shard", type=int, default=0)
    parser.add_argument("--num-config-shards", type=int, default=1)
    parser.add_argument("--repeats", type=int, default=21)
    parser.add_argument("--replays", type=int, default=200)
    parser.add_argument("--cache-multiplier", type=float, default=3.0)
    parser.add_argument("--max-buffers", type=int, default=32)
    parser.add_argument("--max-registers", type=int, default=64)
    args = parser.parse_args()
    token_counts = args.m or list(range(1, 17))
    if any(not 1 <= m <= 16 for m in token_counts):
        raise ValueError("expected 1 <= M <= 16")
    if not 0 <= args.config_shard < args.num_config_shards:
        raise ValueError("config shard must be in [0, num_config_shards)")
    torch.accelerator.set_device_index(0)
    if torch.cuda.get_device_capability() != (10, 3):
        raise RuntimeError("this benchmark requires SM103")
    kernel_class = load_kernel_class(args.kernel)
    properties = torch.cuda.get_device_properties(0)
    metadata = {
        "device": properties.name,
        "compute_capability": list(torch.cuda.get_device_capability()),
        "torch_version": torch.__version__,
        "cuda_version": torch.version.cuda,
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    with args.output.open("w", encoding="utf-8") as output_file:
        for m in token_counts:
            configs = candidate_configs(args.mode, args.config, m)
            torch.manual_seed(20260722 + m)
            count = rotating_buffer_count(m, args.cache_multiplier, args.max_buffers)
            activation = torch.randn((m, K), device="cuda", dtype=torch.bfloat16)
            weights = [
                torch.randn((N, K), device="cuda", dtype=torch.bfloat16)
                for _ in range(count)
            ]
            residuals = [
                torch.randn((m, N), device="cuda", dtype=torch.bfloat16)
                for _ in range(count)
            ]
            candidates: list[tuple[str, Config | None]] = [("cublas_addmm", None)]
            candidates.extend(
                ("cute_residual", config)
                for index, config in enumerate(configs)
                if index % args.num_config_shards == args.config_shard
            )
            for backend, config in candidates:
                row: dict[str, Any] = {
                    "m": m,
                    "n": N,
                    "k": K,
                    "backend": backend,
                    "mode": args.mode,
                    "config": dataclasses.asdict(config) if config else {},
                    "num_buffers": count,
                    "cache_multiplier": args.cache_multiplier,
                    **metadata,
                }
                try:
                    if backend == "cublas_addmm":
                        launch = lambda a, b, residual, c: torch.addmm(
                            residual, a, b.t(), out=c
                        )
                    else:
                        if config is None:
                            raise AssertionError("missing CuTe config")
                        compiled = compile_kernel(
                            kernel_class, m, config, args.max_registers
                        )
                        launch = lambda a, b, residual, c, fn=compiled: fn(
                            a, b, residual, c, stream()
                        )
                        row.update(resource_usage(compiled))
                    samples, outputs = graph_samples(
                        launch,
                        activation,
                        weights,
                        residuals,
                        args.repeats,
                        args.replays,
                    )
                    row.update(
                        correctness(outputs[0], activation, weights[0], residuals[0])
                    )
                    row.update(summarize(samples))
                except Exception as error:  # noqa: BLE001
                    row.update(
                        {
                            "valid": False,
                            "error": f"{type(error).__name__}: {error}",
                        }
                    )
                output_file.write(json.dumps(row, sort_keys=True) + "\n")
                output_file.flush()
                print(json.dumps(row, sort_keys=True), flush=True)
            del activation, weights, residuals
            torch.accelerator.empty_cache()
 if __name__ == "__main__":
    main()
@@ -1,806 +0,0 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Benchmark the Kimi K3 latent-MoE tail and its up-projection kernels.
 The ``up-projection`` subcommand isolates the TP-local dynamic and static-M
 skinny GEMMs. It rotates weights through a working set larger than L2 to model
 successive model layers.
 The ``whole-tail`` subcommand measures the distributed operator. Its reference
 path includes two AllReduces, RMSNorm, the replicated up-projection, and the
 final add. CUDA-event samples report the slowest rank so cross-rank skew is
 included.
 Examples:
 .. code-block:: console
    .venv/bin/python \
      benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py up-projection
    torchrun --nproc-per-node=8 \
      benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py whole-tail
 For multi-node runs, launch one ``torchrun`` agent per node and use a shared
 rendezvous endpoint.
 """
 from __future__ import annotations
 import argparse
 import json
 import math
 import os
 import statistics
 from collections.abc import Callable, Sequence
 from dataclasses import asdict
 from pathlib import Path
 from typing import Any
 import cutlass
 import cutlass.utils as utils
 import torch
 import torch.distributed as dist
 import torch.nn.functional as F
 from cuda.bindings import driver as cuda
 from vllm.distributed import get_tp_group
 from vllm.distributed.parallel_state import (
    init_distributed_environment,
    initialize_model_parallel,
    set_custom_all_reduce,
 )
 from vllm.model_executor.warmup.cutedsl_warmup import cutedsl_warmup
 from vllm.models.kimi_k3.nvidia.ops import latent_moe_tail
 from vllm.models.kimi_k3.nvidia.ops.cute_dsl.latent_moe_tail import (
    fused_add_multicast_gemm,
    fused_add_multicast_skinny_gemm,
 )
 HIDDEN_SIZE = 7168
 LATENT_SIZE = 3584
 RMS_EPS = 0.1
 MAX_NUM_TOKENS = 16
 MMA_TILER_MN = (64, 32)
 CLUSTER_SHAPE_MN = (1, 8)
 B_PRIME_STAGES = 2
 def parse_up_projection_config(
    value: str,
 ) -> fused_add_multicast_skinny_gemm.SkinnyConfig:
    try:
        values = [int(part) for part in value.split(",")]
    except ValueError as error:
        raise argparse.ArgumentTypeError(
            "config must be BLOCK,OUTPUTS,K_UNROLL[,VECTOR_WIDTH[,PREFETCH_B]]"
        ) from error
    if len(values) in (3, 4):
        return fused_add_multicast_skinny_gemm.SkinnyConfig(*values)
    if len(values) == 5 and values[4] in (0, 1):
        return fused_add_multicast_skinny_gemm.SkinnyConfig(
            *values[:4],
            prefetch_b_before_pdl=bool(values[4]),
        )
    raise argparse.ArgumentTypeError(
        "config must be BLOCK,OUTPUTS,K_UNROLL"
        "[,VECTOR_WIDTH[,PREFETCH_B]], where PREFETCH_B is 0 or 1"
    )
 def parse_tail_skinny_config(
    value: str,
 ) -> tuple[int, fused_add_multicast_skinny_gemm.SkinnyConfig]:
    try:
        values = [int(part) for part in value.split(",")]
    except ValueError as error:
        raise argparse.ArgumentTypeError(
            "config must be M,BLOCK,OUTPUTS,K_UNROLL[,VECTOR_WIDTH[,PREFETCH_B]]"
        ) from error
    if len(values) == 4:
        num_tokens, *config = values
        return num_tokens, fused_add_multicast_skinny_gemm.SkinnyConfig(*config)
    if len(values) == 5:
        num_tokens, *config = values
        return num_tokens, fused_add_multicast_skinny_gemm.SkinnyConfig(*config)
    if len(values) == 6 and values[5] in (0, 1):
        num_tokens, block, outputs, unroll, vector_width, prefetch = values
        return num_tokens, fused_add_multicast_skinny_gemm.SkinnyConfig(
            block,
            outputs,
            unroll,
            vector_width,
            bool(prefetch),
        )
    raise argparse.ArgumentTypeError(
        "config must be M,BLOCK,OUTPUTS,K_UNROLL"
        "[,VECTOR_WIDTH[,PREFETCH_B]], where PREFETCH_B is 0 or 1"
    )
 def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    subparsers = parser.add_subparsers(dest="scope", required=True)
    up_projection = subparsers.add_parser(
        "up-projection",
        help="Benchmark the isolated TP-local up-projection kernels.",
    )
    up_projection.add_argument(
        "--backend",
        choices=("dynamic", "skinny", "both"),
        default="both",
    )
    up_projection.add_argument("--tp-size", type=int, default=16)
    up_projection.add_argument(
        "--num-tokens",
        type=int,
        nargs="+",
        default=[*range(1, 9), 16],
    )
    up_projection.add_argument(
        "--skinny-config",
        type=parse_up_projection_config,
        action="append",
        help="Benchmark a static-M config for every selected token count.",
    )
    up_projection.add_argument("--cache-multiplier", type=float, default=2.0)
    up_projection.add_argument("--max-weights", type=int, default=64)
    up_projection.add_argument("--warmup-replays", type=int, default=10)
    up_projection.add_argument("--samples", type=int, default=31)
    up_projection.add_argument("--output", type=Path)
    whole_tail = subparsers.add_parser(
        "whole-tail",
        help="Benchmark the distributed latent-MoE tail operator.",
    )
    whole_tail.add_argument(
        "--backend",
        choices=("reference", "fused", "both"),
        default="both",
    )
    whole_tail.add_argument(
        "--num-tokens",
        type=int,
        nargs="+",
        default=[1, 5, 8, 16],
    )
    whole_tail.add_argument("--warmup-replays", type=int, default=20)
    whole_tail.add_argument("--samples", type=int, default=51)
    whole_tail.add_argument(
        "--skinny-max-num-tokens",
        type=int,
        nargs="+",
        help="Override the fused operator's static-M cutoff; use 0 for dynamic-only.",
    )
    whole_tail.add_argument(
        "--skinny-config",
        type=parse_tail_skinny_config,
        action="append",
        help="Override one static-M config for tuning.",
    )
    whole_tail.add_argument("--output", type=Path)
    return parser.parse_args()
 def percentile(samples: Sequence[float], fraction: float) -> float:
    ordered = sorted(samples)
    position = fraction * (len(ordered) - 1)
    lower = math.floor(position)
    upper = math.ceil(position)
    if lower == upper:
        return ordered[lower]
    upper_weight = position - lower
    return ordered[lower] * (1.0 - upper_weight) + ordered[upper] * upper_weight
 def summarize(samples_us: Sequence[float]) -> dict[str, Any]:
    mean_us = statistics.mean(samples_us)
    return {
        "median_us": statistics.median(samples_us),
        "p10_us": percentile(samples_us, 0.1),
        "p90_us": percentile(samples_us, 0.9),
        "mean_us": mean_us,
        "cv_pct": statistics.pstdev(samples_us) / mean_us * 100.0,
        "samples_us": list(samples_us),
    }
 def rotating_weight_count(
    shard_size: int,
    cache_multiplier: float,
    limit: int,
 ) -> int:
    properties = torch.cuda.get_device_properties(
        torch.accelerator.current_device_index()
    )
    weight_bytes = shard_size * LATENT_SIZE * 2
    target_bytes = math.ceil(properties.L2_cache_size * cache_multiplier)
    return max(2, min(limit, math.ceil(target_bytes / weight_bytes)))
 def capture_up_projection_graph(
    launches: Sequence[Callable[[], None]],
 ) -> torch.cuda.CUDAGraph:
    for launch in launches:
        launch()
    torch.accelerator.synchronize()
    graph = torch.cuda.CUDAGraph()
    with torch.cuda.graph(graph):
        for launch in launches:
            launch()
    torch.accelerator.synchronize()
    return graph
 def benchmark_up_projection_graph(
    graph: torch.cuda.CUDAGraph,
    *,
    operations_per_replay: int,
    warmup_replays: int,
    samples: int,
 ) -> dict[str, Any]:
    for _ in range(warmup_replays):
        graph.replay()
    torch.accelerator.synchronize()
    samples_us = []
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)
    for _ in range(samples):
        start.record()
        graph.replay()
        end.record()
        end.synchronize()
        samples_us.append(start.elapsed_time(end) * 1000.0 / operations_per_replay)
    return summarize(samples_us)
 class DynamicKernel:
    def __init__(
        self,
        shard_size: int,
        mailbox: torch.Tensor,
        shared_shard: torch.Tensor,
    ) -> None:
        self.shard_size = shard_size
        self.mailbox = mailbox
        self.mailbox_c = fused_add_multicast_gemm._as_cute(mailbox)
        compile_latent = torch.empty(
            (1, MAX_NUM_TOKENS, LATENT_SIZE),
            dtype=torch.bfloat16,
            device=mailbox.device,
        )
        compile_weight = torch.empty(
            (1, shard_size, LATENT_SIZE),
            dtype=torch.bfloat16,
            device=mailbox.device,
        )
        cluster_size = math.prod(CLUSTER_SHAPE_MN)
        max_active_clusters = utils.HardwareInfo().get_max_active_clusters(cluster_size)
        self.compiled = fused_add_multicast_gemm.compile_kernel(
            (MAX_NUM_TOKENS, shard_size, LATENT_SIZE, 1),
            fused_add_multicast_gemm._as_cute(
                compile_latent,
                dynamic_m=True,
            ),
            fused_add_multicast_gemm._as_cute(compile_weight),
            self.mailbox_c,
            fused_add_multicast_gemm._as_cute(shared_shard),
            HIDDEN_SIZE,
            shard_size,
            MMA_TILER_MN,
            CLUSTER_SHAPE_MN,
            max_active_clusters,
            B_PRIME_STAGES,
        )
    def launch(
        self,
        latent: torch.Tensor,
        weight: torch.Tensor,
        shared_shard: torch.Tensor,
    ) -> None:
        stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
        self.compiled(
            fused_add_multicast_gemm._as_cute(
                latent.unsqueeze(0),
                dynamic_m=True,
            ),
            fused_add_multicast_gemm._as_cute(weight.unsqueeze(0)),
            self.mailbox_c,
            fused_add_multicast_gemm._as_cute(shared_shard),
            cutlass.Int64(latent.shape[0]),
            cutlass.Int64(self.mailbox.data_ptr()),
            stream,
        )
 class SkinnyKernel:
    def __init__(
        self,
        num_tokens: int,
        shard_size: int,
        config: fused_add_multicast_skinny_gemm.SkinnyConfig,
    ) -> None:
        self.compiled = fused_add_multicast_skinny_gemm.compile_kernel(
            num_rows=num_tokens,
            latent_dim=LATENT_SIZE,
            hidden_dim=HIDDEN_SIZE,
            shard_dim=shard_size,
            config=config,
        )
    def launch(
        self,
        latent: torch.Tensor,
        weight: torch.Tensor,
        shared_shard: torch.Tensor,
        mailbox: torch.Tensor,
    ) -> None:
        self.compiled(
            fused_add_multicast_skinny_gemm._as_cute(latent),
            fused_add_multicast_skinny_gemm._as_cute(weight),
            fused_add_multicast_skinny_gemm._as_cute(shared_shard),
            cutlass.Int64(mailbox.data_ptr()),
            cuda.CUstream(torch.cuda.current_stream().cuda_stream),
        )
 def check_up_projection_output(
    actual: torch.Tensor,
    latent: torch.Tensor,
    weight: torch.Tensor,
    shared_shard: torch.Tensor,
 ) -> None:
    gemm = F.linear(latent.float(), weight.float()).to(torch.bfloat16)
    expected = (gemm.float() + shared_shard.float()).to(torch.bfloat16)
    torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2)
 def make_up_projection_launches(
    launch: Callable[[torch.Tensor, torch.Tensor, torch.Tensor], None],
    latent: torch.Tensor,
    weights: Sequence[torch.Tensor],
    shared_shard: torch.Tensor,
 ) -> list[Callable[[], None]]:
    return [
        lambda weight=weight: launch(latent, weight, shared_shard) for weight in weights
    ]
 def benchmark_up_projection(args: argparse.Namespace) -> None:
    if args.tp_size <= 0 or HIDDEN_SIZE % args.tp_size:
        raise ValueError("TP size must be positive and divide the hidden size")
    if any(not 1 <= num_tokens <= MAX_NUM_TOKENS for num_tokens in args.num_tokens):
        raise ValueError("--num-tokens values must be in [1, 16]")
    if args.cache_multiplier <= 0 or args.max_weights <= 0:
        raise ValueError("cache multiplier and max weights must be positive")
    if args.warmup_replays < 0 or args.samples <= 0:
        raise ValueError("warmup replays must be nonnegative and samples positive")
    torch.accelerator.set_device_index(0)
    device = torch.device("cuda", 0)
    if torch.cuda.get_device_capability(device)[0] != 10:
        raise RuntimeError("Kimi K3 latent-MoE tail requires SM100")
    shard_size = HIDDEN_SIZE // args.tp_size
    weight_count = rotating_weight_count(
        shard_size,
        args.cache_multiplier,
        args.max_weights,
    )
    torch.manual_seed(20260726)
    weights = [
        torch.randn(
            (shard_size, LATENT_SIZE),
            dtype=torch.bfloat16,
            device=device,
        )
        / LATENT_SIZE**0.5
        for _ in range(weight_count)
    ]
    mailbox = torch.empty(
        (1, MAX_NUM_TOKENS, HIDDEN_SIZE),
        dtype=torch.bfloat16,
        device=device,
    )
    shared = torch.randn(
        (MAX_NUM_TOKENS, HIDDEN_SIZE),
        dtype=torch.bfloat16,
        device=device,
    )
    shared_shard = shared[:, :shard_size]
    use_dynamic = args.backend in ("dynamic", "both")
    use_skinny = args.backend in ("skinny", "both")
    dynamic_kernel = (
        DynamicKernel(shard_size, mailbox, shared_shard) if use_dynamic else None
    )
    results = []
    for num_tokens in args.num_tokens:
        latent = torch.randn(
            (num_tokens, LATENT_SIZE),
            dtype=torch.bfloat16,
            device=device,
        )
        result: dict[str, Any] = {"num_tokens": num_tokens}
        if dynamic_kernel is not None:
            launches = make_up_projection_launches(
                dynamic_kernel.launch,
                latent,
                weights,
                shared_shard,
            )
            graph = capture_up_projection_graph(launches)
            result["dynamic"] = benchmark_up_projection_graph(
                graph,
                operations_per_replay=len(launches),
                warmup_replays=args.warmup_replays,
                samples=args.samples,
            )
            check_up_projection_output(
                mailbox[0, :num_tokens, :shard_size],
                latent,
                weights[-1],
                shared_shard[:num_tokens],
            )
        if use_skinny:
            configs = args.skinny_config or [
                fused_add_multicast_skinny_gemm.config_for_m(
                    num_tokens,
                    shard_size,
                )
            ]
            skinny_results = []
            for config in configs:
                skinny_kernel = SkinnyKernel(num_tokens, shard_size, config)
                def launch_skinny(
                    latent: torch.Tensor,
                    weight: torch.Tensor,
                    shared_shard: torch.Tensor,
                    *,
                    skinny_kernel: SkinnyKernel = skinny_kernel,
                    num_tokens: int = num_tokens,
                ) -> None:
                    skinny_kernel.launch(
                        latent,
                        weight,
                        shared_shard[:num_tokens],
                        mailbox,
                    )
                launches = make_up_projection_launches(
                    launch_skinny,
                    latent,
                    weights,
                    shared_shard,
                )
                graph = capture_up_projection_graph(launches)
                timing = benchmark_up_projection_graph(
                    graph,
                    operations_per_replay=len(launches),
                    warmup_replays=args.warmup_replays,
                    samples=args.samples,
                )
                check_up_projection_output(
                    mailbox[0, :num_tokens, :shard_size],
                    latent,
                    weights[-1],
                    shared_shard[:num_tokens],
                )
                skinny_results.append(
                    {
                        "config": asdict(config),
                        **timing,
                    }
                )
            result["skinny"] = skinny_results
        results.append(result)
    properties = torch.cuda.get_device_properties(device)
    report = {
        "scope": "up-projection",
        "device": properties.name,
        "compute_capability": list(torch.cuda.get_device_capability(device)),
        "tp_size": args.tp_size,
        "shard_size": shard_size,
        "weight_count": weight_count,
        "cache_multiplier": args.cache_multiplier,
        "warmup_replays": args.warmup_replays,
        "samples": args.samples,
        "results": results,
    }
    rendered = json.dumps(report, indent=2)
    print(rendered, flush=True)
    if args.output is not None:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(rendered + "\n", encoding="utf-8")
 def capture_tail_graph(
    operation: Callable[[], torch.Tensor],
    cpu_group: dist.ProcessGroup,
 ) -> tuple[torch.cuda.CUDAGraph, torch.Tensor]:
    for _ in range(3):
        dist.barrier(group=cpu_group)
        output = operation()
    torch.accelerator.synchronize()
    dist.barrier(group=cpu_group)
    graph = torch.cuda.CUDAGraph()
    with torch.cuda.graph(graph):
        output = operation()
    torch.accelerator.synchronize()
    return graph, output
 def benchmark_tail_graph(
    graph: torch.cuda.CUDAGraph,
    *,
    warmup_replays: int,
    samples: int,
    device_group: dist.ProcessGroup,
    cpu_group: dist.ProcessGroup,
 ) -> dict[str, Any]:
    for _ in range(warmup_replays):
        graph.replay()
    torch.accelerator.synchronize()
    dist.barrier(group=cpu_group)
    starts = [torch.cuda.Event(enable_timing=True) for _ in range(samples + 1)]
    ends = [torch.cuda.Event(enable_timing=True) for _ in range(samples + 1)]
    for start, end in zip(starts, ends):
        start.record()
        graph.replay()
        end.record()
    torch.accelerator.synchronize()
    samples_us = torch.tensor(
        [start.elapsed_time(end) * 1000.0 for start, end in zip(starts, ends)],
        dtype=torch.float64,
        device=torch.accelerator.current_device_index(),
    )
    dist.all_reduce(samples_us, op=dist.ReduceOp.MAX, group=device_group)
    return summarize(samples_us[1:].tolist())
 def make_inputs(
    num_tokens: int,
    rank: int,
    device: torch.device,
 ) -> tuple[torch.Tensor, torch.Tensor]:
    torch.manual_seed(20260726 + 100 * num_tokens + rank)
    routed = torch.randn(
        (num_tokens, LATENT_SIZE),
        dtype=torch.bfloat16,
        device=device,
    ).mul_(0.01)
    shared = torch.randn(
        (num_tokens, HIDDEN_SIZE),
        dtype=torch.bfloat16,
        device=device,
    )
    return routed, shared
 def make_reference(
    routed: torch.Tensor,
    shared: torch.Tensor,
    rms_weight: torch.Tensor,
    up_weight: torch.Tensor,
    device_group: dist.ProcessGroup,
 ) -> Callable[[], torch.Tensor]:
    routed_workspace = torch.empty_like(routed)
    shared_workspace = torch.empty_like(shared)
    def reference() -> torch.Tensor:
        routed_workspace.copy_(routed)
        dist.all_reduce(routed_workspace, group=device_group)
        normalized = F.rms_norm(
            routed_workspace,
            (LATENT_SIZE,),
            rms_weight,
            RMS_EPS,
        )
        projected = F.linear(normalized, up_weight)
        shared_workspace.copy_(shared)
        dist.all_reduce(shared_workspace, group=device_group)
        return projected.add(shared_workspace)
    return reference
 def check_fused_output(
    fused_output: torch.Tensor,
    reference: Callable[[], torch.Tensor],
    cpu_group: dist.ProcessGroup,
 ) -> None:
    dist.barrier(group=cpu_group)
    expected = reference()
    torch.testing.assert_close(fused_output, expected, atol=8e-2, rtol=3e-2)
 def benchmark_whole_tail(args: argparse.Namespace) -> None:
    if any(not 1 <= num_tokens <= 16 for num_tokens in args.num_tokens):
        raise ValueError("--num-tokens values must be in [1, 16]")
    if args.warmup_replays < 0 or args.samples <= 0:
        raise ValueError("warmup replays must be nonnegative and samples positive")
    if args.skinny_max_num_tokens is not None and any(
        not 0 <= cutoff <= 8 for cutoff in args.skinny_max_num_tokens
    ):
        raise ValueError("--skinny-max-num-tokens must be in [0, 8]")
    skinny_configs = dict(args.skinny_config or ())
    if len(skinny_configs) != len(args.skinny_config or ()):
        raise ValueError("--skinny-config must not repeat an M value")
    if any(not 1 <= num_tokens <= 8 for num_tokens in skinny_configs):
        raise ValueError("--skinny-config M values must be in [1, 8]")
    if not {"RANK", "WORLD_SIZE", "LOCAL_RANK"} <= os.environ.keys():
        raise RuntimeError("launch this benchmark with torchrun")
    rank = int(os.environ["RANK"])
    world_size = int(os.environ["WORLD_SIZE"])
    local_rank = int(os.environ["LOCAL_RANK"])
    device = torch.device("cuda", local_rank)
    torch.accelerator.set_device_index(device)
    init_distributed_environment()
    if world_size > 8:
        set_custom_all_reduce(False)
    initialize_model_parallel(tensor_model_parallel_size=world_size)
    device_group = get_tp_group().device_group
    cpu_group = dist.new_group(backend="gloo")
    if torch.cuda.get_device_capability(device)[0] != 10:
        raise RuntimeError("Kimi K3 latent-MoE tail requires SM100")
    torch.manual_seed(20260726)
    rms_weight = 1 + 0.1 * torch.randn(
        LATENT_SIZE,
        dtype=torch.bfloat16,
        device=device,
    )
    up_weight = (
        torch.randn(
            (HIDDEN_SIZE, LATENT_SIZE),
            dtype=torch.bfloat16,
            device=device,
        )
        / LATENT_SIZE**0.5
    )
    use_reference = args.backend in ("reference", "both")
    use_fused = args.backend in ("fused", "both")
    fused_ops = []
    if use_fused:
        production_config_for_m = fused_add_multicast_skinny_gemm.config_for_m
        def config_for_m(
            num_rows: int,
            shard_dim: int = 896,
        ) -> fused_add_multicast_skinny_gemm.SkinnyConfig:
            config = skinny_configs.get(num_rows)
            if config is not None:
                return config
            return production_config_for_m(num_rows, shard_dim)
        fused_add_multicast_skinny_gemm.config_for_m = config_for_m
        cutoffs = args.skinny_max_num_tokens or [latent_moe_tail._SKINNY_MAX_NUM_TOKENS]
        for cutoff in cutoffs:
            latent_moe_tail._SKINNY_MAX_NUM_TOKENS = cutoff
            latent_moe_tail.KimiK3LatentMoETailOp._instances.clear()
            fused_ops.append(
                (
                    cutoff,
                    latent_moe_tail.KimiK3LatentMoETailOp.initialize(
                        hidden_size=HIDDEN_SIZE,
                        latent_size=LATENT_SIZE,
                        dtype=torch.bfloat16,
                        device=device,
                        rms_eps=RMS_EPS,
                    ),
                )
            )
        cutedsl_warmup()
    results = []
    for num_tokens in args.num_tokens:
        routed, shared = make_inputs(num_tokens, rank, device)
        reference = make_reference(
            routed,
            shared,
            rms_weight,
            up_weight,
            device_group,
        )
        result: dict[str, Any] = {"num_tokens": num_tokens}
        if use_reference:
            reference_graph, _ = capture_tail_graph(reference, cpu_group)
            result["reference"] = benchmark_tail_graph(
                reference_graph,
                warmup_replays=args.warmup_replays,
                samples=args.samples,
                device_group=device_group,
                cpu_group=cpu_group,
            )
        for cutoff, fused_op in fused_ops:
            def fused(
                routed: torch.Tensor = routed,
                shared: torch.Tensor = shared,
                fused_op: latent_moe_tail.KimiK3LatentMoETailOp = fused_op,
            ) -> torch.Tensor:
                return fused_op(routed, shared, rms_weight, up_weight)
            fused_graph, fused_output = capture_tail_graph(fused, cpu_group)
            fused_key = "fused" if len(fused_ops) == 1 else f"fused_skinny_max_{cutoff}"
            result[fused_key] = benchmark_tail_graph(
                fused_graph,
                warmup_replays=args.warmup_replays,
                samples=args.samples,
                device_group=device_group,
                cpu_group=cpu_group,
            )
            check_fused_output(fused_output, reference, cpu_group)
            if "reference" in result:
                speedup = (
                    result["reference"]["median_us"] / result[fused_key]["median_us"]
                )
                if len(fused_ops) == 1:
                    result["speedup"] = speedup
                else:
                    result[f"{fused_key}_speedup"] = speedup
        results.append(result)
    properties = torch.cuda.get_device_properties(device)
    report = {
        "scope": "whole-tail",
        "device": properties.name,
        "compute_capability": list(torch.cuda.get_device_capability(device)),
        "world_size": world_size,
        "torch_version": torch.__version__,
        "cuda_version": torch.version.cuda,
        "warmup_replays": args.warmup_replays,
        "samples": args.samples,
        "skinny_max_num_tokens": [cutoff for cutoff, _ in fused_ops],
        "skinny_configs": {
            str(num_tokens): asdict(config)
            for num_tokens, config in skinny_configs.items()
        },
        "timing_scope": {
            "reference": (
                "two input copies, two AllReduces, RMSNorm, full replicated "
                "up-projection GEMM, and final add"
            ),
            "fused": (
                "routed AllReduce/RMSNorm plus shared ReduceScatter, sharded "
                "up-projection/multicast, and Lamport copy"
            ),
        },
        "results": results,
    }
    if rank == 0:
        rendered = json.dumps(report, indent=2)
        print(rendered, flush=True)
        if args.output is not None:
            args.output.parent.mkdir(parents=True, exist_ok=True)
            args.output.write_text(rendered + "\n", encoding="utf-8")
    dist.barrier(group=cpu_group)
 def main() -> None:
    args = parse_args()
    if args.scope == "up-projection":
        benchmark_up_projection(args)
        return
    from vllm.config import VllmConfig, set_current_vllm_config
    with set_current_vllm_config(VllmConfig()):
        benchmark_whole_tail(args)
 if __name__ == "__main__":
    main()
@@ -1,239 +0,0 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import argparse
 import json
 import os
 import statistics
 from collections.abc import Callable
 import torch
 import torch.distributed as dist
 import vllm._custom_ops as ops
 from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce
 def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--tokens", type=int, nargs="+", default=[8, 32, 128, 1024])
    parser.add_argument("--hidden-size", type=int, default=7168)
    parser.add_argument("--graph-repeats", type=int, default=20)
    parser.add_argument("--warmup-replays", type=int, default=5)
    parser.add_argument("--samples", type=int, default=15)
    return parser.parse_args()
 def capture_graph(op: Callable[[], None], repeats: int) -> torch.cuda.CUDAGraph:
    stream = torch.cuda.Stream()
    stream.wait_stream(torch.cuda.current_stream())
    with torch.cuda.stream(stream):
        for _ in range(3):
            op()
    stream.synchronize()
    graph = torch.cuda.CUDAGraph()
    with torch.cuda.graph(graph, stream=stream):
        for _ in range(repeats):
            op()
    torch.cuda.current_stream().wait_stream(stream)
    return graph
 def max_rank_graph_time(
    graph: torch.cuda.CUDAGraph,
    repeats: int,
    warmup_replays: int,
    samples: int,
    device_group: dist.ProcessGroup,
    cpu_group: dist.ProcessGroup,
 ) -> float:
    for _ in range(warmup_replays):
        graph.replay()
    torch.accelerator.synchronize()
    timings = []
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)
    for _ in range(samples):
        dist.barrier(group=cpu_group)
        start.record()
        graph.replay()
        end.record()
        end.synchronize()
        elapsed = torch.tensor(
            start.elapsed_time(end) / repeats,
            dtype=torch.float64,
            device=torch.accelerator.current_device_index(),
        )
        dist.all_reduce(elapsed, op=dist.ReduceOp.MAX, group=device_group)
        timings.append(elapsed.item())
    return statistics.median(timings)
 def check_outputs(
    comm: CustomAllreduce,
    local: torch.Tensor,
    reduce_input: torch.Tensor,
    device_group: dist.ProcessGroup,
 ) -> None:
    expected_gather = torch.empty(
        (local.shape[0] * dist.get_world_size(), local.shape[1]),
        dtype=local.dtype,
        device=local.device,
    )
    dist.all_gather_into_tensor(expected_gather, local, group=device_group)
    gathered = comm.custom_all_gather(local)
    assert gathered is not None
    torch.testing.assert_close(gathered, expected_gather)
    expected_scatter = torch.empty_like(local)
    dist.reduce_scatter_tensor(
        expected_scatter,
        reduce_input.clone(),
        group=device_group,
    )
    scattered = comm.custom_reduce_scatter(reduce_input)
    assert scattered is not None
    torch.testing.assert_close(scattered, expected_scatter)
 def benchmark_shape(
    comm: CustomAllreduce,
    global_tokens: int,
    hidden_size: int,
    graph_repeats: int,
    warmup_replays: int,
    samples: int,
    device_group: dist.ProcessGroup,
    cpu_group: dist.ProcessGroup,
 ) -> dict[str, float | int]:
    world_size = dist.get_world_size()
    rank = dist.get_rank()
    padded_tokens = (global_tokens + world_size - 1) // world_size * world_size
    local_tokens = padded_tokens // world_size
    local = torch.full(
        (local_tokens, hidden_size),
        rank + 1,
        dtype=torch.bfloat16,
        device=torch.accelerator.current_device_index(),
    )
    reduce_input = torch.full(
        (padded_tokens, hidden_size),
        rank + 1,
        dtype=torch.bfloat16,
        device=local.device,
    )
    check_outputs(comm, local, reduce_input, device_group)
    custom_gather_out = torch.empty(
        (padded_tokens, hidden_size),
        dtype=local.dtype,
        device=local.device,
    )
    custom_scatter_out = torch.empty_like(local)
    nccl_gather_out = torch.empty_like(custom_gather_out)
    nccl_scatter_out = torch.empty_like(local)
    def custom_ag() -> None:
        ops.mnnvl_lamport_all_gather(
            comm._ptr,
            local,
            custom_gather_out,
            comm.mnnvl_lamport_ag_local_ptr,
            comm.mnnvl_lamport_ag_multicast_ptr,
            comm.mnnvl_lamport_ag_epoch_ptr,
            comm.mnnvl_buffer_size,
        )
    def custom_rs() -> None:
        ops.mnnvl_lamport_reduce_scatter(
            comm._ptr,
            reduce_input,
            custom_scatter_out,
            comm.mnnvl_lamport_rs_local_ptr,
            comm.mnnvl_lamport_rs_epoch_ptr,
            comm.mnnvl_buffer_size,
        )
    def nccl_ag() -> None:
        dist.all_gather_into_tensor(nccl_gather_out, local, group=device_group)
    def nccl_rs() -> None:
        dist.reduce_scatter_tensor(
            nccl_scatter_out,
            reduce_input,
            group=device_group,
        )
    graphs = {
        "custom_ag_us": capture_graph(custom_ag, graph_repeats),
        "nccl_ag_us": capture_graph(nccl_ag, graph_repeats),
        "custom_rs_us": capture_graph(custom_rs, graph_repeats),
        "nccl_rs_us": capture_graph(nccl_rs, graph_repeats),
    }
    times = {
        name: max_rank_graph_time(
            graph,
            graph_repeats,
            warmup_replays,
            samples,
            device_group,
            cpu_group,
        )
        * 1000
        for name, graph in graphs.items()
    }
    torch.testing.assert_close(custom_gather_out, nccl_gather_out)
    torch.testing.assert_close(custom_scatter_out, nccl_scatter_out)
    return {
        "global_tokens": global_tokens,
        "padded_tokens": padded_tokens,
        "local_bytes": local.nbytes,
        "full_bytes": reduce_input.nbytes,
        **times,
        "ag_speedup": times["nccl_ag_us"] / times["custom_ag_us"],
        "rs_speedup": times["nccl_rs_us"] / times["custom_rs_us"],
    }
 def main() -> None:
    args = parse_args()
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.accelerator.set_device_index(local_rank)
    dist.init_process_group("nccl")
    device_group = dist.group.WORLD
    cpu_group = dist.new_group(backend="gloo")
    comm = CustomAllreduce(
        group=cpu_group,
        device=torch.device("cuda", local_rank),
    )
    assert not comm.disabled
    assert comm.world_size == 16
    assert comm.mnnvl_only
    assert comm.mnnvl_multicast_ptr
    results = [
        benchmark_shape(
            comm,
            tokens,
            args.hidden_size,
            args.graph_repeats,
            args.warmup_replays,
            args.samples,
            device_group,
            cpu_group,
        )
        for tokens in args.tokens
    ]
    if dist.get_rank() == 0:
        print(json.dumps(results, indent=2), flush=True)
    comm.close()
    dist.destroy_process_group(cpu_group)
    dist.destroy_process_group()
 if __name__ == "__main__":
    main()
@@ -1,201 +0,0 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """
 Benchmark the RDNAHybridW4A16LinearKernel across decode and prefill shapes.
 Usage:
    python benchmark_int4_gemm.py
    python benchmark_int4_gemm.py --models Qwen/Qwen3-4B
    python benchmark_int4_gemm.py --group-size 128
 """
 import argparse
 import copy
 import itertools
 import os
 import torch
 from vllm.triton_utils import triton
 # ---------------------------------------------------------------------------
 # Weight shapes: [K, N], TP_SPLIT_DIM
 # ---------------------------------------------------------------------------
 WEIGHT_SHAPES = {
    "Qwen/Qwen3-4B": [
        ([2560, 3840], 1),  # qkv_proj
        ([2560, 2560], 0),  # o_proj
        ([2560, 19456], 1),  # gate_up_proj
        ([9728, 2560], 0),  # down_proj
    ],
    "Qwen/Qwen2.5-7B-Instruct": [
        ([3584, 4608], 1),
        ([3584, 3584], 0),
        ([3584, 37888], 1),
        ([18944, 3584], 0),
    ],
    "trymirai/SmolLM2-1.7B-Instruct-AWQ": [
        ([2048, 6144], 1),  # qkv_proj
        ([2048, 2048], 0),  # o_proj
        ([2048, 16384], 1),  # gate_up_proj
        ([8192, 2048], 0),  # down_proj
    ],
    "RedHatAI/Qwen3-8B-quantized.w4a16": [
        ([4096, 6144], 1),  # qkv_proj
        ([4096, 4096], 0),  # o_proj
        ([4096, 24576], 1),  # gate_up_proj
        ([12288, 4096], 0),  # down_proj
    ],
 }
 # ---------------------------------------------------------------------------
 # Weight packing
 # ---------------------------------------------------------------------------
 def prepare_hybrid_weights(K, N, group_size, device="cuda"):
    """Create random weights for benchmarking.
    Returns (w_q_skinny, w_s_skinny, w_fp16, w_zp). The triton path derives
    its int32 view from w_q_skinny, so no separate int32 buffer is returned.
    """
    num_groups = K // group_size
    # Random packed weights — actual values don't matter for throughput
    w_q_skinny_i32 = torch.randint(
        0, 2**31, (N, K // 8), dtype=torch.int32, device=device
    )
    w_q_skinny = w_q_skinny_i32.view(torch.int8).contiguous()
    w_s_skinny = torch.randn(N, num_groups, dtype=torch.float16, device=device) * 0.01
    # Raw per-group zero-points for asymmetric benchmarks
    w_zp = torch.randint(0, 16, (N, num_groups), dtype=torch.int32, device=device).to(
        torch.float16
    )
    # FP16 baseline for F.linear
    w_fp16 = torch.randn(N, K, dtype=torch.float16, device=device) * 0.01
    return w_q_skinny, w_s_skinny, w_fp16, w_zp
 # ---------------------------------------------------------------------------
 # Benchmark
 # ---------------------------------------------------------------------------
 PROVIDERS = ["torch-fp16", "hybrid-w4a16", "hybrid-w4a16-zp"]
@triton.testing.perf_report(
    triton.testing.Benchmark(
        x_names=["batch_size"],
        x_vals=[1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096],
        x_log=False,
        line_arg="provider",
        line_vals=PROVIDERS,
        line_names=PROVIDERS,
        ylabel="TFLOP/s (larger is better)",
        plot_name="FP16 vs Hybrid W4A16",
        args={},
    )
 )
 def benchmark(batch_size, provider, N, K, group_size, weights):
    M = batch_size
    device = "cuda"
    dtype = torch.float16
    a = torch.randn((M, K), device=device, dtype=dtype)
    quantiles = [0.5, 0.2, 0.8]
    if provider == "torch-fp16":
        w_fp16 = weights["w_fp16"]
        ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
            lambda: torch.nn.functional.linear(a, w_fp16),
            quantiles=quantiles,
        )
    elif provider in ("hybrid-w4a16", "hybrid-w4a16-zp"):
        from vllm.model_executor.kernels.linear.mixed_precision import (
            rdna_hybrid_w4a16 as _k,
        )
        _rdna_hybrid_w4a16_apply_impl = _k._rdna_hybrid_w4a16_apply_impl
        from vllm.utils.platform_utils import num_compute_units
        w = weights
        cu_count = num_compute_units()
        use_zp = provider == "hybrid-w4a16-zp"
        def run():
            return _rdna_hybrid_w4a16_apply_impl(
                a,
                w["w_q_skinny"],
                w["w_s_skinny"],
                w["w_zp"] if use_zp else None,
                None,  # bias
                cu_count,
                group_size,
            )
        ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
            run,
            quantiles=quantiles,
        )
    else:
        return 0.0, 0.0, 0.0
    to_tflops = lambda t_ms: (2 * M * N * K) * 1e-12 / (t_ms * 1e-3)
    return to_tflops(ms), to_tflops(max_ms), to_tflops(min_ms)
 def prepare_shapes(args):
    KN_model_names = []
    for model, tp_size in itertools.product(args.models, args.tp_sizes):
        for KN, tp_dim in copy.deepcopy(WEIGHT_SHAPES[model]):
            KN[tp_dim] //= tp_size
            KN.append(model)
            KN_model_names.append(KN)
    return KN_model_names
 if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Benchmark RDNAHybridW4A16LinearKernel"
    )
    parser.add_argument(
        "--models",
        nargs="+",
        type=str,
        default=["Qwen/Qwen3-4B"],
        choices=list(WEIGHT_SHAPES.keys()),
    )
    parser.add_argument("--tp-sizes", nargs="+", type=int, default=[1])
    parser.add_argument("--group-size", type=int, default=128)
    parser.add_argument("--save-path", type=str, default=None)
    args = parser.parse_args()
    for K, N, model in prepare_shapes(args):
        group_size = args.group_size
        print(f"\n{'=' * 70}")
        print(f"{model}, N={N} K={K}, group_size={group_size}")
        print(f"{'=' * 70}")
        w_q_skinny, w_s_skinny, w_fp16, w_zp = prepare_hybrid_weights(K, N, group_size)
        weights = {
            "w_q_skinny": w_q_skinny,
            "w_s_skinny": w_s_skinny,
            "w_fp16": w_fp16,
            "w_zp": w_zp,
        }
        save_path = args.save_path or f"bench_int4_res_n{N}_k{K}"
        os.makedirs(save_path, exist_ok=True)
        benchmark.run(
            print_data=True,
            show_plots=False,
            save_path=save_path,
            N=N,
            K=K,
            group_size=group_size,
            weights=weights,
        )
    print("\nBenchmark finished!")
@@ -154,7 +154,7 @@ def main(
                scale=scale,
                causal=True,
                alibi_slopes=None,
-                sliding_window=window_size if sliding_window is not None else -1,
+                sliding_window=window_size,
                block_table=block_tables,
                softcap=0,
                scheduler_metadata=metadata,
@@ -1,267 +0,0 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """End-to-end autoregressive decode benchmark: ReplaySSM vs the standard SSM kernel.
 Loads a hybrid Mamba2 model, replicates one prompt across the batch, and times a
 long greedy decode (CUDA graphs on) once with the standard kernel and once with
 ReplaySSM, then reports the per-step / throughput speedup. The two modes run in
 separate subprocesses so each gets a clean CUDA context.
 The FlashInfer FP4-MoE autotuner is disabled by default (it is unstable under
 CUDA-graph capture on the pre-release Blackwell FP4 path); pass
 --no-disable-flashinfer-autotune for non-FP4 models.
 Examples:
    python e2e_decode_speedup.py --model-id nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16
    python e2e_decode_speedup.py --dtype auto --buffer-len 16 \
        --model-id nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4   # B300 NVFP4
 """
 import argparse
 import json
 import os
 import subprocess
 import sys
 import time
 DEFAULT_PROMPT = "My cat wrote all this CUDA code for a new language model and"
 MODE_LABEL = {"standard": "standard", "replayssm": "ReplaySSM"}
 def parse_args():
    p = argparse.ArgumentParser(
        description="E2E decode speedup: ReplaySSM vs the standard SSM kernel."
    )
    p.add_argument("--model-id", default="nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16")
    p.add_argument("--prompt", default=DEFAULT_PROMPT)
    p.add_argument("--batch-size", type=int, default=256)
    p.add_argument("--num-steps", type=int, default=1000)
    p.add_argument("--warmup-steps", type=int, default=128)
    p.add_argument("--repeats", type=int, default=1)
    p.add_argument(
        "--buffer-len", type=int, default=16, help="ReplaySSM input-buffer length."
    )
    p.add_argument(
        "--dtype",
        default="bfloat16",
        choices=["bfloat16", "float16", "float32", "auto"],
    )
    p.add_argument("--gpu-memory-utilization", type=float, default=0.9)
    p.add_argument("--max-model-len", type=int, default=None)
    p.add_argument(
        "--disable-flashinfer-autotune",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="Disable the FlashInfer FP4-MoE autotuner (default: on). "
        "It is unstable under CUDA-graph capture on the "
        "pre-release Blackwell FP4 path; pass "
        "--no-disable-flashinfer-autotune for non-FP4 models.",
    )
    p.add_argument(
        "--mamba-ssm-cache-dtype",
        default="auto",
        choices=["auto", "float32", "float16", "bfloat16"],
        help="SSM state dtype (both modes). 'auto' = config-driven; "
        "'float32' = fp32 state, 'bfloat16' = s16 state.",
    )
    p.add_argument(
        "--baseline-ssm-config",
        default="",
        help="Pin the STANDARD baseline's SSM launch config as "
        "'bsm,nw' via override_ssm_config (forces the in-process "
        "engine so the override reaches the kernel). Empty = off.",
    )
    p.add_argument(
        "--worker",
        choices=["standard", "replayssm"],
        default=None,
        help=argparse.SUPPRESS,
    )
    return p.parse_args()
 def resolve_max_model_len(args) -> int:
    if args.max_model_len is not None:
        return args.max_model_len
    return args.num_steps + 256
 def run_worker(args):
    # override_ssm_config is a module global; it only reaches the model if the
    # engine runs in-process (default V1 spawns a separate EngineCore). Force it.
    if args.worker == "standard" and args.baseline_ssm_config:
        os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
    import torch
    from vllm import LLM, SamplingParams
    mode = args.worker
    max_model_len = resolve_max_model_len(args)
    llm_kwargs = dict(
        model=args.model_id,
        tensor_parallel_size=1,
        dtype=args.dtype,
        max_model_len=max_model_len,
        trust_remote_code=True,
        enable_prefix_caching=False,
        enable_chunked_prefill=False,
        max_num_seqs=args.batch_size,
        max_num_batched_tokens=max(max_model_len, args.batch_size * 64),
        enforce_eager=False,
        disable_log_stats=True,
        gpu_memory_utilization=args.gpu_memory_utilization,
        # SSM state dtype (applies to both standard and ReplaySSM).
        mamba_ssm_cache_dtype=args.mamba_ssm_cache_dtype,
    )
    if args.disable_flashinfer_autotune:
        # FP4-MoE autotuner is unstable under CUDA-graph capture on Blackwell;
        # re-enable (--no-disable-flashinfer-autotune) only for non-FP4 models.
        llm_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False}
    if mode == "replayssm":
        llm_kwargs.update(use_replayssm=True, replayssm_buffer_len=args.buffer_len)
    _ssm_cm = None
    if mode == "standard" and args.baseline_ssm_config:
        from vllm.model_executor.layers.mamba.ops.mamba_ssm import override_ssm_config
        _bsm, _nw = (int(x) for x in args.baseline_ssm_config.split(","))
        _ssm_cm = override_ssm_config((_bsm, _nw))
        _ssm_cm.__enter__()  # active through LLM() graph capture + decode
        print(
            f"[{mode}] override_ssm_config -> (BLOCK_SIZE_M={_bsm}, num_warps={_nw})",
            flush=True,
        )
    llm = LLM(**llm_kwargs)
    prompts = [args.prompt] * args.batch_size
    def timed_generate(n_tokens):
        sp = SamplingParams(
            n=1,
            temperature=0.0,
            ignore_eos=True,
            min_tokens=n_tokens,
            max_tokens=n_tokens,
        )
        if torch.accelerator.is_available():
            torch.accelerator.synchronize()
        t0 = time.perf_counter()
        outs = llm.generate(prompts, sp, use_tqdm=False)
        if torch.accelerator.is_available():
            torch.accelerator.synchronize()
        elapsed = time.perf_counter() - t0
        produced = min(len(o.outputs[0].token_ids) for o in outs)
        assert produced == n_tokens, f"expected {n_tokens} tokens, got {produced}"
        return elapsed
    timed_generate(args.warmup_steps)
    best = None
    for _ in range(args.repeats):
        elapsed = timed_generate(args.num_steps)
        tok_s = args.batch_size * args.num_steps / elapsed
        per_step_ms = elapsed / args.num_steps * 1e3
        print(
            f"[{mode}] {elapsed:.3f}s  {tok_s:,.0f} tok/s  {per_step_ms:.3f} ms/step",
            flush=True,
        )
        if best is None or elapsed < best["elapsed_s"]:
            best = {
                "mode": mode,
                "elapsed_s": elapsed,
                "tok_s": tok_s,
                "per_step_ms": per_step_ms,
            }
    print("RESULT_JSON " + json.dumps(best), flush=True)
    if _ssm_cm is not None:
        _ssm_cm.__exit__(None, None, None)
 def run_one_mode(args, mode) -> dict:
    cmd = [
        sys.executable,
        __file__,
        "--worker",
        mode,
        "--model-id",
        args.model_id,
        "--prompt",
        args.prompt,
        "--batch-size",
        str(args.batch_size),
        "--num-steps",
        str(args.num_steps),
        "--warmup-steps",
        str(args.warmup_steps),
        "--repeats",
        str(args.repeats),
        "--buffer-len",
        str(args.buffer_len),
        "--dtype",
        args.dtype,
        "--gpu-memory-utilization",
        str(args.gpu_memory_utilization),
        "--mamba-ssm-cache-dtype",
        args.mamba_ssm_cache_dtype,
        "--baseline-ssm-config",
        args.baseline_ssm_config,
    ]
    cmd.append(
        "--disable-flashinfer-autotune"
        if args.disable_flashinfer_autotune
        else "--no-disable-flashinfer-autotune"
    )
    if args.max_model_len is not None:
        cmd += ["--max-model-len", str(args.max_model_len)]
    result = None
    proc = subprocess.Popen(
        cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1
    )
    for line in proc.stdout:
        sys.stdout.write(line)
        sys.stdout.flush()
        if line.startswith("RESULT_JSON "):
            result = json.loads(line[len("RESULT_JSON ") :])
    proc.wait()
    if proc.returncode != 0:
        raise RuntimeError(f"mode '{mode}' worker exited with {proc.returncode}")
    if result is None:
        raise RuntimeError(f"mode '{mode}' produced no RESULT_JSON line")
    return result
 def main():
    args = parse_args()
    if args.worker is not None:
        run_worker(args)
        return
    print(
        f"model={args.model_id}  batch_size={args.batch_size}  "
        f"steps={args.num_steps}  buffer_len={args.buffer_len}  dtype={args.dtype}"
    )
    std = run_one_mode(args, "standard")
    fla = run_one_mode(args, "replayssm")
    speedup = std["per_step_ms"] / fla["per_step_ms"]
    print()
    header = f"{'mode':<10}{'ms/step':>12}{'tok/s':>16}{'wall (s)':>12}"
    print(header)
    print("-" * len(header))
    for r in (std, fla):
        print(
            f"{MODE_LABEL[r['mode']]:<10}{r['per_step_ms']:>12.3f}"
            f"{r['tok_s']:>16,.0f}{r['elapsed_s']:>12.3f}"
        )
    print("-" * len(header))
    print(f"speedup (standard / ReplaySSM, per step): {speedup:.3f}x")
 if __name__ == "__main__":
    main()
@@ -15,7 +15,6 @@ endif()
 #
 set(ENABLE_X86_ISA $ENV{VLLM_CPU_X86})
 set(ENABLE_ARM_BF16 $ENV{VLLM_CPU_ARM_BF16})
 set(ENABLE_ARM_I8MM $ENV{VLLM_CPU_ARM_I8MM})
 set(ENABLE_RVV_BF16 $ENV{VLLM_CPU_RVV_BF16})
 include_directories("${CMAKE_SOURCE_DIR}/csrc")
@@ -97,14 +96,12 @@ if (MACOSX_FOUND AND CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64")
    set(ENABLE_NUMA OFF)
    check_sysctl(hw.optional.neon ASIMD_FOUND)
    check_sysctl(hw.optional.arm.FEAT_BF16 ARM_BF16_FOUND)
    check_sysctl(hw.optional.arm.FEAT_I8MM ARM_I8MM_FOUND)
 else()
    find_isa(${CPUINFO} "Power11" POWER11_FOUND)
    find_isa(${CPUINFO} "POWER10" POWER10_FOUND)
    find_isa(${CPUINFO} "POWER9" POWER9_FOUND)
    find_isa(${CPUINFO} "asimd" ASIMD_FOUND) # Check for ARM NEON support
    find_isa(${CPUINFO} "bf16" ARM_BF16_FOUND) # Check for ARM BF16 support
    find_isa(${CPUINFO} "i8mm" ARM_I8MM_FOUND) # Check for ARM I8MM support
    find_isa(${CPUINFO} "S390" S390_FOUND)
    find_isa(${CPUINFO} "zvfhmin" RVV_FP16_FOUND) # Check for RISC-V Vector FP16 support
    find_isa(${CPUINFO} "zvfbfmin" RVV_BF16_FOUND) # Check for RISC-V Vector BF16 support
@@ -114,11 +111,6 @@ else()
        set(ARM_BF16_FOUND ON)
        message(STATUS "ARM BF16 support enabled via VLLM_CPU_ARM_BF16 environment variable")
    endif()
    if (ENABLE_ARM_I8MM)
        set(ARM_I8MM_FOUND ON)
        message(STATUS
            "ARM I8MM support enabled via VLLM_CPU_ARM_I8MM environment variable")
    endif()
    # Some kernels (e.g. Bianbu on Spacemit X100) do not report zvfbfmin
    # in /proc/cpuinfo despite hardware support. VLLM_CPU_RVV_BF16=1
    # overrides the detection result.
@@ -174,11 +166,6 @@ elseif (ASIMD_FOUND)
        message(WARNING "BF16 functionality is not available")
        set(MARCH_FLAGS "-march=armv8.2-a+dotprod+fp16")  
    endif()
    if(ARM_I8MM_FOUND)
        message(STATUS "I8MM extension detected")
        string(APPEND MARCH_FLAGS "+i8mm")
        add_compile_definitions(ARM_I8MM_SUPPORT)
    endif()
    list(APPEND CXX_COMPILE_FLAGS ${MARCH_FLAGS})     
 elseif (S390_FOUND)
    message(STATUS "S390 detected")
@@ -443,7 +430,6 @@ set(VLLM_EXT_SRC
    "csrc/cpu/layernorm.cpp"
    "csrc/cpu/mla_decode.cpp"
    "csrc/cpu/pos_encoding.cpp"
    "csrc/cpu/mamba_cpu.cpp"
    "csrc/moe/dynamic_4bit_int_moe_cpu.cpp"
    "csrc/cpu/cpu_attn.cpp"
    "csrc/cpu/torch_bindings.cpp")
@@ -460,13 +446,8 @@ if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
        "csrc/cpu/shm.cpp"
        "csrc/cpu/activation_lut_bf16.cpp"
        "csrc/cpu/cpu_tanhf_neon.hpp"
        "csrc/cpu/cpu_fused_moe.cpp"
        ${VLLM_EXT_SRC})
    if (ARM_BF16_FOUND)
        set(VLLM_EXT_SRC "csrc/cpu/cpu_fused_moe.cpp" ${VLLM_EXT_SRC})
        if (ARM_I8MM_FOUND)
            set(VLLM_EXT_SRC "csrc/cpu/cpu_fused_moe_int8.cpp" ${VLLM_EXT_SRC})
        endif()
    endif()
 endif()
 if (POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND)	
@@ -508,7 +489,6 @@ if (ENABLE_X86_ISA)
        "csrc/cpu/spec_decode_utils.cpp"
        "csrc/cpu/cpu_attn.cpp"
        "csrc/cpu/dnnl_kernels.cpp"
        "csrc/cpu/mamba_cpu.cpp"
        "csrc/cpu/torch_bindings.cpp"
        # TODO: Remove these files
        "csrc/cpu/activation.cpp"
@@ -522,7 +502,6 @@ if (ENABLE_X86_ISA)
        "csrc/cpu/utils.cpp"
        "csrc/cpu/spec_decode_utils.cpp"
        "csrc/cpu/cpu_attn.cpp"
        "csrc/cpu/mamba_cpu.cpp"
        "csrc/cpu/dnnl_kernels.cpp"
        "csrc/cpu/torch_bindings.cpp"
        # TODO: Remove these files
@@ -68,9 +68,6 @@ endif()
 if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9)
    list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0f")
    if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.4)
      list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.7f")
    endif()
  else()
    list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0a")
  endif()
@@ -1,74 +0,0 @@
 include(FetchContent)
 if(DEFINED ENV{FLASH_KDA_SRC_DIR})
  set(FLASH_KDA_SRC_DIR $ENV{FLASH_KDA_SRC_DIR})
 endif()
 if(FLASH_KDA_SRC_DIR)
  FetchContent_Declare(
    flashkda
    SOURCE_DIR ${FLASH_KDA_SRC_DIR}
  )
 else()
  FetchContent_Declare(
    flashkda
    GIT_REPOSITORY https://github.com/vllm-project/FlashKDA.git
    GIT_TAG a3e42bbbece3bb38f7c426b880315294a336e82f
    GIT_PROGRESS TRUE
    GIT_SUBMODULES cutlass
  )
 endif()
 FetchContent_MakeAvailable(flashkda)
 message(STATUS "FlashKDA is available at ${flashkda_SOURCE_DIR}")
 set(FLASH_KDA_SUPPORT_ARCHS)
 if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0)
  list(APPEND FLASH_KDA_SUPPORT_ARCHS "9.0a")
 endif()
 if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
  list(APPEND FLASH_KDA_SUPPORT_ARCHS "10.0f" "12.0f")
 elseif(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9)
  list(APPEND FLASH_KDA_SUPPORT_ARCHS "10.0a" "10.3a" "12.0a")
 endif()
 cuda_archs_loose_intersection(
  FLASH_KDA_ARCHS "${FLASH_KDA_SUPPORT_ARCHS}" "${CUDA_ARCHS}")
 if(FLASH_KDA_ARCHS)
  message(STATUS "FlashKDA CUDA architectures: ${FLASH_KDA_ARCHS}")
  set(FLASH_KDA_SOURCES
    csrc/flashkda_registration.cpp
    ${flashkda_SOURCE_DIR}/csrc/flash_kda.cpp
    ${flashkda_SOURCE_DIR}/csrc/smxx/fwd_launch.cu)
  set(FLASH_KDA_INCLUDES
    ${flashkda_SOURCE_DIR}/csrc
    ${flashkda_SOURCE_DIR}/cutlass/include
    ${flashkda_SOURCE_DIR}/cutlass/examples/common
    ${flashkda_SOURCE_DIR}/cutlass/tools/util/include)
  set_gencode_flags_for_srcs(
    SRCS "${FLASH_KDA_SOURCES}"
    CUDA_ARCHS "${FLASH_KDA_ARCHS}")
  define_extension_target(
    _flashkda_C
    DESTINATION vllm
    LANGUAGE ${VLLM_GPU_LANG}
    SOURCES ${FLASH_KDA_SOURCES}
    COMPILE_FLAGS ${VLLM_GPU_FLAGS}
    ARCHITECTURES ${VLLM_GPU_ARCHES}
    INCLUDE_DIRECTORIES ${FLASH_KDA_INCLUDES}
    USE_SABI 3
    WITH_SOABI)
  target_compile_options(_flashkda_C PRIVATE
    $<$<COMPILE_LANGUAGE:CUDA>:-UPy_LIMITED_API --expt-relaxed-constexpr --expt-extended-lambda --use_fast_math -O3>
    $<$<COMPILE_LANGUAGE:CXX>:-UPy_LIMITED_API>)
 else()
  message(STATUS
    "FlashKDA will not compile: CUDA >=12.0 and a supported architecture "
    "(SM90, SM10x, or SM12x) are required")
  add_custom_target(_flashkda_C)
 endif()
@@ -19,7 +19,7 @@ else()
  FetchContent_Declare(
        flashmla
        GIT_REPOSITORY https://github.com/vllm-project/FlashMLA
-        GIT_TAG a8f794d1251cbfd88a5011445dd5582289c727e4
+        GIT_TAG b70aff3d110a2b1a037e62eac295166b5143643a
        GIT_PROGRESS TRUE
        CONFIGURE_COMMAND ""
        BUILD_COMMAND ""
@@ -35,7 +35,7 @@ set(FLASHMLA_VENDOR_DIR "${CMAKE_SOURCE_DIR}/vllm/third_party/flashmla")
 file(MAKE_DIRECTORY "${FLASHMLA_VENDOR_DIR}")
 file(READ "${flashmla_SOURCE_DIR}/flash_mla/flash_mla_interface.py"
     FLASHMLA_INTERFACE_CONTENT)
-string(REPLACE "flash_mla_cuda = torch.ops._flashmla_C"
+string(REPLACE "import flash_mla.cuda as flash_mla_cuda"
               "import vllm._flashmla_C\nflash_mla_cuda = torch.ops._flashmla_C"
               FLASHMLA_INTERFACE_CONTENT
               "${FLASHMLA_INTERFACE_CONTENT}")
@@ -60,9 +60,6 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9)
    # CUDA 12.9 has introduced "Family-Specific Architecture Features"
    # this supports all compute_10x family
    list(APPEND SUPPORT_ARCHS "10.0f")
    if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.4)
        list(APPEND SUPPORT_ARCHS "10.7f")
    endif()
 elseif(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
    list(APPEND SUPPORT_ARCHS "10.0a")
 endif()
@@ -75,7 +72,7 @@ if(FLASH_MLA_ARCHS)
    list(APPEND VLLM_FLASHMLA_GPU_FLAGS "--expt-relaxed-constexpr" "--expt-extended-lambda" "--use_fast_math")
    set(FlashMLA_SOURCES
-        ${flashmla_SOURCE_DIR}/csrc/api/api.cpp
+        ${flashmla_SOURCE_DIR}/csrc/torch_api.cpp
        # Misc kernels for decoding
        ${flashmla_SOURCE_DIR}/csrc/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu
@@ -131,7 +128,6 @@ if(FLASH_MLA_ARCHS)
    set(FlashMLA_Extension_INCLUDES
        ${flashmla_SOURCE_DIR}/csrc
        ${flashmla_SOURCE_DIR}/csrc/kerutils/include
        ${flashmla_SOURCE_DIR}/csrc/extension/sm90/dense_fp8/
        ${flashmla_SOURCE_DIR}/csrc/cutlass/include
        ${flashmla_SOURCE_DIR}/csrc/cutlass/tools/util/include
@@ -156,18 +152,15 @@ if(FLASH_MLA_ARCHS)
        USE_SABI 3
        WITH_SOABI)
-    # Enable C++20 for the FlashMLA sources (required for std::span, requires, etc.)
+    # Keep Stable ABI for the module, but *not* for CUDA/C++ files.
    # This prevents Py_LIMITED_API from affecting nvcc and C++ compiles.
    # Also enable C++20 for the FlashMLA sources (required for std::span, requires, etc.)
    target_compile_options(_flashmla_C PRIVATE
        $<$<COMPILE_LANGUAGE:CUDA>:-UPy_LIMITED_API>
        $<$<COMPILE_LANGUAGE:CXX>:-UPy_LIMITED_API>
        $<$<COMPILE_LANGUAGE:CXX>:-std=c++20>
        $<$<COMPILE_LANGUAGE:CUDA>:-std=c++20>)
    # _flashmla_C is now ABI-stable torch 2.11+
    target_compile_definitions(_flashmla_C PRIVATE
        TORCH_TARGET_VERSION=0x020B000000000000ULL)
    if(VLLM_GPU_LANG STREQUAL "CUDA")
        target_compile_definitions(_flashmla_C PRIVATE USE_CUDA)
    endif()
    define_extension_target(
        _flashmla_extension_C
        DESTINATION vllm
@@ -179,15 +172,15 @@ if(FLASH_MLA_ARCHS)
        USE_SABI 3
        WITH_SOABI)
-    # _flashmla_extension_C is now ABI-stable w/ torch 2.11+
+    # Keep Stable ABI for the module, but *not* for CUDA/C++ files.
-    target_compile_definitions(_flashmla_extension_C PRIVATE
+    # This prevents Py_LIMITED_API from affecting nvcc and C++ compiles.
-        TORCH_TARGET_VERSION=0x020B000000000000ULL)
+    target_compile_options(_flashmla_extension_C PRIVATE
-    if(VLLM_GPU_LANG STREQUAL "CUDA")
+        $<$<COMPILE_LANGUAGE:CUDA>:-UPy_LIMITED_API>
-        target_compile_definitions(_flashmla_extension_C PRIVATE USE_CUDA)
+        $<$<COMPILE_LANGUAGE:CXX>:-UPy_LIMITED_API>)
    endif()
 else()
    message(STATUS "FlashMLA will not compile: unsupported CUDA architecture ${CUDA_ARCHS}")
    # Create empty targets for setup.py on unsupported systems
    add_custom_target(_flashmla_C)
    add_custom_target(_flashmla_extension_C)
 endif()
@@ -17,7 +17,7 @@ else()
  FetchContent_Declare(
    fmha_sm100
    GIT_REPOSITORY https://github.com/vllm-project/MSA.git
-    GIT_TAG 890aaa1a37a598ad17ccff0827fea21540d381fa
+    GIT_TAG 2e63ec37a0fc29bc20f39cd1a52e0f5affc33a73
    GIT_PROGRESS TRUE
    CONFIGURE_COMMAND ""
    BUILD_COMMAND ""
@@ -22,7 +22,7 @@ if(QUTLASS_SRC_DIR)
  set(qutlass_BINARY_DIR "${CMAKE_BINARY_DIR}/qutlass-binary-dir-unused")
 else()
  set(_QUTLASS_UPSTREAM_REPO "https://github.com/IST-DASLab/qutlass.git")
-  set(_QUTLASS_UPSTREAM_TAG "e74319e3405ce6d71965732880f5dc1f52371f64")
+  set(_QUTLASS_UPSTREAM_TAG "830d2c4537c7396e14a02a46fbddd18b5d107c65")
  set(_qutlass_fc_root "${FETCHCONTENT_BASE_DIR}")
  if(NOT _qutlass_fc_root)
@@ -55,11 +55,7 @@ message(STATUS "[QUTLASS] QuTLASS is available at ${qutlass_SOURCE_DIR}")
 if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
  cuda_archs_loose_intersection(QUTLASS_SM120_ARCHS "12.0f" "${CUDA_ARCHS}")
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.4)
+  cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0f" "${CUDA_ARCHS}")
    cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0f;10.7f" "${CUDA_ARCHS}")
  else()
    cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0f" "${CUDA_ARCHS}")
  endif()
 else()
  cuda_archs_loose_intersection(QUTLASS_SM120_ARCHS "12.0a;12.1a" "${CUDA_ARCHS}")
  cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0a;10.3a" "${CUDA_ARCHS}")
@@ -129,6 +125,8 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
    CUDA_ARCHS "${QUTLASS_ARCHS}"
  )
  # QuTLASS uses legacy ATen headers and cannot be built with TORCH_TARGET_VERSION.
  # Keep it as its own extension (registers torch.ops._qutlass_C).
  define_extension_target(
    _qutlass_C
    DESTINATION vllm
@@ -141,11 +139,9 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
    WITH_SOABI)
  target_compile_definitions(_qutlass_C PRIVATE
-    QUTLASS_MINIMAL_BUILD=1
+    QUTLASS_DISABLE_PYBIND=1
    TARGET_CUDA_ARCH=${QUTLASS_TARGET_CC}
-    CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1
+    CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
    TORCH_TARGET_VERSION=0x020B000000000000ULL
    USE_CUDA)
  set_property(SOURCE ${QUTLASS_SOURCES} APPEND PROPERTY COMPILE_OPTIONS
    $<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr --use_fast_math -O3>
@@ -1,50 +0,0 @@
 include(FetchContent)
 if(DEFINED ENV{TML_FA4_SRC_DIR})
  set(TML_FA4_SRC_DIR $ENV{TML_FA4_SRC_DIR})
 endif()
 if(TML_FA4_SRC_DIR)
  FetchContent_Declare(
    tml_fa4
    SOURCE_DIR ${TML_FA4_SRC_DIR}
    CONFIGURE_COMMAND ""
    BUILD_COMMAND "")
 else()
  FetchContent_Declare(
    tml_fa4
    GIT_REPOSITORY https://github.com/vllm-project/tml-fa4.git
    GIT_TAG b206834606ed5b5f21f8eed6b0683f528ea9cf7d
    GIT_PROGRESS TRUE
    CONFIGURE_COMMAND ""
    BUILD_COMMAND "")
 endif()
 FetchContent_GetProperties(tml_fa4)
 if(NOT tml_fa4_POPULATED)
  FetchContent_Populate(tml_fa4)
 endif()
 message(STATUS "tml-fa4 is available at ${tml_fa4_SOURCE_DIR}")
 add_custom_target(tml_fa4)
 # Install into a private namespace so this implementation cannot shadow the
 # flash_attn package used by vLLM's standard attention backends.
 install(CODE "
  file(GLOB_RECURSE TML_FA4_PY_FILES
    \"${tml_fa4_SOURCE_DIR}/flash_attn/cute/*.py\")
  foreach(SRC_FILE \${TML_FA4_PY_FILES})
    file(RELATIVE_PATH REL_PATH
      \"${tml_fa4_SOURCE_DIR}/flash_attn/cute\" \${SRC_FILE})
    set(DST_FILE
      \"\${CMAKE_INSTALL_PREFIX}/vllm/third_party/tml_fa4/\${REL_PATH}\")
    get_filename_component(DST_DIR \${DST_FILE} DIRECTORY)
    file(MAKE_DIRECTORY \${DST_DIR})
    file(READ \${SRC_FILE} FILE_CONTENTS)
    string(REPLACE
      \"flash_attn.cute\"
      \"vllm.third_party.tml_fa4\"
      FILE_CONTENTS \"\${FILE_CONTENTS}\")
    file(WRITE \${DST_FILE} \"\${FILE_CONTENTS}\")
  endforeach()
 " COMPONENT tml_fa4)
@@ -39,7 +39,7 @@ else()
  FetchContent_Declare(
          vllm-flash-attn
          GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
-          GIT_TAG ed4b7342bc8f0489dd9b649d5288867e35fc6a32
+          GIT_TAG bb9a72e7dde0dc614ffc663e052cd6a19ce73a42
          GIT_PROGRESS TRUE
          # Don't share the vllm-flash-attn build between build types
          BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
@@ -241,15 +241,14 @@ endmacro()
 # `<major>.<minor>`, dedupes them and then sorts them in ascending order and 
 # stores them in `OUT_ARCHES`.
 #
-# Prefer `code=sm_*`; fall back to `arch=compute_*` for PTX-only flags.
+# Example:
-# This handles mismatches such as `arch=compute_20,code=sm_121`.
+#   CUDA_ARCH_FLAGS="-gencode arch=compute_75,code=sm_75;...;-gencode arch=compute_90a,code=sm_90a" 
 #   extract_unique_cuda_archs_ascending(OUT_ARCHES CUDA_ARCH_FLAGS)
 #   OUT_ARCHES="7.5;...;9.0"
 function(extract_unique_cuda_archs_ascending OUT_ARCHES CUDA_ARCH_FLAGS)
  set(_CUDA_ARCHES)
  foreach(_ARCH ${CUDA_ARCH_FLAGS})
-    string(REGEX MATCH "code=sm_\([0-9]+[af]?\)" _COMPUTE ${_ARCH})
+    string(REGEX MATCH "arch=compute_\([0-9]+[af]?\)" _COMPUTE ${_ARCH})
    if (NOT _COMPUTE)
      string(REGEX MATCH "arch=compute_\([0-9]+[af]?\)" _COMPUTE ${_ARCH})
    endif()
    if (_COMPUTE)
      set(_COMPUTE ${CMAKE_MATCH_1})
    endif()
@@ -397,24 +396,14 @@ function(cuda_archs_loose_intersection OUT_CUDA_ARCHS SRC_CUDA_ARCHS TGT_CUDA_AR
  # match — e.g. SRC="12.0f" matches TGT="12.1a" since SM121 is in the SM12x
  # family. The output uses TGT's value to preserve the user's compilation flags.
  set(_CUDA_ARCHS)
  # Resolve exact base matches before family fallbacks so a generic entry such
  # as 10.0f cannot consume a 10.7 target that has a 10.7f source entry.
  foreach(_arch ${_SRC_CUDA_ARCHS})
    if(_arch MATCHES "[af]$")
      string(REGEX REPLACE "[af]$" "" _base "${_arch}")
      if("${_base}" IN_LIST _TGT_CUDA_ARCHS)
        list(REMOVE_ITEM _SRC_CUDA_ARCHS "${_arch}")
        list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_base}")
        list(APPEND _CUDA_ARCHS "${_arch}")
      endif()
    endif()
  endforeach()
  foreach(_arch ${_SRC_CUDA_ARCHS})
    if(_arch MATCHES "[af]$")
      list(REMOVE_ITEM _SRC_CUDA_ARCHS "${_arch}")
      string(REGEX REPLACE "[af]$" "" _base "${_arch}")
-      if("${_base}a" IN_LIST _TGT_CUDA_ARCHS)
+      if ("${_base}" IN_LIST TGT_CUDA_ARCHS)
        list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_base}")
        list(APPEND _CUDA_ARCHS "${_arch}")
      elseif("${_base}a" IN_LIST _TGT_CUDA_ARCHS)
        list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_base}a")
        list(APPEND _CUDA_ARCHS "${_base}a")
      elseif("${_base}f" IN_LIST _TGT_CUDA_ARCHS)
@@ -498,7 +487,7 @@ endfunction()
 function(cuda_archs_sm90plus OUT_CUDA_ARCHS TGT_CUDA_ARCHS)
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(_archs "9.0a;10.0f;10.7f;11.0f;12.0f" "${TGT_CUDA_ARCHS}")
+    cuda_archs_loose_intersection(_archs "9.0a;10.0f;11.0f;12.0f" "${TGT_CUDA_ARCHS}")
  else()
    cuda_archs_loose_intersection(_archs "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${TGT_CUDA_ARCHS}")
  endif()
@@ -67,8 +67,9 @@ void cp_gather_and_upconvert_fp8_kv_cache(
    torch::Tensor const& src_cache,         // [NUM_BLOCKS, BLOCK_SIZE, 656]
    torch::Tensor const& dst,               // [TOT_TOKENS, 576]
    torch::Tensor const& block_table,       // [BATCH, BLOCK_INDICES]
    torch::Tensor const& seq_lens,          // [BATCH]
    torch::Tensor const& workspace_starts,  // [BATCH]
-    int64_t batch_size, std::optional<torch::Tensor> seq_starts = std::nullopt);
+    int64_t batch_size);
 // Indexer K quantization and cache function
 void indexer_k_quant_and_cache(
@@ -172,15 +172,4 @@
 #endif  // __riscv_v
 // Power VSX
 #ifdef __powerpc__
  // FP32Vec16::exp() in cpu_types_vsx.hpp delegates to FP32Vec8::exp(), which
  // implements a vectorised 5-term minimax polynomial using VSX intrinsics.
  #define DEFINE_FAST_EXP                                                     \
    auto fast_exp = [&](const vec_op::FP32Vec16& vec)                         \
                        __attribute__((always_inline)) { return vec.exp(); }; \
    auto fast_exp_f16 = fast_exp;
 #endif  // __powerpc__
 #endif
@@ -30,8 +30,7 @@ torch::Tensor get_scheduler_metadata(
    const torch::Tensor& query_start_loc, const bool causal,
    const int64_t window_size, const std::string& isa_hint,
    const bool enable_kv_split,
-    const std::optional<torch::Tensor>& dynamic_causal,
+    const std::optional<torch::Tensor>& dynamic_causal) {
    const std::string& kv_cache_dtype) {
  cpu_attention::ISA isa;
  if (isa_hint == "amx") {
    isa = cpu_attention::ISA::AMX;
@@ -66,11 +65,9 @@ torch::Tensor get_scheduler_metadata(
  input.dynamic_causal =
      dynamic_causal.has_value() ? dynamic_causal->data_ptr<bool>() : nullptr;
  const int64_t kv_cache_idx =
      static_cast<int64_t>(parse_fp8_kv_dtype(kv_cache_dtype));
  VLLM_DISPATCH_FLOATING_TYPES(dtype, "get_scheduler_metadata", [&]() {
-    CPU_ATTN_DISPATCH(head_dim, isa, kv_cache_idx, [&]() {
+    CPU_ATTN_DISPATCH(head_dim, isa, 0, [&]() {
-      input.elem_size = sizeof(attn_impl::kv_cache_t);
+      input.elem_size = sizeof(scalar_t);
      input.q_buffer_elem_size = sizeof(attn_impl::q_buffer_t);
      input.logits_buffer_elem_size = sizeof(attn_impl::logits_buffer_t);
      input.output_buffer_elem_size =
@@ -102,9 +102,7 @@ class TileGemm82 {
    kv_cache_t* __restrict__ curr_b = b_tile;
    for (int32_t k = 0; k < dynamic_k_size; ++k) {
-      auto fp32_b_regs = load_b_pair_vec(curr_b);
+      auto [fp32_b_0_reg, fp32_b_1_reg] = load_b_pair_vec(curr_b);
      auto fp32_b_0_reg = fp32_b_regs.first;
      auto fp32_b_1_reg = fp32_b_regs.second;
      float* __restrict__ curr_m_a = curr_a;
      vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
@@ -1,6 +1,5 @@
 #include "cpu/cpu_types.hpp"
 #include "cpu/utils.hpp"
 #include "cpu/cpu_fused_moe_activations.hpp"
 #include "cpu/micro_gemm/cpu_micro_gemm_vec.hpp"
 #include "cpu/cpu_arch_macros.h"
@@ -44,9 +43,193 @@
  }()
 namespace {
 enum class FusedMOEAct {
  SiluAndMul,
  SwigluOAIAndMul,
  GeluAndMul,
  GeluTanhAndMul,
 };
-using cpu_fused_moe_utils::apply_gated_act;
+FusedMOEAct get_act_type(const std::string& act) {
-using cpu_fused_moe_utils::FusedMOEAct;
+  if (act == "silu") {
    return FusedMOEAct::SiluAndMul;
  } else if (act == "swigluoai") {
    return FusedMOEAct::SwigluOAIAndMul;
  } else if (act == "gelu") {
    return FusedMOEAct::GeluAndMul;
  } else if (act == "gelu_tanh") {
    return FusedMOEAct::GeluTanhAndMul;
  } else {
    TORCH_CHECK(false, "Invalid act type: " + act);
  }
 }
 template <typename scalar_t>
 void swigluoai_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
                       const int32_t m_size, const int32_t n_size,
                       const int32_t input_stride,
                       const int32_t output_stride) {
  using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
 #if !defined(__aarch64__)
  // For GPT-OSS interleaved gate-up weights
  alignas(64) static int32_t index[16] = {0,  2,  4,  6,  8,  10, 12, 14,
                                          16, 18, 20, 22, 24, 26, 28, 30};
  vec_op::INT32Vec16 index_vec(index);
 #endif
  vec_op::FP32Vec16 gate_up_max_vec(7.0);
  vec_op::FP32Vec16 up_min_vec(-7.0);
  vec_op::FP32Vec16 alpha_vec(1.702);
  vec_op::FP32Vec16 one_vec(1.0);
  DEFINE_FAST_EXP
  for (int32_t m = 0; m < m_size; ++m) {
    for (int32_t n = 0; n < n_size; n += 32) {
      // Note: AdvSIMD does not support gather loads
 #if defined(__aarch64__)
      vec_op::FP32Vec16 gate_vec(vec_op::uninit);
      vec_op::FP32Vec16 up_vec(vec_op::uninit);
      vec_op::FP32Vec16::load_even_odd(input + n, gate_vec, up_vec);
 #else
      vec_op::FP32Vec16 gate_vec(input + n, index_vec);
      vec_op::FP32Vec16 up_vec(input + n + 1, index_vec);
 #endif
      gate_vec = gate_vec.min(gate_up_max_vec);
      up_vec = up_vec.clamp(up_min_vec, gate_up_max_vec);
      auto sigmoid_vec = one_vec / (one_vec + fast_exp(-gate_vec * alpha_vec));
      auto glu = gate_vec * sigmoid_vec;
      auto gated_output_fp32 = (one_vec + up_vec) * glu;
      scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
      gated_output.save(output + n / 2);
    }
    input += input_stride;
    output += output_stride;
  }
 }
 template <typename scalar_t>
 void silu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
                  const int32_t m_size, const int32_t n_size,
                  const int32_t input_stride, const int32_t output_stride) {
  using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
  const int32_t dim = n_size / 2;
  float* __restrict__ gate = input;
  float* __restrict__ up = input + dim;
  vec_op::FP32Vec16 one_vec(1.0);
  DEFINE_FAST_EXP
  for (int32_t m = 0; m < m_size; ++m) {
    for (int32_t n = 0; n < dim; n += 16) {
      vec_op::FP32Vec16 gate_vec(gate + n);
      vec_op::FP32Vec16 up_vec(up + n);
      auto sigmoid_vec = one_vec / (one_vec + fast_exp(-gate_vec));
      auto silu = gate_vec * sigmoid_vec;
      auto gated_output_fp32 = up_vec * silu;
      scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
      gated_output.save(output + n);
    }
    gate += input_stride;
    up += input_stride;
    output += output_stride;
  }
 }
 template <typename scalar_t>
 void gelu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
                  const int32_t m_size, const int32_t n_size,
                  const int32_t input_stride, const int32_t output_stride) {
  using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
  const int32_t dim = n_size / 2;
  float* __restrict__ gate = input;
  float* __restrict__ up = input + dim;
  vec_op::FP32Vec16 one_vec(1.0);
  vec_op::FP32Vec16 w1_vec(M_SQRT1_2);
  vec_op::FP32Vec16 w2_vec(0.5);
  alignas(64) float temp[16];
  DEFINE_FAST_EXP
  for (int32_t m = 0; m < m_size; ++m) {
    for (int32_t n = 0; n < dim; n += 16) {
      vec_op::FP32Vec16 gate_vec(gate + n);
      vec_op::FP32Vec16 up_vec(up + n);
      auto er_input_vec = gate_vec * w1_vec;
      er_input_vec.save(temp);
      for (int32_t i = 0; i < 16; ++i) {
        temp[i] = std::erf(temp[i]);
      }
      vec_op::FP32Vec16 er_vec(temp);
      auto gelu = gate_vec * w2_vec * (one_vec + er_vec);
      auto gated_output_fp32 = up_vec * gelu;
      scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
      gated_output.save(output + n);
    }
    gate += input_stride;
    up += input_stride;
    output += output_stride;
  }
 }
 template <typename scalar_t>
 void gelu_tanh_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
                       const int32_t m_size, const int32_t n_size,
                       const int32_t input_stride,
                       const int32_t output_stride) {
  using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
  const int32_t dim = n_size / 2;
  float* __restrict__ gate = input;
  float* __restrict__ up = input + dim;
  vec_op::FP32Vec16 one_vec(1.0);
  vec_op::FP32Vec16 w1_vec(0.7978845608028654);
  vec_op::FP32Vec16 w2_vec(0.5);
  vec_op::FP32Vec16 w3_vec(0.044715);
  for (int32_t m = 0; m < m_size; ++m) {
    for (int32_t n = 0; n < dim; n += 16) {
      vec_op::FP32Vec16 gate_vec(gate + n);
      vec_op::FP32Vec16 up_vec(up + n);
      auto gate_pow3_vec = gate_vec * gate_vec * gate_vec;
      auto inner_vec = w1_vec * (gate_vec + w3_vec * gate_pow3_vec);
      // Note: can't use fast_exp form because diffusiongemma will generate
      // wrong results
      auto tanh_vec = inner_vec.tanh();
      auto gelu_tanh = gate_vec * w2_vec * (one_vec + tanh_vec);
      auto gated_output_fp32 = up_vec * gelu_tanh;
      scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
      gated_output.save(output + n);
    }
    gate += input_stride;
    up += input_stride;
    output += output_stride;
  }
 }
 template <typename scalar_t>
 FORCE_INLINE void apply_gated_act(const FusedMOEAct act,
                                  float* __restrict__ input,
                                  scalar_t* __restrict__ output,
                                  const int32_t m, const int32_t n,
                                  const int32_t input_stride,
                                  const int32_t output_stride) {
  switch (act) {
    case FusedMOEAct::SwigluOAIAndMul:
      swigluoai_and_mul(input, output, m, n, input_stride, output_stride);
      return;
    case FusedMOEAct::SiluAndMul:
      silu_and_mul(input, output, m, n, input_stride, output_stride);
      return;
    case FusedMOEAct::GeluAndMul:
      gelu_and_mul(input, output, m, n, input_stride, output_stride);
      return;
    case FusedMOEAct::GeluTanhAndMul:
      gelu_tanh_and_mul(input, output, m, n, input_stride, output_stride);
      return;
    default:
      TORCH_CHECK(false, "Unsupported act type.");
  }
 }
 template <typename scalar_t, typename gemm_t>
 void prepack_moe_weight_impl(scalar_t* __restrict__ weight_ptr,
@@ -634,7 +817,6 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input,
    }
  }
 }
 }  // namespace
 void prepack_moe_weight(
@@ -682,7 +864,7 @@ void cpu_fused_moe(
  const int32_t input_size_2 = w2.size(2);
  const int32_t output_size_2 = w2.size(1);
  const int32_t topk_num = topk_id.size(1);
-  const FusedMOEAct act_type = cpu_fused_moe_utils::get_act_type(act);
+  const FusedMOEAct act_type = get_act_type(act);
  cpu_utils::ISA isa_type = cpu_utils::get_isa(isa);
  TORCH_CHECK(!skip_weighted || topk_num == 1,
              "skip_weighted is only supported for topk=1 on CPU");
@@ -1,204 +0,0 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 #ifndef CPU_FUSED_MOE_ACTIVATIONS_HPP
 #define CPU_FUSED_MOE_ACTIVATIONS_HPP
 #include <cmath>
 #include <cstdint>
 #include <string>
 #include "cpu/cpu_arch_macros.h"
 #include "cpu/utils.hpp"
 namespace cpu_fused_moe_utils {
 enum class FusedMOEAct {
  SiluAndMul,
  SwigluOAIAndMul,
  GeluAndMul,
  GeluTanhAndMul,
 };
 inline FusedMOEAct get_act_type(const std::string& act) {
  if (act == "silu") {
    return FusedMOEAct::SiluAndMul;
  } else if (act == "swigluoai") {
    return FusedMOEAct::SwigluOAIAndMul;
  } else if (act == "gelu") {
    return FusedMOEAct::GeluAndMul;
  } else if (act == "gelu_tanh") {
    return FusedMOEAct::GeluTanhAndMul;
  } else {
    TORCH_CHECK(false, "Invalid act type: " + act);
  }
 }
 template <typename scalar_t>
 void swigluoai_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
                       const int32_t m_size, const int32_t n_size,
                       const int32_t input_stride,
                       const int32_t output_stride) {
  using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
 #if !defined(__aarch64__)
  // For GPT-OSS interleaved gate-up weights
  alignas(64) static int32_t index[16] = {0,  2,  4,  6,  8,  10, 12, 14,
                                          16, 18, 20, 22, 24, 26, 28, 30};
  vec_op::INT32Vec16 index_vec(index);
 #endif
  vec_op::FP32Vec16 gate_up_max_vec(7.0);
  vec_op::FP32Vec16 up_min_vec(-7.0);
  vec_op::FP32Vec16 alpha_vec(1.702);
  vec_op::FP32Vec16 one_vec(1.0);
  DEFINE_FAST_EXP
  for (int32_t m = 0; m < m_size; ++m) {
    for (int32_t n = 0; n < n_size; n += 32) {
      // Note: AdvSIMD does not support gather loads
 #if defined(__aarch64__)
      vec_op::FP32Vec16 gate_vec(vec_op::uninit);
      vec_op::FP32Vec16 up_vec(vec_op::uninit);
      vec_op::FP32Vec16::load_even_odd(input + n, gate_vec, up_vec);
 #else
      vec_op::FP32Vec16 gate_vec(input + n, index_vec);
      vec_op::FP32Vec16 up_vec(input + n + 1, index_vec);
 #endif
      gate_vec = gate_vec.min(gate_up_max_vec);
      up_vec = up_vec.clamp(up_min_vec, gate_up_max_vec);
      auto sigmoid_vec = one_vec / (one_vec + fast_exp(-gate_vec * alpha_vec));
      auto glu = gate_vec * sigmoid_vec;
      auto gated_output_fp32 = (one_vec + up_vec) * glu;
      scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
      gated_output.save(output + n / 2);
    }
    input += input_stride;
    output += output_stride;
  }
 }
 template <typename scalar_t>
 void silu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
                  const int32_t m_size, const int32_t n_size,
                  const int32_t input_stride, const int32_t output_stride) {
  using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
  const int32_t dim = n_size / 2;
  float* __restrict__ gate = input;
  float* __restrict__ up = input + dim;
  vec_op::FP32Vec16 one_vec(1.0);
  DEFINE_FAST_EXP
  for (int32_t m = 0; m < m_size; ++m) {
    for (int32_t n = 0; n < dim; n += 16) {
      vec_op::FP32Vec16 gate_vec(gate + n);
      vec_op::FP32Vec16 up_vec(up + n);
      auto sigmoid_vec = one_vec / (one_vec + fast_exp(-gate_vec));
      auto silu = gate_vec * sigmoid_vec;
      auto gated_output_fp32 = up_vec * silu;
      scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
      gated_output.save(output + n);
    }
    gate += input_stride;
    up += input_stride;
    output += output_stride;
  }
 }
 template <typename scalar_t>
 void gelu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
                  const int32_t m_size, const int32_t n_size,
                  const int32_t input_stride, const int32_t output_stride) {
  using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
  const int32_t dim = n_size / 2;
  float* __restrict__ gate = input;
  float* __restrict__ up = input + dim;
  vec_op::FP32Vec16 one_vec(1.0);
  vec_op::FP32Vec16 w1_vec(M_SQRT1_2);
  vec_op::FP32Vec16 w2_vec(0.5);
  alignas(64) float temp[16];
  DEFINE_FAST_EXP
  for (int32_t m = 0; m < m_size; ++m) {
    for (int32_t n = 0; n < dim; n += 16) {
      vec_op::FP32Vec16 gate_vec(gate + n);
      vec_op::FP32Vec16 up_vec(up + n);
      auto er_input_vec = gate_vec * w1_vec;
      er_input_vec.save(temp);
      for (int32_t i = 0; i < 16; ++i) {
        temp[i] = std::erf(temp[i]);
      }
      vec_op::FP32Vec16 er_vec(temp);
      auto gelu = gate_vec * w2_vec * (one_vec + er_vec);
      auto gated_output_fp32 = up_vec * gelu;
      scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
      gated_output.save(output + n);
    }
    gate += input_stride;
    up += input_stride;
    output += output_stride;
  }
 }
 template <typename scalar_t>
 void gelu_tanh_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
                       const int32_t m_size, const int32_t n_size,
                       const int32_t input_stride,
                       const int32_t output_stride) {
  using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
  const int32_t dim = n_size / 2;
  float* __restrict__ gate = input;
  float* __restrict__ up = input + dim;
  vec_op::FP32Vec16 one_vec(1.0);
  vec_op::FP32Vec16 w1_vec(0.7978845608028654);
  vec_op::FP32Vec16 w2_vec(0.5);
  vec_op::FP32Vec16 w3_vec(0.044715);
  for (int32_t m = 0; m < m_size; ++m) {
    for (int32_t n = 0; n < dim; n += 16) {
      vec_op::FP32Vec16 gate_vec(gate + n);
      vec_op::FP32Vec16 up_vec(up + n);
      auto gate_pow3_vec = gate_vec * gate_vec * gate_vec;
      auto inner_vec = w1_vec * (gate_vec + w3_vec * gate_pow3_vec);
      // Note: can't use fast_exp form because diffusiongemma will generate
      // wrong results
      auto tanh_vec = inner_vec.tanh();
      auto gelu_tanh = gate_vec * w2_vec * (one_vec + tanh_vec);
      auto gated_output_fp32 = up_vec * gelu_tanh;
      scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
      gated_output.save(output + n);
    }
    gate += input_stride;
    up += input_stride;
    output += output_stride;
  }
 }
 template <typename scalar_t>
 FORCE_INLINE void apply_gated_act(const FusedMOEAct act,
                                  float* __restrict__ input,
                                  scalar_t* __restrict__ output,
                                  const int32_t m, const int32_t n,
                                  const int32_t input_stride,
                                  const int32_t output_stride) {
  switch (act) {
    case FusedMOEAct::SwigluOAIAndMul:
      swigluoai_and_mul(input, output, m, n, input_stride, output_stride);
      return;
    case FusedMOEAct::SiluAndMul:
      silu_and_mul(input, output, m, n, input_stride, output_stride);
      return;
    case FusedMOEAct::GeluAndMul:
      gelu_and_mul(input, output, m, n, input_stride, output_stride);
      return;
    case FusedMOEAct::GeluTanhAndMul:
      gelu_tanh_and_mul(input, output, m, n, input_stride, output_stride);
      return;
    default:
      TORCH_CHECK(false, "Unsupported act type.");
  }
 }
 }  // namespace cpu_fused_moe_utils
 #endif
--- a/Show More
+++ b/Show More
		`@@ -0,0 +1,2 @@`
							`collect_env.py`
							`vllm/model_executor/layers/fla/ops/*.py`