[Attention][TokenSpeed MLA] Also warm up prefill kernel from decode impl

The prefill backend may be paired with flash_attn / trtllm in production — in that case the prefill backend's __init__ never runs and the prefill kernel's first call pays a 1.5–2 minute JIT cost. Add the same idempotent `warmup_compile_prefill` invocation to TokenspeedMLAImpl.__init__ (the decode-side backend, always present when tokenspeed is selected). The function dedupes by config key, so the double call is a no-op when both backends are tokenspeed. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
update version
2026-05-06 07:45:16 +00:00 · 2026-05-06 07:28:20 +00:00 · 2026-05-06 07:20:48 +00:00 · 2026-05-06 03:14:49 +00:00 · 2026-05-06 02:53:46 +00:00 · 2026-05-06 02:49:09 +00:00
1951 changed files with 62332 additions and 240630 deletions
@@ -8,7 +8,6 @@ run_all_patterns:
  - "CMakeLists.txt"
  - "requirements/common.txt"
  - "requirements/cuda.txt"
-  - "requirements/kv_connectors.txt"
  - "requirements/build/cuda.txt"
  - "requirements/test/cuda.txt"
  - "setup.py"
@@ -17,26 +17,6 @@ steps:
      --target test
      --no-cache
      --progress plain .
-    - |
-      docker run --rm --network=none --entrypoint /bin/bash "rocm/vllm-ci:${BUILDKITE_COMMIT}" -ec '
-        if [ ! -d /vllm-workspace ]; then echo Missing directory: /vllm-workspace >&2; exit 1; fi
-        if [ ! -d /vllm-workspace/tests ]; then echo Missing directory: /vllm-workspace/tests >&2; exit 1; fi
-        if [ ! -d /vllm-workspace/src/vllm ]; then echo Missing directory: /vllm-workspace/src/vllm >&2; exit 1; fi
-        if [ ! -x /vllm-workspace/src/vllm/vllm-rs ]; then echo Missing executable: /vllm-workspace/src/vllm/vllm-rs >&2; exit 1; fi
-        command -v python3
-        command -v uv
-        command -v pytest
-        if ! command -v amd-smi >/dev/null 2>&1 && ! command -v rocminfo >/dev/null 2>&1; then
-          echo No ROCm CLI found in image >&2
-          exit 1
-        fi
-        python3 - <<PY
-      import torch, vllm
-      print(torch.__version__)
-      print(vllm.__version__)
-      PY
-        echo AMD image smoke OK
-      '
    - docker push "rocm/vllm-ci:${BUILDKITE_COMMIT}"
    env:
      DOCKER_BUILDKIT: "1"
@@ -12,21 +12,15 @@ steps:
  - vllm/_custom_ops.py
  - tests/kernels/attention/test_cpu_attn.py
  - tests/kernels/moe/test_cpu_fused_moe.py
-  - tests/kernels/moe/test_cpu_quant_fused_moe.py
  - tests/kernels/test_onednn.py
  - tests/kernels/test_awq_int4_to_int8.py
-  - tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
-  - tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
  commands:
    - |
-      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
+      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 20m "
      pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
      pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py
-      pytest -x -v -s tests/kernels/moe/test_cpu_quant_fused_moe.py
      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/test_awq_int4_to_int8.py"

 - label: CPU-Compatibility Tests
  depends_on: []
@@ -56,41 +50,22 @@ steps:
      pytest -x -v -s tests/models/language/generation -m cpu_model
      pytest -x -v -s tests/models/language/pooling -m cpu_model"

- label: CPU-ModelRunnerV2 Tests
-  depends_on: []
-  device: intel_cpu
-  no_plugin: true
-  soft_fail: true
-  source_file_dependencies:
-  - vllm/v1/worker/cpu/
-  - vllm/v1/worker/gpu/
-  - vllm/v1/sample/ops/topk_topp_triton.py
-  - vllm/v1/sample/ops/topk_topp_sampler.py
-  - tests/v1/sample/test_topk_topp_sampler.py
-  commands:
-    - |
-      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 45m "
-      uv pip install git+https://github.com/triton-lang/triton-cpu.git@270e696d
-      VLLM_USE_V2_MODEL_RUNNER=1 pytest -x -v -s tests/models/language/generation/test_granite.py -m cpu_model
-      # TODO: move to CPU-Kernel Tests once triton-cpu has a pre-built wheel
-      pytest -x -v -s tests/v1/sample/test_topk_topp_sampler.py::TestTritonTopkTopp"
-
 - label: CPU-Quantization Model Tests
  depends_on: []
  device: intel_cpu
  no_plugin: true
  source_file_dependencies:
  - csrc/cpu/
-  - vllm/model_executor/layers/quantization/auto_gptq.py
+  - vllm/model_executor/layers/quantization/cpu_wna16.py
+  - vllm/model_executor/layers/quantization/gptq_marlin.py
  - vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_int8.py
-  - vllm/model_executor/kernels/linear/mixed_precision/cpu.py
-  - vllm/model_executor/kernels/linear/scaled_mm/cpu.py
-  - vllm/model_executor/layers/fused_moe/experts/cpu_moe.py
+  - vllm/model_executor/layers/quantization/kernels/scaled_mm/cpu.py
+  - vllm/model_executor/layers/quantization/kernels/mixed_precision/cpu.py
  - tests/quantization/test_compressed_tensors.py
  - tests/quantization/test_cpu_wna16.py
  commands:
    - |
-      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
+      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 20m "
      pytest -x -v -s tests/quantization/test_compressed_tensors.py::test_compressed_tensors_w8a8_logprobs
      pytest -x -v -s tests/quantization/test_cpu_wna16.py"
      
@@ -8,3 +8,10 @@ steps:
    commands: 
    - bash .buildkite/scripts/hardware_ci/run-hpu-test.sh

+  - label: "Intel GPU Test"
+    depends_on: []
+    soft_fail: true
+    device: intel_gpu
+    no_plugin: true
+    commands: 
+    - bash .buildkite/scripts/hardware_ci/run-xpu-test.sh
@@ -13,60 +13,6 @@ steps:
        - exit_status: -10  # Agent was lost
          limit: 2

-  - label: ":docker: :smoking: Non-root smoke tests"
-    key: image-build-smoke-test
-    depends_on:
-      - image-build
-    commands:
-    # Smoke 1: the default (root) image must still be importable
-    # under a non-root UID via `--user 2000:0`. Validates the `vllm` passwd
-    # entry + group-0-writable /home/vllm + uv path cleanup from #31959.
-    # Uses `import vllm` rather than `vllm serve --help` because the latter
-    # instantiates `VllmConfig` which requires a GPU attached to the
-    # container.
-    - docker run --rm --user 2000:0 --entrypoint python3 "$IMAGE_TAG" -c "import vllm; print(vllm.__version__)"
-    # Smoke 2: assert the non-root enabling invariants are baked
-    # into the image. Runs as UID 2000:0 via a shell so we can verify
-    # filesystem perms + passwd/group file state + wrapper presence without
-    # triggering vLLM's GPU-requiring config-init path. The opt-in
-    # `vllm-openai-nonroot` target adds only `USER vllm`, `WORKDIR
-    # /home/vllm`, and an `ENTRYPOINT` override on top of these invariants;
-    # its build correctness is reviewed at the Dockerfile level. Wrapper
-    # logic is covered separately by the pre-commit hook
-    # `test-nonroot-entrypoint` (see .pre-commit-config.yaml).
-    - |
-      docker run --rm --user 2000:0 --entrypoint /bin/sh "$IMAGE_TAG" -ec '
-        if ! getent passwd 2000 | grep -q ^vllm:; then
-          echo FAIL: UID 2000 != vllm
-          exit 1
-        fi
-        if ! id -gn 2>/dev/null | grep -qx root; then
-          echo FAIL: GID 0 not root group
-          exit 1
-        fi
-        touch /home/vllm/.smoke && rm /home/vllm/.smoke
-        touch /opt/uv/cache/.smoke && rm /opt/uv/cache/.smoke
-        if ! test -x /usr/local/bin/vllm-nonroot-entrypoint.sh; then
-          echo FAIL: wrapper missing
-          exit 1
-        fi
-        if ! test -w /etc/passwd; then
-          echo FAIL: /etc/passwd not group-writable
-          exit 1
-        fi
-        if ! test -w /etc/group; then
-          echo FAIL: /etc/group not group-writable
-          exit 1
-        fi
-        echo non-root invariants OK
-      '
-    retry:
-      automatic:
-        - exit_status: -1  # Agent was lost
-          limit: 2
-        - exit_status: -10  # Agent was lost
-          limit: 2
-
  - label: ":docker: Build CPU image"
    key: image-build-cpu
    depends_on: []
@@ -110,21 +56,3 @@ steps:
          limit: 2
        - exit_status: -10  # Agent was lost
          limit: 2
-
-  - label: ":docker: Build arm64 image"
-    key: arm64-image-build
-    depends_on: []
-    source_file_dependencies:
-      - ".buildkite/image_build/image_build.yaml"
-      - ".buildkite/image_build/image_build_arm64.sh"
-      - "docker/Dockerfile"
-    commands:
-    - .buildkite/image_build/image_build_arm64.sh $REGISTRY $REPO $BUILDKITE_COMMIT
-    env:
-      DOCKER_BUILDKIT: "1"
-    retry:
-      automatic:
-        - exit_status: -1  # Agent was lost
-          limit: 2
-        - exit_status: -10  # Agent was lost
-          limit: 2
@@ -1,37 +0,0 @@
-#!/bin/bash
-set -e
-
-if [[ $# -lt 3 ]]; then
-  echo "Usage: $0 <registry> <repo> <commit>"
-  exit 1
-fi
-
-REGISTRY=$1
-REPO=$2
-BUILDKITE_COMMIT=$3
-
-# authenticate with AWS ECR
-aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" || true
-
-# skip build if image already exists
-if [[ -z $(docker manifest inspect "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-arm64) ]]; then
-  echo "Image not found, proceeding with build..."
-else
-  echo "Image found"
-  exit 0
-fi
-
-# build (Grace/GH200 is the arm64 GPU target; sm_90)
-docker build --file docker/Dockerfile \
-  --platform linux/arm64 \
-  --build-arg max_jobs=16 \
-  --build-arg nvcc_threads=4 \
-  --build-arg torch_cuda_arch_list="9.0" \
-  --build-arg USE_SCCACHE=1 \
-  --build-arg buildkite_commit="$BUILDKITE_COMMIT" \
-  --tag "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-arm64 \
-  --target test \
-  --progress plain .
-
-# push
-docker push "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-arm64
@@ -11,7 +11,7 @@ REPO=$2
 BUILDKITE_COMMIT=$3

 # authenticate with AWS ECR
-aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" || true
+aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY"

 # skip build if image already exists
 if [[ -z $(docker manifest inspect "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-hpu) ]]; then
@@ -11,8 +11,8 @@ REPO=$2
 BUILDKITE_COMMIT=$3

 # authenticate with AWS ECR
-aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" || true
-aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com || true
+aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY"
+aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com

 # skip build if image already exists
 if ! docker manifest inspect "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-xpu &> /dev/null; then
@@ -18,18 +18,17 @@ steps:
    - >-
      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
      'cd tests &&
-      export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
      pytest -v -s lora/test_layers.py &&
      pytest -v -s lora/test_lora_checkpoints.py &&
-      pytest -v -s lora/test_lora_functions.py &&
+      (pytest -v -s lora/test_lora_functions.py --deselect="tests/lora/test_lora_functions.py::test_lora_functions_sync" --deselect="tests/lora/test_lora_functions.py::test_lora_functions_async" || true) &&
      pytest -v -s lora/test_lora_huggingface.py &&
      pytest -v -s lora/test_lora_manager.py &&
      pytest -v -s lora/test_lora_utils.py &&
      pytest -v -s lora/test_peft_helper.py &&
      pytest -v -s lora/test_resolver.py &&
      pytest -v -s lora/test_utils.py &&
-      pytest -v -s lora/test_add_lora.py  &&
-      pytest -v -s lora/test_worker.py'
+      (pytest -v -s lora/test_add_lora.py --deselect="tests/lora/test_add_lora.py::test_add_lora" || true) &&
+      (pytest -v -s lora/test_worker.py --deselect="tests/lora/test_worker.py::test_worker_apply_lora" || true)'

 - label: LoRA Fused/MoE Kernels
  timeout_in_minutes: 45
@@ -47,9 +46,8 @@ steps:
    - >-
      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
      'cd tests &&
-      export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
      pytest -v -s lora/test_fused_moe_lora_kernel.py && 
-      pytest -v -s lora/test_moe_lora_align_sum.py --deselect="tests/lora/test_moe_lora_align_sum.py::test_moe_lora_align_block_size_mixed_base_and_lora[1]"'
+      pytest -v -s lora/test_moe_lora_align_sum.py'

 - label: LoRA Punica Kernels
  timeout_in_minutes: 45
@@ -67,9 +65,8 @@ steps:
    - >-
      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
      'cd tests &&
-      export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
      set -o pipefail &&
-      pytest -v -s lora/test_punica_ops.py --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype0-3-43264-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype1-1-2049-64-128-16]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-128-1-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-256-1-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-256-8-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[expand-0-xpu:0-dtype0-3-2049-128-8-16]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-128-8-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels[expand-0-xpu:0-dtype1-1-2049-256-128-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype0-3-64256-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype1-2-29696-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype1-3-49408-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype0-2-16384-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype0-2-51328-32-4-4]"'
+      pytest -v -s lora/test_punica_ops.py --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-2-2049-64-32-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype1-2-64000-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-128-1-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-256-1-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-256-8-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[expand-0-xpu:0-dtype0-3-2049-128-8-16]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-128-8-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels[expand-0-xpu:0-dtype1-1-2049-256-128-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype0-3-64256-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype1-2-29696-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype1-3-49408-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype0-2-16384-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype0-2-51328-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype1-1-102656-32-4-4]"'

 - label: LoRA Punica FP8/XPU Ops
  timeout_in_minutes: 45
@@ -87,7 +84,6 @@ steps:
    - >-
      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
      'cd tests &&
-      export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
      pytest -v -s lora/test_punica_ops_fp8.py &&
      pytest -v -s lora/test_punica_xpu_ops.py'

@@ -107,12 +103,10 @@ steps:
    - >-
      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
      'cd tests &&
-      export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
      (pytest -v -s lora/test_mixtral.py --deselect="tests/lora/test_mixtral.py::test_mixtral_lora[4]" || true) &&
      pytest -v -s lora/test_quant_model.py --deselect="tests/lora/test_quant_model.py::test_quant_model_lora[model0]" --deselect="tests/lora/test_quant_model.py::test_quant_model_lora[model1]" --deselect="tests/lora/test_quant_model.py::test_quant_model_tp_equality[model0]" &&
-      pytest -v -s lora/test_transformers_model.py &&
-      pytest -v -s lora/test_chatglm3_tp.py &&
-      pytest -s -v lora/test_minicpmv_tp.py'
+      pytest -v -s lora/test_qwen35_densemodel_lora.py &&
+      pytest -v -s lora/test_transformers_model.py'

 - label: LoRA Multimodal
  timeout_in_minutes: 45
@@ -130,6 +124,6 @@ steps:
    - >-
      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
      'cd tests &&
-      export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
      pytest -v -s lora/test_default_mm_loras.py && 
+      (pytest -v -s lora/test_qwen3_unembed.py || true) &&
      pytest -v -s lora/test_whisper.py'
@@ -49,7 +49,7 @@ steps:
      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
      'export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
      cd tests &&
-      pytest -v -s v1/logits_processors --ignore=v1/logits_processors/test_custom_online.py --ignore=v1/logits_processors/test_custom_offline.py &&
+      pytest -v -s v1/logits_processors &&
      pytest -v -s v1/test_oracle.py &&
      pytest -v -s v1/test_request.py &&
      pytest -v -s v1/test_outputs.py'
@@ -36,12 +36,9 @@ steps:
        python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --attention-backend=TRITON_ATTN &&
        python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --quantization fp8 &&
        python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --kv-cache-dtype fp8 &&
-        python3 examples/basic/offline_inference/generate.py --model nvidia/Llama-3.1-8B-Instruct-FP8 --block-size 64 --enforce-eager --quantization modelopt --kv-cache-dtype fp8 --attention-backend TRITON_ATTN --max-model-len 4096 &&
        python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4 --block-size 64 --enforce-eager --max-model-len 8192 &&
        python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2 &&
-        python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2 --enable-expert-parallel &&
-        python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4 --max-model-len 8192
-        '
+        python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2 --enable-expert-parallel'
  - label: "XPU V1 test"
    depends_on:
      - image-build-xpu
@@ -64,24 +61,5 @@ 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/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 &&
-        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: 30
-    device: intel_gpu
-    no_plugin: true
-    env:
-      REGISTRY: "public.ecr.aws/q9t5s3a7"
-      REPO: "vllm-ci-test-repo"
-    source_file_dependencies:
-      - vllm/
-      - .buildkite/intel_jobs/test-intel.yaml 
-    commands:
-      - >-
-        bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-        'pip install av &&
-        cd tests &&
-        pytest -v -s entrypoints/openai/chat_completion/test_audio_in_video.py &&
-        pytest -v -s benchmarks/test_serve_cli.py'
+        pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_tree_attention.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.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'
@@ -2,7 +2,7 @@
 # We can use this script to compute baseline accuracy on chartqa for vllm.
 #
 # Make sure you have lm-eval-harness installed:
-#   pip install "lm-eval[api]>=0.4.12"
+#   pip install "lm-eval[api]>=0.4.11"

 usage() {
    echo``
@@ -2,7 +2,7 @@
 # We can use this script to compute baseline accuracy on GSM for transformers.
 #
 # Make sure you have lm-eval-harness installed:
-#   pip install "lm-eval[api]>=0.4.12"
+#   pip install "lm-eval[api]>=0.4.11"

 usage() {
    echo``
@@ -3,7 +3,7 @@
 # We use this for fp8, which HF does not support.
 #
 # Make sure you have lm-eval-harness installed:
-#   pip install "lm-eval[api]>=0.4.12"
+#   pip install "lm-eval[api]>=0.4.11"

 usage() {
    echo``
@@ -3,7 +3,7 @@
 # We use this for fp8, which HF does not support.
 #
 # Make sure you have lm-eval-harness installed:
-#   pip install "lm-eval[api]>=0.4.12"
+#   pip install "lm-eval[api]>=0.4.11"

 usage() {
    echo``
@@ -31,9 +31,30 @@
    }
  },
  "tests": [
+    {
+      "test_name": "serving_llama8B_tp1_sharegpt",
+      "server_parameters": {
+        "tensor_parallel_size": 1
+      },
+      "client_parameters": {
+        "dataset_name": "sharegpt",
+        "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json"
+      }
+    },
+    {
+      "test_name": "serving_llama8B_tp2_sharegpt",
+      "server_parameters": {
+        "tensor_parallel_size": 2
+      },
+      "client_parameters": {
+        "dataset_name": "sharegpt",
+        "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json"
+      }
+    },
    {
      "test_name": "serving_llama8B_tp1_random_128_128",
      "server_parameters": {
+        "tensor_parallel_size": 1
      },
      "client_parameters": {
        "dataset_name": "random",
@@ -41,10 +62,158 @@
        "random-output-len": 128
      }
    },
+    {
+      "test_name": "serving_llama8B_tp2_random_128_128",
+      "server_parameters": {
+        "tensor_parallel_size": 2
+      },
+      "client_parameters": {
+        "dataset_name": "random",
+        "random-input-len": 128,
+        "random-output-len": 128
+      }
+    },
+    {
+      "test_name": "serving_llama8B_tp4_random_128_128",
+      "server_parameters": {
+        "tensor_parallel_size": 4
+      },
+      "client_parameters": {
+        "dataset_name": "random",
+        "random-input-len": 128,
+        "random-output-len": 128
+      }
+    },
+    {
+      "test_name": "serving_llama8B_tp1_random_128_2048",
+      "server_parameters": {
+        "tensor_parallel_size": 1
+      },
+      "client_parameters": {
+        "dataset_name": "random",
+        "random-input-len": 128,
+        "random-output-len": 2048
+      }
+    },
+    {
+      "test_name": "serving_llama8B_tp2_random_128_2048",
+      "server_parameters": {
+        "tensor_parallel_size": 2
+      },
+      "client_parameters": {
+        "dataset_name": "random",
+        "random-input-len": 128,
+        "random-output-len": 2048
+      }
+    },
+    {
+      "test_name": "serving_llama8B_tp4_random_128_2048",
+      "server_parameters": {
+        "tensor_parallel_size": 4
+      },
+      "client_parameters": {
+        "dataset_name": "random",
+        "random-input-len": 128,
+        "random-output-len": 2048
+      }
+    },
+    {
+      "test_name": "serving_llama8B_tp1_random_2048_128",
+      "server_parameters": {
+        "tensor_parallel_size": 1
+      },
+      "client_parameters": {
+        "dataset_name": "random",
+        "random-input-len": 2048,
+        "random-output-len": 128
+      }
+    },
+    {
+      "test_name": "serving_llama8B_tp2_random_2048_128",
+      "server_parameters": {
+        "tensor_parallel_size": 2
+      },
+      "client_parameters": {
+        "dataset_name": "random",
+        "random-input-len": 2048,
+        "random-output-len": 128
+      }
+    },
+    {
+      "test_name": "serving_llama8B_tp4_random_2048_128",
+      "server_parameters": {
+        "tensor_parallel_size": 4
+      },
+      "client_parameters": {
+        "dataset_name": "random",
+        "random-input-len": 2048,
+        "random-output-len": 128
+      }
+    },
+    {
+      "test_name": "serving_llama8B_tp1_random_2048_2048",
+      "server_parameters": {
+        "tensor_parallel_size": 1
+      },
+      "client_parameters": {
+        "dataset_name": "random",
+        "random-input-len": 2048,
+        "random-output-len": 2048
+      }
+    },
+    {
+      "test_name": "serving_llama8B_tp2_random_2048_2048",
+      "server_parameters": {
+        "tensor_parallel_size": 2
+      },
+      "client_parameters": {
+        "dataset_name": "random",
+        "random-input-len": 2048,
+        "random-output-len": 2048
+      }
+    },
+    {
+      "test_name": "serving_llama8B_tp4_random_2048_2048",
+      "server_parameters": {
+        "tensor_parallel_size": 4
+      },
+      "client_parameters": {
+        "dataset_name": "random",
+        "random-input-len": 2048,
+        "random-output-len": 2048
+      }
+    },
    {
      "test_name": "serving_llama8B_int4_tp1_random_128_128",
      "server_parameters": {
-        "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4"
+        "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
+        "tensor_parallel_size": 1
+      },
+      "client_parameters": {
+        "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
+        "dataset_name": "random",
+        "random-input-len": 128,
+        "random-output-len": 128
+      }
+    },
+    {
+      "test_name": "serving_llama8B_int4_tp2_random_128_128",
+      "server_parameters": {
+        "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
+        "tensor_parallel_size": 2
+      },
+      "client_parameters": {
+        "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
+        "dataset_name": "random",
+        "random-input-len": 128,
+        "random-output-len": 128
+      }
+    },
+    {
+      "test_name": "serving_llama8B_int4_tp4_random_128_128",
+      "server_parameters": {
+        "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
+        "tensor_parallel_size": 4
      },
      "client_parameters": {
        "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
@@ -56,7 +225,8 @@
    {
      "test_name": "serving_llama8B_int8_tp1_random_128_128",
      "server_parameters": {
-        "model": "RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8"
+        "model": "RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8",
+        "tensor_parallel_size": 1
      },
      "client_parameters": {
        "model": "RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8",
@@ -66,12 +236,26 @@
      }
    },
    {
-      "test_name": "serving_llama1B_tp1_random_128_128",
+      "test_name": "serving_llama8B_int8_tp2_random_128_128",
      "server_parameters": {
-        "model": "meta-llama/Llama-3.2-1B"
+        "model": "RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8",
+        "tensor_parallel_size": 2
      },
      "client_parameters": {
-        "model": "meta-llama/Llama-3.2-1B",
+        "model": "RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8",
+        "dataset_name": "random",
+        "random-input-len": 128,
+        "random-output-len": 128
+      }
+    },
+    {
+      "test_name": "serving_llama8B_int8_tp4_random_128_128",
+      "server_parameters": {
+        "model": "RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8",
+        "tensor_parallel_size": 4
+      },
+      "client_parameters": {
+        "model": "RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8",
        "dataset_name": "random",
        "random-input-len": 128,
        "random-output-len": 128
@@ -80,7 +264,8 @@
    {
      "test_name": "serving_llama3B_tp1_random_128_128",
      "server_parameters": {
-        "model": "meta-llama/Llama-3.2-3B-Instruct"
+        "model": "meta-llama/Llama-3.2-3B-Instruct",
+        "tensor_parallel_size": 1
      },
      "client_parameters": {
        "model": "meta-llama/Llama-3.2-3B-Instruct",
@@ -89,22 +274,11 @@
        "random-output-len": 128
      }
    },
-    {
-      "test_name": "serving_llama70B_tp1_random_128_128",
-      "server_parameters": {
-        "model": "meta-llama/Llama-3.3-70B-Instruct"
-      },
-      "client_parameters": {
-        "model": "meta-llama/Llama-3.3-70B-Instruct",
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 128
-      }
-    },
    {
      "test_name": "serving_granite2B_tp1_random_128_128",
      "server_parameters": {
-        "model": "ibm-granite/granite-3.2-2b-instruct"
+        "model": "ibm-granite/granite-3.2-2b-instruct",
+        "tensor_parallel_size": 1
      },
      "client_parameters": {
        "model": "ibm-granite/granite-3.2-2b-instruct",
@@ -116,7 +290,8 @@
    {
      "test_name": "serving_qwen1.7B_tp1_random_128_128",
      "server_parameters": {
-        "model": "Qwen/Qwen3-1.7B"
+        "model": "Qwen/Qwen3-1.7B",
+        "tensor_parallel_size": 1
      },
      "client_parameters": {
        "model": "Qwen/Qwen3-1.7B",
@@ -128,7 +303,8 @@
    {
      "test_name": "serving_qwen4B_tp1_random_128_128",
      "server_parameters": {
-        "model": "Qwen/Qwen3-4B"
+        "model": "Qwen/Qwen3-4B",
+        "tensor_parallel_size": 1
      },
      "client_parameters": {
        "model": "Qwen/Qwen3-4B",
@@ -140,7 +316,8 @@
    {
      "test_name": "serving_qwen8B_tp1_random_128_128",
      "server_parameters": {
-        "model": "Qwen/Qwen3-8B"
+        "model": "Qwen/Qwen3-8B",
+        "tensor_parallel_size": 1
      },
      "client_parameters": {
        "model": "Qwen/Qwen3-8B",
@@ -149,34 +326,11 @@
        "random-output-len": 128
      }
    },
-    {
-      "test_name": "serving_qwen14B_tp1_random_128_128",
-      "server_parameters": {
-        "model": "Qwen/Qwen3-14B"
-      },
-      "client_parameters": {
-        "model": "Qwen/Qwen3-14B",
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_qwen30B_tp1_random_128_128",
-      "server_parameters": {
-        "model": "Qwen/Qwen3-30B-A3B"
-      },
-      "client_parameters": {
-        "model": "Qwen/Qwen3-30B-A3B",
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 128
-      }
-    },
    {
      "test_name": "serving_glm9B_tp1_random_128_128",
      "server_parameters": {
-        "model": "zai-org/glm-4-9b-hf"
+        "model": "zai-org/glm-4-9b-hf",
+        "tensor_parallel_size": 1
      },
      "client_parameters": {
        "model": "zai-org/glm-4-9b-hf",
@@ -188,7 +342,8 @@
    {
      "test_name": "serving_gemma7B_tp1_random_128_128",
      "server_parameters": {
-        "model": "google/gemma-7b"
+        "model": "google/gemma-7b",
+        "tensor_parallel_size": 1
      },
      "client_parameters": {
        "model": "google/gemma-7b",
@@ -196,94 +351,6 @@
        "random-input-len": 128,
        "random-output-len": 128
      }
-    },
-    {
-      "test_name": "serving_gemma3-4b_tp1_random_128_128",
-      "server_environment_variables": {
-	"VLLM_CPU_SGL_KERNEL": 0
-      },
-      "server_parameters": {
-        "model": "google/gemma-3-4b-it"
-      },
-      "client_parameters": {
-        "model": "google/gemma-3-4b-it",
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_gemma3-12b_tp1_random_128_128",
-      "server_environment_variables": {
-	"VLLM_CPU_SGL_KERNEL": 0
-      },
-      "server_parameters": {
-        "model": "google/gemma-3-12b-it"
-      },
-      "client_parameters": {
-        "model": "google/gemma-3-12b-it",
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_gemma4-4b_tp1_random_128_128",
-      "server_environment_variables": {
-	"VLLM_CPU_SGL_KERNEL": 0
-      },
-      "server_parameters": {
-        "model": "google/gemma-4-E4B-it"
-      },
-      "client_parameters": {
-        "model": "google/gemma-4-E4B-it",
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_gemma4-2b_tp1_random_128_128",
-      "server_environment_variables": {
-	"VLLM_CPU_SGL_KERNEL": 0
-      },
-      "server_parameters": {
-        "model": "google/gemma-4-E2B-it"
-      },
-      "client_parameters": {
-        "model": "google/gemma-4-E2B-it",
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_gemma4-26b_tp1_random_128_128",
-      "server_environment_variables": {
-	"VLLM_CPU_SGL_KERNEL": 0,
-        "VLLM_CPU_ATTN_SPLIT_KV": 0
-      },
-      "server_parameters": {
-        "model": "google/gemma-4-26B-A4B-it"
-      },
-      "client_parameters": {
-        "model": "google/gemma-4-26B-A4B-it",
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_phi4_tp1_random_128_128",
-      "server_parameters": {
-        "model": "microsoft/Phi-4-reasoning"
-      },
-      "client_parameters": {
-        "model": "microsoft/Phi-4-reasoning",
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 128
-      }
    }
  ]
 }
@@ -1,112 +1,77 @@
-{
-  "defaults": {
-    "qps_list": [
-      "inf"
-    ],
-    "max_concurrency_list": [12, 16, 24, 32, 64, 128, 200],
-    "server_parameters": {
-      "model": "meta-llama/Llama-3.1-8B-Instruct",
-      "tensor_parallel_size": 1,
-      "dtype": "bfloat16"
+[
+    {
+        "test_name": "serving_llama8B_tp1_sharegpt",
+        "qps_list": [1, 4, 16, "inf"],
+        "server_parameters": {
+            "model": "meta-llama/Meta-Llama-3.1-8B-Instruct",
+            "tensor_parallel_size": 1,
+            "disable_log_stats": "",
+            "load_format": "dummy"
+        },
+        "client_parameters": {
+            "model": "meta-llama/Meta-Llama-3.1-8B-Instruct",
+            "backend": "vllm",
+            "dataset_name": "sharegpt",
+            "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json",
+            "temperature": 0,
+            "num_prompts": 200
+        }
    },
-    "client_parameters": {
-      "model": "meta-llama/Llama-3.1-8B-Instruct",
-      "backend": "vllm",
-      "ignore-eos": "",
-      "temperature": 0,
-      "num_prompts": 200
+    {
+        "test_name": "serving_llama70B_tp4_sharegpt",
+        "qps_list": [1, 4, 16, "inf"],
+        "server_parameters": {
+            "model": "meta-llama/Meta-Llama-3.1-70B-Instruct",
+            "tensor_parallel_size": 4,
+            "disable_log_stats": "",
+            "load_format": "dummy"
+        },
+        "client_parameters": {
+            "model": "meta-llama/Meta-Llama-3.1-70B-Instruct",
+            "backend": "vllm",
+            "dataset_name": "sharegpt",
+            "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json",
+            "temperature": 0,
+            "num_prompts": 200
+        }
+    },
+    {
+        "test_name": "serving_mixtral8x7B_tp2_sharegpt",
+        "qps_list": [1, 4, 16, "inf"],
+        "server_parameters": {
+            "model": "mistralai/Mixtral-8x7B-Instruct-v0.1",
+            "tensor_parallel_size": 2,
+            "disable_log_stats": "",
+            "load_format": "dummy"
+        },
+        "client_parameters": {
+            "model": "mistralai/Mixtral-8x7B-Instruct-v0.1",
+            "backend": "vllm",
+            "dataset_name": "sharegpt",
+            "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json",
+            "temperature": 0,
+            "num_prompts": 200
+        }
+    },
+    {
+        "test_name": "serving_llama70B_tp4_sharegpt_specdecode",
+        "qps_list": [2],
+        "server_parameters": {
+            "model": "meta-llama/Meta-Llama-3.1-70B-Instruct", 
+            "tensor_parallel_size": 4,
+            "speculative_config": {
+                "model": "turboderp/Qwama-0.5B-Instruct",
+                "num_speculative_tokens": 4,
+                "draft_tensor_parallel_size": 1
+            }
+        },
+        "client_parameters": {
+            "model": "meta-llama/Meta-Llama-3.1-70B-Instruct",
+            "backend": "vllm",
+            "dataset_name": "sharegpt",
+            "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json",
+            "temperature": 0,
+            "num_prompts": 200 
+        }
    }
-  },
-  "tests": [
-    {
-      "test_name": "serving_llama8B_tp1_sharegpt",
-      "server_parameters": {
-        "tensor_parallel_size": 1
-      },
-      "client_parameters": {
-        "dataset_name": "sharegpt",
-        "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": {
-        "tensor_parallel_size": 1
-      },
-      "client_parameters": {
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_llama8B_tp1_random_128_2048",
-      "server_parameters": {
-        "tensor_parallel_size": 1
-      },
-      "client_parameters": {
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 2048
-      }
-    },
-    {
-      "test_name": "serving_llama8B_tp1_random_2048_128",
-      "server_parameters": {
-        "tensor_parallel_size": 1
-      },
-      "client_parameters": {
-        "dataset_name": "random",
-        "random-input-len": 2048,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_llama8B_tp1_random_2048_2048",
-      "server_parameters": {
-        "tensor_parallel_size": 1
-      },
-      "client_parameters": {
-        "dataset_name": "random",
-        "random-input-len": 2048,
-        "random-output-len": 2048
-      }
-    },
-    {
-      "test_name": "serving_llama70B_tp4_random_128_128",
-      "server_parameters": {
-        "model": "meta-llama/Llama-3.3-70B-Instruct",
-        "async_scheduling": "",
-	"no_enable_prefix_caching": "",
-        "max_num_batched_tokens": 8192
-      },
-      "client_parameters": {
-        "model": "meta-llama/Llama-3.3-70B-Instruct",
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_gemma4-e4b_tp1_random_128_128",
-      "server_parameters": {
-        "model": "google/gemma-4-E4B-it",
-	"enable_auto_tool_choice": "",
-	"tool_call_parser": "gemma4",
-	"chat_template": "examples/tool_chat_template_gemma4.jinja",
-	"reasoning_parser": "gemma4"
-      },
-      "client_parameters": {
-        "model": "google/gemma-4-E4B-it",
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 128
-      }
-    }
-  ]
-}
+]
@@ -7,9 +7,6 @@ env:
  CUDA_ARCH_AARCH64: "8.0 8.7 8.9 9.0 10.0 11.0 12.0+PTX"
  CUDA_ARCH_X86_CU129: "7.5 8.0 8.6 8.9 9.0 10.0 12.0"
  CUDA_ARCH_AARCH64_CU129: "8.0 8.7 8.9 9.0 10.0 12.0"
-  MOONCAKE_WHEEL_AARCH64_2_35: "https://vllm-wheels.s3.amazonaws.com/mooncake/mooncake_transfer_engine-0.3.10.post2-0da9dfea3-cp312-cp312-manylinux_2_35_aarch64.whl"
-  MOONCAKE_WHEEL_AARCH64_2_39: "https://vllm-wheels.s3.amazonaws.com/mooncake/mooncake_transfer_engine-0.3.10.post2-0da9dfea3-cp312-cp312-manylinux_2_39_aarch64.whl"
-  MOONCAKE_WHEEL_X86_64: "https://vllm-wheels.s3.amazonaws.com/mooncake/mooncake_transfer_engine-0.3.10.post2-0da9dfea3-cp312-cp312-manylinux_2_35_x86_64.whl"

 steps:
  - input: "Provide Release version here"
@@ -27,11 +24,10 @@ steps:
        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=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64_CU129}\" --build-arg BUILD_OS=manylinux --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda12.9 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
+          - "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=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64_CU129}\" --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"

@@ -45,7 +41,6 @@ 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)"'
        env:
          DOCKER_BUILDKIT: "1"

@@ -59,7 +54,6 @@ 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)"'
        env:
          DOCKER_BUILDKIT: "1"

@@ -69,11 +63,10 @@ steps:
        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=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86_CU129}\" --build-arg BUILD_OS=manylinux --build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda12.9 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
+          - "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=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86_CU129}\" --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"

@@ -87,7 +80,6 @@ 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)"'
        env:
          DOCKER_BUILDKIT: "1"

@@ -101,7 +93,6 @@ 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)"'
        env:
          DOCKER_BUILDKIT: "1"

@@ -139,8 +130,6 @@ steps:
              --build-arg CUDA_VERSION=13.0.2 \
              --build-arg torch_cuda_arch_list="${CUDA_ARCH_X86}" \
              --build-arg INSTALL_KV_CONNECTORS=true \
-              --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_35}" \
-              --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
              --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 \
              --target vllm-openai \
              --progress plain \
@@ -149,7 +138,6 @@ 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)"'

      - label: "Build release image - aarch64 - CUDA 13.0"
        depends_on: ~
@@ -167,14 +155,11 @@ steps:
              --build-arg CUDA_VERSION=13.0.2 \
              --build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64}" \
              --build-arg INSTALL_KV_CONNECTORS=true \
-              --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_35}" \
-              --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
              --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 \
              --target vllm-openai \
              --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)"'

      - label: "Build release image - x86_64 - CUDA 12.9"
        depends_on: ~
@@ -192,8 +177,6 @@ steps:
              --build-arg CUDA_VERSION=12.9.1 \
              --build-arg torch_cuda_arch_list="${CUDA_ARCH_X86_CU129}" \
              --build-arg INSTALL_KV_CONNECTORS=true \
-              --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_35}" \
-              --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
              --target vllm-openai \
              --progress plain \
              -f docker/Dockerfile .
@@ -201,7 +184,6 @@ 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"'

      - label: "Build release image - aarch64 - CUDA 12.9"
        depends_on: ~
@@ -219,13 +201,10 @@ steps:
              --build-arg CUDA_VERSION=12.9.1 \
              --build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64_CU129}" \
              --build-arg INSTALL_KV_CONNECTORS=true \
-              --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_35}" \
-              --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
              --target vllm-openai \
              --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"'

      - label: "Build release image - x86_64 - CUDA 13.0 - Ubuntu 24.04"
        depends_on: ~
@@ -245,8 +224,6 @@ steps:
              --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
              --build-arg torch_cuda_arch_list="${CUDA_ARCH_X86}" \
              --build-arg INSTALL_KV_CONNECTORS=true \
-              --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_39}" \
-              --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
              --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu24.04 \
              --target vllm-openai \
              --progress plain \
@@ -254,7 +231,6 @@ 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"'

      - label: "Build release image - aarch64 - CUDA 13.0 - Ubuntu 24.04"
        depends_on: ~
@@ -274,14 +250,11 @@ steps:
              --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
              --build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64}" \
              --build-arg INSTALL_KV_CONNECTORS=true \
-              --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_39}" \
-              --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
              --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu24.04 \
              --target vllm-openai \
              --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"'

      - label: "Build release image - x86_64 - CUDA 12.9 - Ubuntu 24.04"
        depends_on: ~
@@ -301,15 +274,12 @@ steps:
              --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
              --build-arg torch_cuda_arch_list="${CUDA_ARCH_X86_CU129}" \
              --build-arg INSTALL_KV_CONNECTORS=true \
-              --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_39}" \
-              --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
              --target vllm-openai \
              --progress plain \
              -f docker/Dockerfile .
          - "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"'

      - label: "Build release image - aarch64 - CUDA 12.9 - Ubuntu 24.04"
        depends_on: ~
@@ -329,13 +299,10 @@ steps:
              --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
              --build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64_CU129}" \
              --build-arg INSTALL_KV_CONNECTORS=true \
-              --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_39}" \
-              --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
              --target vllm-openai \
              --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"'

      - block: "Build release image for x86_64 CPU"
        key: block-cpu-release-image-build
@@ -353,7 +320,6 @@ 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)"'
        env:
          DOCKER_BUILDKIT: "1"

@@ -373,7 +339,6 @@ 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)"'
        env:
          DOCKER_BUILDKIT: "1"

@@ -391,7 +356,15 @@ 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"'
+
+      - label: "Annotate release workflow - CUDA 13.0"
+        depends_on:
+          - create-multi-arch-manifest
+        id: annotate-release-workflow
+        agents:
+          queue: small_cpu_queue_release
+        commands:
+          - "bash .buildkite/scripts/annotate-release.sh"

      - label: "Create multi-arch manifest - CUDA 12.9"
        depends_on:
@@ -404,7 +377,6 @@ 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"'

      - label: "Create multi-arch manifest - CUDA 13.0 - Ubuntu 24.04"
        depends_on:
@@ -417,7 +389,6 @@ 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"'

      - label: "Create multi-arch manifest - CUDA 12.9 - Ubuntu 24.04"
        depends_on:
@@ -430,7 +401,6 @@ 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"'

      - label: "Publish nightly multi-arch image to DockerHub"
        depends_on:
@@ -468,6 +438,59 @@ steps:
          DOCKER_BUILDKIT: "1"
          DOCKERHUB_USERNAME: "vllmbot"

+      - block: "Publish release images to DockerHub"
+        key: block-publish-release-images
+        depends_on:
+          - create-multi-arch-manifest
+          - create-multi-arch-manifest-cuda-12-9
+          - create-multi-arch-manifest-ubuntu2404
+          - create-multi-arch-manifest-cuda-12-9-ubuntu2404
+          - build-rocm-release-image
+          - input-release-version
+          # Wait for CPU builds if their block steps were unblocked, so publish
+          # doesn't race the in-progress CPU build. allow_failure lets publish
+          # proceed when the operator legitimately leaves the CPU block steps
+          # unblocked or the CPU build fails.
+          - step: build-cpu-release-image-x86
+            allow_failure: true
+          - step: build-cpu-release-image-arm64
+            allow_failure: true
+        if: build.env("NIGHTLY") != "1"
+
+      - label: "Publish release images to DockerHub"
+        depends_on:
+          - block-publish-release-images
+        key: publish-release-images-dockerhub
+        agents:
+          queue: small_cpu_queue_release
+        commands:
+          - "bash .buildkite/scripts/publish-release-images.sh"
+        plugins:
+          - docker-login#v3.0.0:
+              username: vllmbot
+              password-env: DOCKERHUB_TOKEN
+        env:
+          DOCKER_BUILDKIT: "1"
+          DOCKERHUB_USERNAME: "vllmbot"
+
+  - group: "Publish wheels"
+    key: "publish-wheels"
+    steps:
+      - block: "Confirm update release wheels to PyPI (experimental, use with caution)?"
+        key: block-upload-release-wheels
+        depends_on:
+          - input-release-version
+          - build-wheels
+
+      - label: "Upload release wheels to PyPI"
+        depends_on:
+          - block-upload-release-wheels
+        id: upload-release-wheels
+        agents:
+          queue: small_cpu_queue_release
+        commands:
+          - "bash .buildkite/scripts/upload-release-wheels-pypi.sh"
+
  # =============================================================================
  # ROCm Release Pipeline (x86_64 only)
  # =============================================================================
@@ -581,7 +604,7 @@ steps:
          echo ""
          echo " Build complete - Image and wheels cached"
        fi
-
+        
    artifact_paths:
      - "artifacts/rocm-base-wheels/*.whl"
    env:
@@ -737,7 +760,7 @@ steps:
      - "bash tools/vllm-rocm/generate-rocm-wheels-root-index.sh"
    env:
      S3_BUCKET: "vllm-wheels"
-      VARIANT: "rocm723"
+      VARIANT: "rocm722"

  # ROCm Job 6: Build ROCm Release Docker Image
  - label: ":docker: Build release image - x86_64 - ROCm"
@@ -797,7 +820,7 @@ steps:
        
        # Push to ECR
        docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm
-
+        
        echo ""
        echo " Successfully built and pushed ROCm release image"
        echo "   Image: public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm"
@@ -824,60 +847,3 @@ steps:
    env:
      DOCKER_BUILDKIT: "1"
      DOCKERHUB_USERNAME: "vllmbot"
-
-  # =============================================================================
-  # Publish to DockerHub and PyPI (at the end so all builds complete first)
-  # =============================================================================
-
-  - block: "Publish release images to DockerHub"
-    key: block-publish-release-images
-    depends_on:
-      - create-multi-arch-manifest
-      - create-multi-arch-manifest-cuda-12-9
-      - create-multi-arch-manifest-ubuntu2404
-      - create-multi-arch-manifest-cuda-12-9-ubuntu2404
-      - build-rocm-release-image
-      - input-release-version
-      # Wait for CPU builds if their block steps were unblocked, so publish
-      # doesn't race the in-progress CPU build. allow_failure lets publish
-      # proceed when the operator legitimately leaves the CPU block steps
-      # unblocked or the CPU build fails.
-      - step: build-cpu-release-image-x86
-        allow_failure: true
-      - step: build-cpu-release-image-arm64
-        allow_failure: true
-    if: build.env("NIGHTLY") != "1"
-
-  - label: "Publish release images to DockerHub"
-    depends_on:
-      - block-publish-release-images
-    key: publish-release-images-dockerhub
-    agents:
-      queue: small_cpu_queue_release
-    commands:
-      - "bash .buildkite/scripts/publish-release-images.sh"
-    plugins:
-      - docker-login#v3.0.0:
-          username: vllmbot
-          password-env: DOCKERHUB_TOKEN
-    env:
-      DOCKER_BUILDKIT: "1"
-      DOCKERHUB_USERNAME: "vllmbot"
-
-  - group: "Publish wheels"
-    key: "publish-wheels"
-    steps:
-      - block: "Confirm update release wheels to PyPI (experimental, use with caution)?"
-        key: block-upload-release-wheels
-        depends_on:
-          - input-release-version
-          - build-wheels
-
-      - label: "Upload release wheels to PyPI"
-        depends_on:
-          - block-upload-release-wheels
-        id: upload-release-wheels
-        agents:
-          queue: small_cpu_queue_release
-        commands:
-          - "bash .buildkite/scripts/upload-release-wheels-pypi.sh"
@@ -1,9 +0,0 @@
-#!/bin/bash
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-#
-# Append a build artifact line to the Buildkite annotation.
-# Usage: annotate-build-artifact.sh <label> <value>
-set -e
-echo "- **${1}**: \`${2}\`" | \
-  buildkite-agent annotate --append --style 'info' --context 'release-artifacts'
@@ -0,0 +1,27 @@
+#!/bin/bash
+
+set -ex
+
+# Get release version, default to 1.0.0.dev for nightly/per-commit builds
+RELEASE_VERSION=$(buildkite-agent meta-data get release-version 2>/dev/null | sed 's/^v//')
+if [ -z "${RELEASE_VERSION}" ]; then
+  RELEASE_VERSION="1.0.0.dev"
+fi
+
+buildkite-agent annotate --style 'info' --context 'release-workflow' << EOF
+To download the wheel (by commit):
+\`\`\`
+aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux_2_35_x86_64.whl .
+aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux_2_35_aarch64.whl .
+
+(Optional) For CUDA 12.9:
+aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu129-cp38-abi3-manylinux_2_31_x86_64.whl .
+aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu129-cp38-abi3-manylinux_2_31_aarch64.whl .
+
+(Optional) For CPU:
+aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cpu-cp38-abi3-manylinux_2_35_x86_64.whl .
+aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cpu-cp38-abi3-manylinux_2_35_aarch64.whl .
+\`\`\`
+
+Docker images are published automatically by the "Publish release images to DockerHub" pipeline step.
+EOF
@@ -9,13 +9,6 @@
 # Find <build_number> and <job_uuid> via:
 #   gh pr checks <PR> --repo vllm-project/vllm
 # Each failing row's URL is .../builds/<build_number>#<job_uuid>.
-#
-# Default output path: ci-<build>-<uuid_first_13_chars>.log (e.g.
-# ci-68478-019e6b07-daae.log). Jobs in the same build share the UUID's
-# first 8 chars, so the second segment is needed for uniqueness when
-# fetching multiple jobs in parallel. The script refuses to overwrite an
-# existing output file; pass an explicit path or set CI_FETCH_LOG_FORCE=1
-# to override.

 set -euo pipefail

@@ -33,12 +26,12 @@ if [ $# -lt 1 ]; then usage; fi
 if [[ "$1" == https://* ]]; then
    BUILD=$(echo "$1" | sed -nE 's#.*/builds/([0-9]+).*#\1#p')
    JOB=$(echo "$1" | grep -oE '[0-9a-f]{8}-[0-9a-f-]+' | head -n 1)
-    OUT="${2:-}"
+    OUT="${2:-ci-${BUILD}-${JOB:0:8}.log}"
 else
    if [ $# -lt 2 ]; then usage; fi
    BUILD="$1"
    JOB="$2"
-    OUT="${3:-}"
+    OUT="${3:-ci-${BUILD}-${JOB:0:8}.log}"
 fi

 if [ -z "$BUILD" ] || [ -z "$JOB" ]; then
@@ -46,18 +39,6 @@ if [ -z "$BUILD" ] || [ -z "$JOB" ]; then
    usage
 fi

-# Jobs in the same build share the UUID's first segment, so include the
-# second segment (chars 9-13, e.g. "019e6b07-daae") to keep default filenames
-# unique when fetching multiple jobs from one build in parallel.
-if [ -z "$OUT" ]; then
-    OUT="ci-${BUILD}-${JOB:0:13}.log"
-fi
-
-if [ -e "$OUT" ] && [ -z "${CI_FETCH_LOG_FORCE:-}" ]; then
-    echo "Refusing to overwrite existing $OUT (set CI_FETCH_LOG_FORCE=1 or pass an explicit output path)." >&2
-    exit 1
-fi
-
 COOKIES=$(mktemp)
 trap 'rm -f "$COOKIES"' EXIT

@@ -35,9 +35,25 @@ export PYTHONPATH=".."
 # Helper Functions
 ###############################################################################

-report_docker_usage() {
-  echo "--- Docker usage"
-  docker system df || true
+cleanup_docker() {
+  # Get Docker's root directory
+  docker_root=$(docker info -f '{{.DockerRootDir}}')
+  if [ -z "$docker_root" ]; then
+    echo "Failed to determine Docker root directory."
+    exit 1
+  fi
+  echo "Docker root directory: $docker_root"
+
+  disk_usage=$(df "$docker_root" | tail -1 | awk '{print $5}' | sed 's/%//')
+  threshold=70
+  if [ "$disk_usage" -gt "$threshold" ]; then
+    echo "Disk usage is above $threshold%. Cleaning up Docker images and volumes..."
+    docker image prune -f
+    docker volume prune -f && docker system prune --force --filter "until=72h" --all
+    echo "Docker images and volumes cleanup completed."
+  else
+    echo "Disk usage is below $threshold%. No cleanup needed."
+  fi
 }

 cleanup_network() {
@@ -98,7 +114,8 @@ handle_pytest_exit() {
 # unquoted since they have no spaces and work fine.
 #
 # Already-quoted expressions (containing literal single quotes) are passed
-# through untouched to avoid double-quoting well-formed shell fragments.
+# through untouched to avoid double-quoting values injected by
+# apply_rocm_test_overrides.
 #
 # NOTE: This ONLY fixes -m/-k flags. It cannot recover arbitrary inner
 # double-quotes stripped by the calling shell (see header comment).
@@ -230,6 +247,102 @@ re_quote_pytest_markers() {
  echo "${output% }"
 }

+###############################################################################
+# ROCm-specific pytest command rewrites
+#
+# These apply ignore flags and environment overrides for tests that are not
+# yet supported or behave differently on ROCm hardware. Kept as a single
+# function so new exclusions are easy to add in one place.
+###############################################################################
+
+apply_rocm_test_overrides() {
+  local cmds="$1"
+
+  # --- Model registry filter ---
+  if [[ $cmds == *"pytest -v -s models/test_registry.py"* ]]; then
+    cmds=${cmds//"pytest -v -s models/test_registry.py"/"pytest -v -s models/test_registry.py -k 'not BambaForCausalLM and not GritLM and not Mamba2ForCausalLM and not Zamba2ForCausalLM'"}
+  fi
+
+  # --- LoRA: disable custom paged attention ---
+  if [[ $cmds == *"pytest -v -s lora"* ]]; then
+    cmds=${cmds//"pytest -v -s lora"/"pytest -v -s lora"}
+  fi
+
+  # --- Kernel ignores ---
+  if [[ $cmds == *" kernels/core"* ]]; then
+    cmds="${cmds} \
+    --ignore=kernels/core/test_fused_quant_layernorm.py \
+    --ignore=kernels/core/test_permute_cols.py"
+  fi
+
+  if [[ $cmds == *" kernels/attention"* ]]; then
+    cmds="${cmds} \
+    --ignore=kernels/attention/test_attention_selector.py \
+    --ignore=kernels/attention/test_encoder_decoder_attn.py \
+    --ignore=kernels/attention/test_flash_attn.py \
+    --ignore=kernels/attention/test_flashinfer.py \
+    --ignore=kernels/attention/test_prefix_prefill.py \
+    --ignore=kernels/attention/test_cascade_flash_attn.py \
+    --ignore=kernels/attention/test_mha_attn.py \
+    --ignore=kernels/attention/test_lightning_attn.py \
+    --ignore=kernels/attention/test_attention.py"
+  fi
+
+  if [[ $cmds == *" kernels/quantization"* ]]; then
+    cmds="${cmds} \
+    --ignore=kernels/quantization/test_int8_quant.py \
+    --ignore=kernels/quantization/test_machete_mm.py \
+    --ignore=kernels/quantization/test_block_fp8.py \
+    --ignore=kernels/quantization/test_block_int8.py \
+    --ignore=kernels/quantization/test_marlin_gemm.py \
+    --ignore=kernels/quantization/test_cutlass_scaled_mm.py \
+    --ignore=kernels/quantization/test_int8_kernel.py"
+  fi
+
+  if [[ $cmds == *" kernels/mamba"* ]]; then
+    cmds="${cmds} \
+    --ignore=kernels/mamba/test_mamba_mixer2.py \
+    --ignore=kernels/mamba/test_causal_conv1d.py \
+    --ignore=kernels/mamba/test_mamba_ssm_ssd.py"
+  fi
+
+  if [[ $cmds == *" kernels/moe"* ]]; then
+    cmds="${cmds} \
+    --ignore=kernels/moe/test_moe.py \
+    --ignore=kernels/moe/test_cutlass_moe.py"
+  fi
+
+  # --- Entrypoint ignores ---
+  if [[ $cmds == *" entrypoints/openai "* ]]; then
+    cmds=${cmds//" entrypoints/openai "/" entrypoints/openai \
+    --ignore=entrypoints/openai/chat_completion/test_audio.py \
+    --ignore=entrypoints/openai/completion/test_shutdown.py \
+    --ignore=entrypoints/openai/test_completion.py \
+    --ignore=entrypoints/openai/models/test_models.py \
+    --ignore=entrypoints/openai/test_return_tokens_as_ids.py \
+    --ignore=entrypoints/openai/chat_completion/test_root_path.py \
+    --ignore=entrypoints/openai/completion/test_prompt_validation.py "}
+  fi
+
+  if [[ $cmds == *" entrypoints/serve"* ]]; then
+    cmds="${cmds} \
+    --ignore=entrypoints/serve/lora/test_lora_adapters.py"
+  fi
+
+  if [[ $cmds == *" entrypoints/llm "* ]]; then
+    cmds=${cmds//" entrypoints/llm "/" entrypoints/llm \
+    --ignore=entrypoints/llm/test_chat.py \
+    --ignore=entrypoints/llm/test_accuracy.py \
+    --ignore=entrypoints/llm/test_init.py \
+    --ignore=entrypoints/llm/test_prompt_validation.py "}
+  fi
+
+  # Clean up escaped newlines from --ignore appends
+  cmds=$(echo "$cmds" | sed 's/ \\ / /g')
+
+  echo "$cmds"
+}
+
 ###############################################################################
 # Main
 ###############################################################################
@@ -238,8 +351,8 @@ re_quote_pytest_markers() {
 echo "--- ROCm info"
 rocminfo

-# --- Docker status ---
-report_docker_usage
+# --- Docker housekeeping ---
+cleanup_docker

 # --- Pull test image ---
 echo "--- Pulling container"
@@ -248,17 +361,9 @@ container_name="rocm_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | hea
 docker pull "${image_name}"

 remove_docker_container() {
-  # docker run uses --rm, so the container is normally already gone when the
-  # EXIT trap runs. Cleanup is best-effort and must not affect the test result.
-  docker rm -f "${container_name}" >/dev/null 2>&1 || true
+  docker rm -f "${container_name}" || docker image rm -f "${image_name}" || true
 }
-
-on_exit() {
-  local exit_code=$?
-  remove_docker_container
-  exit "$exit_code"
-}
-trap on_exit EXIT
+trap remove_docker_container EXIT

 # --- Prepare commands ---
 echo "--- Running container"
@@ -273,11 +378,9 @@ HF_MOUNT="/root/.cache/huggingface"
 # double-quotes will have been stripped by the calling shell.
 if [[ -n "${VLLM_TEST_COMMANDS:-}" ]]; then
  commands="${VLLM_TEST_COMMANDS}"
-  commands_source="env"
  echo "Commands sourced from VLLM_TEST_COMMANDS (quoting preserved)"
 else
  commands="$*"
-  commands_source="argv"
  if [[ -z "$commands" ]]; then
    echo "Error: No test commands provided." >&2
    echo "Usage:" >&2
@@ -294,16 +397,11 @@ fi

 echo "Raw commands: $commands"

-# Only try to repair stripped pytest -m/-k quoting in legacy argv mode.
-# VLLM_TEST_COMMANDS preserves inner quoting already, and re-quoting that path
-# can corrupt embedded echo strings or otherwise well-formed shell fragments.
-if [[ "$commands_source" == "argv" ]]; then
-  commands=$(re_quote_pytest_markers "$commands")
-  echo "After re-quoting: $commands"
-else
-  echo "Skipping re-quoting for VLLM_TEST_COMMANDS input"
-fi
+# Fix quoting before ROCm overrides (so overrides see correct structure)
+commands=$(re_quote_pytest_markers "$commands")
+echo "After re-quoting: $commands"

+commands=$(apply_rocm_test_overrides "$commands")
 echo "Final commands: $commands"

 MYPYTHONPATH=".."
@@ -31,22 +31,6 @@ function cpu_tests() {
    set -e
    pip list"

-  # Run kernel tests
-  docker exec cpu-test bash -c "
-    set -e
-    pytest -x -v -s tests/kernels/test_onednn.py
-    pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
-    pytest -x -v -s tests/kernels/core/test_cpu_activation.py
-    pytest -x -v -s tests/kernels/moe/test_moe.py -k test_cpu_fused_moe_basic
-    pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py"
-
-  # skip tests requiring model downloads if HF_TOKEN is not set
-  # due to rate-limits
-  if [ -z "$HF_TOKEN" ]; then
-    echo "Warning: HF_TOKEN is not set. Skipping tests that require model downloads."
-    return
-  fi
-
  # offline inference
  docker exec cpu-test bash -c "
    set -e
@@ -62,6 +46,13 @@ function cpu_tests() {
    set -e
    pytest -x -v -s tests/quantization/test_compressed_tensors.py::test_compressed_tensors_w8a8_logprobs"

+  # Run kernel tests
+  docker exec cpu-test bash -c "
+    set -e
+    pytest -x -v -s tests/kernels/test_onednn.py
+    pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
+    pytest -x -v -s tests/kernels/core/test_cpu_activation.py
+    pytest -x -v -s tests/kernels/moe/test_moe.py -k test_cpu_fused_moe_basic"

  # basic online serving
  docker exec cpu-test bash -c '
@@ -76,21 +67,6 @@ function cpu_tests() {
      --num-prompts 20 \
      --endpoint /v1/completions
    kill -s SIGTERM $server_pid &'
-
-  # smoke test for Gated DeltaNet
-  docker exec cpu-test bash -c '
-    set -e
-    VLLM_CPU_OMP_THREADS_BIND=$E2E_OMP_THREADS vllm serve Qwen/Qwen3.5-0.8B --max-model-len 2048 &
-    server_pid=$!
-    timeout 600 bash -c "until curl localhost:8000/v1/models; do sleep 1; done" || exit 1
-    vllm bench serve \
-      --backend vllm \
-      --dataset-name random \
-      --model Qwen/Qwen3.5-0.8B \
-      --num-prompts 20 \
-      --endpoint /v1/completions
-    kill -s SIGTERM $server_pid &'
-
 }

 # All of CPU tests are expected to be finished less than 40 mins.
@@ -352,31 +352,17 @@ if [[ -z "${ZE_AFFINITY_MASK:-}" ]]; then
  echo "Warning: ZE_AFFINITY_MASK is not set. Proceeding without device affinity." >&2
 fi

-export CMDS="${commands}"
-export HF_TOKEN ZE_AFFINITY_MASK
-
-{
-  flock 9
-  if ! docker image inspect "${IMAGE}" >/dev/null 2>&1; then
-    echo 'Image missing before container creation, pulling again...'
-    timeout 900 docker pull "${IMAGE}"
-  fi
-
-  docker create \
+docker run \
    --device /dev/dri:/dev/dri \
    --net=host \
    --ipc=host \
    --privileged \
    -v /dev/dri/by-path:/dev/dri/by-path \
-    -v "${HOME}/.cache/huggingface:/root/.cache/huggingface" \
-    --entrypoint='' \
-    -e HF_TOKEN \
-    -e ZE_AFFINITY_MASK \
-    -e CMDS \
+    -v ${HOME}/.cache/huggingface:/root/.cache/huggingface \
+    --entrypoint="" \
+    -e "HF_TOKEN=${HF_TOKEN:-}" \
+    -e "ZE_AFFINITY_MASK=${ZE_AFFINITY_MASK:-}" \
+    -e "CMDS=${commands}" \
    --name "${container_name}" \
-    "${IMAGE}" \
-    bash -c 'set -e; echo "ZE_AFFINITY_MASK is ${ZE_AFFINITY_MASK:-}"; eval "$CMDS"' \
-    >/dev/null
-} 9>/tmp/docker-pull.lock
-
-docker start -a "${container_name}"
+    "${image_name}" \
+    bash -c 'set -e; echo "ZE_AFFINITY_MASK is ${ZE_AFFINITY_MASK:-}"; eval "$CMDS"'
@@ -61,7 +61,7 @@ echo "Results will be stored in: $RESULTS_DIR"
 echo "--- Installing Python dependencies ---"
 python3 -m pip install --progress-bar off git+https://github.com/thuml/depyf.git \
    && python3 -m pip install --progress-bar off pytest pytest-asyncio tpu-info \
-    && python3 -m pip install --progress-bar off "lm-eval[api]>=0.4.12" \
+    && python3 -m pip install --progress-bar off "lm-eval[api]>=0.4.11" \
    && python3 -m pip install --progress-bar off hf-transfer tblib==3.1.0
 echo "--- Python dependencies installed ---"

@@ -61,7 +61,7 @@ echo "Results will be stored in: $RESULTS_DIR"
 echo "--- Installing Python dependencies ---"
 python3 -m pip install --progress-bar off git+https://github.com/thuml/depyf.git \
    && python3 -m pip install --progress-bar off pytest pytest-asyncio tpu-info \
-    && python3 -m pip install --progress-bar off "lm-eval[api]>=0.4.12" \
+    && python3 -m pip install --progress-bar off "lm-eval[api]>=0.4.11" \
    && python3 -m pip install --progress-bar off hf-transfer tblib==3.1.0
 echo "--- Python dependencies installed ---"

@@ -136,6 +136,8 @@ run_and_track_test 3 "test_accuracy.py::test_lm_eval_accuracy_v1_engine" \
    "python3 -m pytest -s -v /workspace/vllm/tests/entrypoints/llm/test_accuracy.py::test_lm_eval_accuracy_v1_engine"
 run_and_track_test 4 "test_quantization_accuracy.py" \
    "python3 -m pytest -s -v /workspace/vllm/tests/tpu/test_quantization_accuracy.py"
+run_and_track_test 5 "examples/offline_inference/tpu.py" \
+    "python3 /workspace/vllm/examples/offline_inference/tpu.py"
 run_and_track_test 6 "test_tpu_model_runner.py" \
    "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/worker/test_tpu_model_runner.py"
 run_and_track_test 7 "test_sampler.py" \
@@ -0,0 +1,53 @@
+#!/bin/bash
+
+# This script build the CPU docker image and run the offline inference inside the container.
+# It serves a sanity check for compilation and basic model usage.
+set -ex
+
+image_name="xpu/vllm-ci:${BUILDKITE_COMMIT}"
+container_name="xpu_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | head -c 10; echo)"
+
+# Try building the docker image
+docker build -t "${image_name}" -f docker/Dockerfile.xpu .
+
+# Setup cleanup
+remove_docker_container() {
+  docker rm -f "${container_name}" || true
+}
+trap remove_docker_container EXIT
+
+# Run the image and test offline inference/tensor parallel
+docker run \
+    --device /dev/dri:/dev/dri \
+    --net=host \
+    --ipc=host \
+    --privileged \
+    -v /dev/dri/by-path:/dev/dri/by-path \
+    --entrypoint="" \
+    -e "HF_TOKEN=${HF_TOKEN}" \
+    -e "ZE_AFFINITY_MASK=${ZE_AFFINITY_MASK}" \
+    --name "${container_name}" \
+    "${image_name}" \
+    bash -c '
+    set -e
+    echo $ZE_AFFINITY_MASK
+    python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager
+    python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 -O3 -cc.cudagraph_mode=NONE
+    python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend ray
+    python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend mp
+    python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --attention-backend=TRITON_ATTN
+    python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --quantization fp8
+    python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4  --block-size 64 --enforce-eager --max-model-len 8192
+    python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b  --block-size 64 --enforce-eager -tp 2
+    python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b  --block-size 64 --enforce-eager -tp 2 --enable-expert-parallel
+    python3 examples/basic/offline_inference/generate.py --model OPEA/Qwen2.5-0.5B-Instruct-int4-sym-inc --block-size 64 --enforce-eager --max-model-len 8192
+    cd tests
+    pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py --ignore=v1/core/test_scheduler_e2e.py
+    pytest -v -s v1/engine
+    pytest -v -s v1/sample --ignore=v1/sample/test_logprobs.py --ignore=v1/sample/test_logprobs_e2e.py
+    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/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_tree_attention.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.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
+    pytest -v -s v1/test_serial_utils.py
+'
@@ -1,39 +0,0 @@
-#!/bin/bash
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-set -euo pipefail
-
-REQUIREMENTS_FILE="${KV_CONNECTORS_REQUIREMENTS:-/vllm-workspace/requirements/kv_connectors.txt}"
-
-uv pip install --system -r "${REQUIREMENTS_FILE}"
-
-NIXL_METADATA=$(python3 - <<'PY'
-import importlib.metadata as metadata
-
-import torch
-
-cuda_version = torch.version.cuda
-if cuda_version is None:
-    raise SystemExit("torch.version.cuda is not set")
-
-print(cuda_version.split(".", 1)[0], metadata.version("nixl"))
-PY
-)
-read -r CUDA_MAJOR NIXL_VERSION <<<"${NIXL_METADATA}"
-
-# nixl>=1.1.0 can install multiple CUDA wheel variants. Keep only the variant
-# matching this CI image so nixl_ep_cpp links against the available libcudart.
-uv pip uninstall --system nixl-cu12 nixl-cu13 2>/dev/null || true
-uv pip install --system --no-deps "nixl-cu${CUDA_MAJOR}==${NIXL_VERSION}"
-
-python3 - <<'PY'
-import importlib.metadata as metadata
-
-for package_name in ("nixl", "nixl-cu12", "nixl-cu13"):
-    try:
-        version = metadata.version(package_name)
-    except metadata.PackageNotFoundError:
-        version = "not installed"
-    print(f"{package_name}: {version}")
-PY
@@ -1,156 +0,0 @@
-#!/usr/bin/env bash
-set -euo pipefail
-
-MODE="${1:-}"
-
-if [[ "$MODE" != "style-clippy" && "$MODE" != "test" ]]; then
-  echo "Usage: $0 {style-clippy|test}" >&2
-  exit 2
-fi
-
-ROOT_DIR="$(git rev-parse --show-toplevel)"
-cd "$ROOT_DIR"
-
-export CARGO_TERM_COLOR="${CARGO_TERM_COLOR:-always}"
-export CARGO_HOME="${CARGO_HOME:-$HOME/.cargo}"
-export RUSTUP_HOME="${RUSTUP_HOME:-$HOME/.rustup}"
-export PATH="$CARGO_HOME/bin:$PATH"
-
-log_section() {
-  echo "--- $*"
-}
-
-install_protoc() {
-  if command -v protoc >/dev/null 2>&1; then
-    return
-  fi
-
-  local version="${PROTOC_VERSION:-31.1}"
-  local arch
-  case "$(uname -m)" in
-    x86_64)
-      arch="x86_64"
-      ;;
-    aarch64|arm64)
-      arch="aarch_64"
-      ;;
-    *)
-      echo "Unsupported protoc architecture: $(uname -m)" >&2
-      return 1
-      ;;
-  esac
-
-  local url="https://github.com/protocolbuffers/protobuf/releases/download/v${version}/protoc-${version}-linux-${arch}.zip"
-  local tmp_dir
-  tmp_dir="$(mktemp -d)"
-
-  log_section "Installing protoc ${version}"
-  curl -L --proto '=https' --tlsv1.2 -sSf "$url" -o "$tmp_dir/protoc.zip"
-  mkdir -p "$CARGO_HOME/bin"
-  unzip -q "$tmp_dir/protoc.zip" bin/protoc 'include/*' -d "$CARGO_HOME"
-  chmod +x "$CARGO_HOME/bin/protoc"
-  rm -rf "$tmp_dir"
-}
-
-rust_toolchain() {
-  awk -F '"' '/channel[[:space:]]*=/ { print $2; exit }' rust-toolchain.toml
-}
-
-install_rust_toolchain() {
-  log_section "Installing Rust toolchain"
-  if ! command -v rustup >/dev/null 2>&1; then
-    curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs \
-      | sh -s -- -y --profile minimal --default-toolchain none
-  fi
-
-  local toolchain
-  toolchain="$(rust_toolchain)"
-  rustup toolchain install "$toolchain" --profile minimal --component rustfmt,clippy
-  rustup component add --toolchain "$toolchain" rustfmt clippy
-}
-
-install_cargo_binstall() {
-  if command -v cargo-binstall >/dev/null 2>&1; then
-    return
-  fi
-
-  log_section "Installing cargo-binstall"
-  curl -L --proto '=https' --tlsv1.2 -sSf \
-    https://raw.githubusercontent.com/cargo-bins/cargo-binstall/main/install-from-binstall-release.sh \
-    | bash
-}
-
-install_cargo_sort() {
-  if command -v cargo-sort >/dev/null 2>&1; then
-    return
-  fi
-
-  log_section "Installing cargo-sort"
-  install_cargo_binstall
-  cargo binstall --no-confirm cargo-sort
-}
-
-install_cargo_nextest() {
-  if command -v cargo-nextest >/dev/null 2>&1; then
-    return
-  fi
-
-  log_section "Installing cargo-nextest"
-  install_cargo_binstall
-  cargo binstall --no-confirm --secure cargo-nextest
-}
-
-install_uv() {
-  if command -v uv >/dev/null 2>&1; then
-    return
-  fi
-
-  log_section "Installing uv"
-  curl -LsSf --proto '=https' --tlsv1.2 https://astral.sh/uv/install.sh \
-    | env UV_INSTALL_DIR="$CARGO_HOME/bin" sh
-}
-
-run_style_clippy() {
-  install_cargo_sort
-
-  log_section "Checking Rust formatting"
-  cargo fmt --manifest-path rust/Cargo.toml --all -- --check
-
-  log_section "Checking Cargo.toml ordering"
-  cargo sort --workspace --check rust
-
-  log_section "Running clippy"
-  cargo clippy \
-    --manifest-path rust/Cargo.toml \
-    --workspace \
-    --all-targets \
-    --all-features \
-    --locked \
-    -- \
-    -D warnings
-}
-
-run_tests() {
-  install_uv
-  install_cargo_nextest
-
-  log_section "Running cargo nextest"
-  cargo nextest run \
-    --manifest-path rust/Cargo.toml \
-    --workspace \
-    --all-features \
-    --locked \
-    --no-fail-fast
-}
-
-install_protoc
-install_rust_toolchain
-
-case "$MODE" in
-  style-clippy)
-    run_style_clippy
-    ;;
-  test)
-    run_tests
-    ;;
-esac
@@ -49,7 +49,6 @@ for BACK in "${BACKENDS[@]}"; do
    --data-parallel-size 2 \
    --enable-expert-parallel \
    --enable-eplb \
-    --eplb-config '{"use_async": false}' \
    --trust-remote-code \
    --max-model-len 2048 \
    --all2all-backend "$BACK" \
@@ -48,7 +48,7 @@ for BACK in "${BACKENDS[@]}"; do
    --enforce-eager \
    --enable-eplb \
    --all2all-backend "$BACK" \
-    --eplb-config '{"window_size":10, "step_interval":100, "num_redundant_experts":0, "log_balancedness":true, "use_async":false}' \
+    --eplb-config '{"window_size":10, "step_interval":100, "num_redundant_experts":0, "log_balancedness":true}' \
    --tensor-parallel-size "${TENSOR_PARALLEL_SIZE}" \
    --data-parallel-size "${DATA_PARALLEL_SIZE}" \
    --enable-expert-parallel \
@@ -70,7 +70,7 @@ echo "============================================"
 # ---- Install bfcl-eval if missing ----
 if ! python3 -c "import bfcl_eval" 2>/dev/null; then
    echo "Installing bfcl-eval..."
-    uv pip install "bfcl-eval>=2025.10.20.1,<2026"
+    pip install "bfcl-eval>=2025.10.20.1,<2026"
 fi

 # ---- Cleanup handler ----
@@ -100,7 +100,7 @@ SERVE_ARGS=(
    --tensor-parallel-size "$TP_SIZE"
    --max-model-len "$MAX_MODEL_LEN"
    --enforce-eager
-    --enable-prefix-caching
+    --no-enable-prefix-caching
 )

 # Append reasoning parser if specified
@@ -39,17 +39,11 @@ fi

 set -x # avoid printing secrets above

-# install uv if not already available
-if ! command -v uv &> /dev/null; then
-  curl -LsSf https://astral.sh/uv/install.sh | UV_VERSION=0.11.14 sh
-  export PATH="$HOME/.local/bin:$PATH"
-fi
-
-# install twine and sdist build prerequisites using uv with Python 3.12
-uv venv --python 3.12 /tmp/vllm-release-env
+# install twine and sdist build prerequisites from pypi
+python3 -m venv /tmp/vllm-release-env
 source /tmp/vllm-release-env/bin/activate
-uv pip install twine
-uv pip install -r requirements/build/cuda.txt
+pip install twine
+pip install -r requirements/build/cuda.txt
 python3 -m twine --version

 # copy release wheels to local directory
@@ -139,6 +139,19 @@ steps:
  - pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py
  - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest models/multimodal/generation/test_whisper.py -v -s -m 'distributed(num_gpus=2)'

+#--------------------------------------------------------  mi250 · benchmarks  ---------------------------------------------------------#
+
+- label: Benchmarks # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
+  agent_pool: mi250_1
+  working_dir: "/vllm-workspace/.buildkite"
+  source_file_dependencies:
+  - benchmarks/
+  - vllm/platforms/rocm.py
+  commands:
+  - bash scripts/run-benchmarks.sh
+
 #----------------------------------------------------------  mi250 · compile  ----------------------------------------------------------#

 - label: PyTorch Compilation Unit Tests # TBD
@@ -217,6 +230,7 @@ steps:
  - tests/entrypoints/llm/test_collective_rpc.py
  - vllm/platforms/rocm.py
  commands:
+  - export TORCH_NCCL_BLOCKING_WAIT=1
  - pytest -v -s entrypoints/llm/test_collective_rpc.py
  - pytest -v -s ./compile/fullgraph/test_basic_correctness.py
  - pytest -v -s ./compile/test_wrapper.py
@@ -258,6 +272,7 @@ steps:
  - tests/v1/worker/test_worker_memory_snapshot.py
  - vllm/platforms/rocm.py
  commands:
+  - export TORCH_NCCL_BLOCKING_WAIT=1
  - VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
  - VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
  - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
@@ -381,8 +396,8 @@ steps:
    - python3 pooling/embed/vision_embedding_offline.py --seed 0
    # Features demo
    - python3 features/automatic_prefix_caching/prefix_caching_offline.py
-    - python3 deployment/llm_engine_example.py
-    - python3 features/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 features/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
+    - python3 offline_inference/llm_engine_example.py
+    - python3 others/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 others/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
    - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle --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 2048
    - 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

@@ -447,7 +462,7 @@ steps:
  - tests/lora
  - vllm/platforms/rocm.py
  commands:
-  - pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_qwen3_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py --ignore=lora/test_qwen35_densemodel_lora.py
+  - pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_llm_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py --ignore=lora/test_qwen35_densemodel_lora.py

 #------------------------------------------------------  mi250 · model_executor  -------------------------------------------------------#

@@ -472,7 +487,7 @@ steps:
  - pytest -v -s model_executor -m '(not slow_test)'
  - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py

-#------------------------------------------------------  mi250 · models / basic  -------------------------------------------------------#
+#----------------------------------------------------------  mi250 · models  -----------------------------------------------------------#

 - label: Basic Models Test (Other CPU) # TBD
  timeout_in_minutes: 180
@@ -533,8 +548,6 @@ steps:
  commands:
  - pytest -v -s models/test_terratorch.py models/test_transformers.py models/test_registry.py

-#-----------------------------------------------------  mi250 · models / language  -----------------------------------------------------#
-
 - label: Language Models Test (MTEB) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
@@ -577,10 +590,9 @@ steps:
  - vllm/platforms/rocm.py
  commands:
  - pip freeze | grep -E 'torch'
+  - export TORCH_NCCL_BLOCKING_WAIT=1
  - pytest -v -s models/language -m 'core_model and slow_test' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB

-#----------------------------------------------------  mi250 · models / multimodal  ----------------------------------------------------#
-
 - label: Multi-Modal Models (Extended Generation 2) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
@@ -609,7 +621,6 @@ steps:
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
  agent_pool: mi250_1
  torch_nightly: true
-  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -789,6 +800,32 @@ steps:
  commands:
  - pytest -v -s v1/attention

+- label: V1 Core + KV + Metrics # TBD
+  timeout_in_minutes: 60
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
+  agent_pool: mi250_1
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/
+  - tests/v1/core
+  - tests/v1/executor
+  - tests/v1/kv_offload
+  - tests/v1/worker
+  - tests/v1/kv_connector/unit
+  - tests/v1/metrics
+  - tests/entrypoints/openai/correctness/test_lmeval.py
+  commands:
+  - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
+  - pytest -v -s -m 'not cpu_test' v1/core
+  - pytest -v -s v1/executor
+  - pytest -v -s v1/kv_offload
+  - pytest -v -s v1/worker
+  - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
+  - pytest -v -s -m 'not cpu_test' v1/metrics
+  - pip install -U git+https://github.com/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api
+  - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
+
 - label: V1 Sample + Logits # TBD
  timeout_in_minutes: 60
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
@@ -827,6 +864,7 @@ steps:
  - tests/entrypoints/openai/test_multi_api_servers.py
  - vllm/platforms/rocm.py
  commands:
+  - export TORCH_NCCL_BLOCKING_WAIT=1
  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py
  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
@@ -845,7 +883,7 @@ steps:
  - vllm/platforms/rocm.py
  commands:
  - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
-  - ATTENTION_BACKEND=ROCM_ATTN bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
+  - ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh

 - label: V1 e2e (2 GPUs) # TBD
  timeout_in_minutes: 180
@@ -892,9 +930,7 @@ steps:
  - tests/renderers
  - tests/standalone_tests/lazy_imports.py
  - tests/tokenizers_
-  - tests/reasoning
  - tests/tool_parsers
-  - tests/parser
  - tests/transformers_utils
  - tests/config
  commands:
@@ -906,9 +942,8 @@ steps:
  - pytest -v -s -m 'cpu_test' multimodal
  - pytest -v -s renderers
  - pytest -v -s tokenizers_
-  - pytest -v -s reasoning --ignore=reasoning/test_seedoss_reasoning_parser.py --ignore=reasoning/test_glm4_moe_reasoning_parser.py
+  - pytest -v -s reasoning --ignore=reasoning/test_seedoss_reasoning_parser.py --ignore=reasoning/test_glm4_moe_reasoning_parser.py --ignore=reasoning/test_gemma4_reasoning_parser.py
  - pytest -v -s tool_parsers
-  - pytest -v -s parser
  - pytest -v -s transformers_utils
  - pytest -v -s config

@@ -967,6 +1002,18 @@ steps:

 #--------------------------------------------------------  mi300 · benchmarks  ---------------------------------------------------------#

+- label: Benchmarks # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  agent_pool: mi300_1
+  optional: true
+  working_dir: "/vllm-workspace/.buildkite"
+  source_file_dependencies:
+  - benchmarks/
+  - vllm/platforms/rocm.py
+  commands:
+  - bash scripts/run-benchmarks.sh
+
 - label: Benchmarks CLI Test # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -1055,13 +1102,13 @@ steps:
  - vllm/compilation/
  - vllm/model_executor/layers
  - tests/compile/passes/distributed/
-  - tests/compile/fusions_e2e/
  - vllm/_aiter_ops.py
  - vllm/platforms/rocm.py
  commands:
  - export VLLM_TEST_CLEAN_GPU_MEMORY=1
  - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/passes/distributed/test_async_tp.py
-  - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py::test_tp2_ar_rms_fusions
+  - pytest -v -s tests/compile/passes/distributed/test_sequence_parallelism.py
+  - pytest -v -s tests/compile/passes/distributed/test_tp2_ar_rms.py::test_tp2_ar_rms_fusions

 #-----------------------------------------------------------  mi300 · cuda  ------------------------------------------------------------#

@@ -1126,6 +1173,7 @@ steps:
  - vllm/_aiter_ops.py
  - vllm/platforms/rocm.py
  commands:
+  - export TORCH_NCCL_BLOCKING_WAIT=1
  - pytest -v -s tests/distributed/test_context_parallel.py
  - VLLM_LOGGING_LEVEL=DEBUG python3 examples/features/data_parallel/data_parallel_offline.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=allgather_reducescatter --disable-nccl-for-dp-synchronization

@@ -1139,6 +1187,7 @@ steps:
  source_file_dependencies:
  - vllm/
  commands:
+  - export TORCH_NCCL_BLOCKING_WAIT=1
  - pytest -v -s distributed/test_custom_all_reduce.py
  - torchrun --nproc_per_node=2 distributed/test_ca_buffer_sharing.py
  - TARGET_TEST_SUITE=A100 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
@@ -1158,6 +1207,7 @@ steps:
  - tests/examples/features/data_parallel/data_parallel_offline.py
  - vllm/platforms/rocm.py
  commands:
+  - export TORCH_NCCL_BLOCKING_WAIT=1
  - torchrun --nproc-per-node=4 distributed/test_torchrun_example.py
  - PP_SIZE=2 torchrun --nproc-per-node=4 distributed/test_torchrun_example.py
  - TP_SIZE=4 torchrun --nproc-per-node=4 distributed/test_torchrun_example_moe.py
@@ -1203,6 +1253,7 @@ steps:
  - vllm/platforms/rocm.py
  commands:
  - export VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1
+  - export TORCH_NCCL_BLOCKING_WAIT=1
  - pytest -v -s distributed/test_ray_v2_executor.py
  - pytest -v -s distributed/test_ray_v2_executor_e2e.py
  - pytest -v -s distributed/test_pipeline_parallel.py -k "ray"
@@ -1224,6 +1275,7 @@ steps:
  - vllm/v1/worker/gpu_worker.py
  - vllm/platforms/rocm.py
  commands:
+  - export TORCH_NCCL_BLOCKING_WAIT=1
  - torchrun --nproc-per-node=8 ../examples/features/torchrun/torchrun_dp_example_offline.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep

 #--------------------------------------------------------  mi300 · entrypoints  --------------------------------------------------------#
@@ -1232,17 +1284,19 @@ steps:
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
-  optional: true
  fast_check: true
  torch_nightly: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
-  - tests/entrypoints/serve
+  - tests/entrypoints/rpc
+  - tests/entrypoints/serve/instrumentator
+  - tests/tool_use
  commands:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc
-  - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
+  - pytest -v -s entrypoints/serve/instrumentator
+  - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc
+  - pytest -v -s tool_use

 - label: Entrypoints Integration (API Server openai - Part 1) # TBD
  timeout_in_minutes: 180
@@ -1250,7 +1304,6 @@ steps:
  agent_pool: mi300_1
  fast_check: true
  torch_nightly: true
-  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -1258,13 +1311,12 @@ steps:
  - tests/entrypoints/test_chat_utils
  commands:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
+  - pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py

 - label: Entrypoints Integration (API Server openai - Part 2) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
-  optional: true
  fast_check: true
  torch_nightly: true
  working_dir: "/vllm-workspace/tests"
@@ -1272,20 +1324,16 @@ steps:
  - vllm/
  - tests/entrypoints/openai
  - tests/entrypoints/test_chat_utils
-  - tests/entrypoints/generate
-  - tests/tool_use
  commands:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
  - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
+  - pytest -v -s entrypoints/openai/speech_to_text/
  - pytest -v -s entrypoints/test_chat_utils.py
-  - pytest -v -s entrypoints/generate
-  - pytest -v -s tool_use

 - label: Entrypoints Integration (API Server openai - Part 3) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
-  optional: true
  fast_check: true
  torch_nightly: true
  working_dir: "/vllm-workspace/tests"
@@ -1295,21 +1343,7 @@ steps:
  - tests/entrypoints/test_chat_utils
  commands:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py
-
- label: Entrypoints Integration (Speech to Text) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  agent_pool: mi300_1
-  fast_check: true
-  torch_nightly: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/
-  - tests/entrypoints/speech_to_text
-  commands:
-  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/speech_to_text
+  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/speech_to_text/ --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py

 - label: Entrypoints Integration (LLM) # TBD
  timeout_in_minutes: 180
@@ -1369,7 +1403,7 @@ steps:
  - vllm/platforms/rocm.py
  commands:
  - pytest -v -s entrypoints/openai/tool_parsers
-  - pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/offline_mode --ignore=entrypoints/openai --ignore=entrypoints/serve --ignore=entrypoints/test_chat_utils.py  --ignore=entrypoints/pooling --ignore=entrypoints/speech_to_text --ignore=tests/entrypoints/generate
+  - pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/rpc --ignore=entrypoints/sleep --ignore=entrypoints/serve/instrumentator --ignore=entrypoints/openai --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py  --ignore=entrypoints/pooling

 - label: OpenAI API correctness # TBD
  timeout_in_minutes: 180
@@ -1485,7 +1519,7 @@ steps:
  commands:
  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx-fp8-and-mixed.txt

- label: DeepSeek V2-Lite Sync EPLB Accuracy (4xH100-4xMI300) # TBD
+- label: DeepSeek V2-Lite Accuracy (4xH100-4xMI300) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_4
@@ -1527,7 +1561,7 @@ steps:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
  - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large.txt --tp-size=4

- label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy (4xH100-4xMI300) # TBD
+- label: Qwen3-30B-A3B-FP8-block Accuracy (4xH100-4xMI300) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_4
@@ -1622,8 +1656,8 @@ steps:
    - python3 pooling/embed/vision_embedding_offline.py --seed 0
    # Features demo
    - python3 features/automatic_prefix_caching/prefix_caching_offline.py
-    - python3 deployment/llm_engine_example.py
-    - python3 features/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 features/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
+    - python3 offline_inference/llm_engine_example.py
+    - python3 others/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 others/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
    - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle --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 2048
    - 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

@@ -1686,7 +1720,6 @@ steps:
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
-  optional: true
  parallelism: 2
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -1734,18 +1767,17 @@ steps:
  - export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
  - pytest -v -s -x lora/test_chatglm3_tp.py
  - pytest -v -s -x lora/test_llama_tp.py
-  - pytest -v -s -x lora/test_qwen3_with_multi_loras.py
+  - pytest -v -s -x lora/test_llm_with_multi_loras.py
  - pytest -v -s -x lora/test_olmoe_tp.py
  - pytest -v -s -x lora/test_gptoss_tp.py
  - pytest -v -s -x lora/test_qwen35_densemodel_lora.py

-#-----------------------------------------------------  mi300 · models / language  -----------------------------------------------------#
+#----------------------------------------------------------  mi300 · models  -----------------------------------------------------------#

 - label: Language Models Test (Extended Pooling)  # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
-  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -1767,23 +1799,19 @@ steps:
  - pip freeze | grep -E 'torch'
  - pytest -v -s models/language -m 'core_model and (not slow_test)'

-#----------------------------------------------------  mi300 · models / multimodal  ----------------------------------------------------#
-
 - label: Multi-Modal Models (Extended Generation 1) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
-  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
  - tests/models/multimodal/generation
  - tests/models/multimodal/test_mapping.py
  commands:
-  - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@rocm-7.0-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 --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
-
+  - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
+  - pytest -v -s models/multimodal/generation -m 'not core_model' --ignore models/multimodal/generation/test_common.py
+  - pytest -v -s models/multimodal/test_mapping.py

 - label: Multi-Modal Models (Extended Generation 2) # TBD
  timeout_in_minutes: 180
@@ -1795,10 +1823,8 @@ steps:
  - vllm/
  - tests/models/multimodal/generation
  commands:
-  - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@rocm-7.0-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)'
-
+  - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
+  - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=0) and not core_model'

 - label: Multi-Modal Models (Extended Generation 3) # TBD
  timeout_in_minutes: 180
@@ -1818,7 +1844,6 @@ steps:
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
  torch_nightly: true
-  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -1874,11 +1899,10 @@ steps:
  - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
  - pytest -v -s models/multimodal/processing/test_tensor_schema.py

- label: Multi-Modal Processor (CPU) %N # TBD
+- label: Multi-Modal Processor (CPU) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
-  parallelism: 4
  no_gpu: true
  optional: true
  working_dir: "/vllm-workspace/tests"
@@ -1888,9 +1912,7 @@ steps:
  - tests/models/registry.py
  commands:
  - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
-  - pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
-
-#-----------------------------------------------------  mi300 · models / quantized  -----------------------------------------------------#
+  - pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py

 - label: Quantized Models Test # TBD
  timeout_in_minutes: 180
@@ -1906,31 +1928,7 @@ steps:
  commands:
  - pytest -v -s models/quantization

-#--------------------------------------------------  mi300 · models / transformers  ---------------------------------------------------#
-
- label: Transformers Nightly Models (Shardable) %N # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  agent_pool: mi300_1
-  parallelism: 4
-  optional: true
-  working_dir: "/vllm-workspace/"
-  source_file_dependencies:
-  - vllm/model_executor/models/
-  - vllm/model_executor/model_loader/
-  - vllm/multimodal/
-  - vllm/model_executor/layers/
-  - vllm/v1/attention/backends/
-  - vllm/v1/attention/selector.py
-  - vllm/_aiter_ops.py
-  - vllm/platforms/rocm.py
-  - tests/models/
-  commands:
-  - pip install --upgrade git+https://github.com/huggingface/transformers
-  - pytest -v -s tests/models/test_initialization.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
-  - pytest -v -s tests/models/multimodal/processing/ --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
-
- label: Transformers Nightly Models (Single) # TBD
+- label: Transformers Nightly Models # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
@@ -1949,7 +1947,9 @@ steps:
  - examples/
  commands:
  - pip install --upgrade git+https://github.com/huggingface/transformers
+  - pytest -v -s tests/models/test_initialization.py
  - pytest -v -s tests/models/test_transformers.py
+  - pytest -v -s tests/models/multimodal/processing/
  - pytest -v -s tests/models/multimodal/test_mapping.py
  - python3 examples/basic/offline_inference/chat.py
  - python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl
@@ -2010,7 +2010,7 @@ steps:
  - vllm/model_executor/layers
  - vllm/sampling_metadata.py
  - vllm/v1/sample/
-  - vllm/entrypoints/generate/beam_search/
+  - vllm/beam_search.py
  - tests/samplers
  - tests/conftest.py
  - vllm/_aiter_ops.py
@@ -2204,6 +2204,7 @@ steps:
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
+  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -2280,6 +2281,7 @@ steps:
  - tests/entrypoints/openai/test_multi_api_servers.py
  - vllm/platforms/rocm.py
  commands:
+  - export TORCH_NCCL_BLOCKING_WAIT=1
  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py
  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
@@ -2299,6 +2301,7 @@ steps:
  - vllm/_aiter_ops.py
  - vllm/platforms/rocm.py
  commands:
+  - export TORCH_NCCL_BLOCKING_WAIT=1
  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py
  - VLLM_LOGGING_LEVEL=DEBUG python3 examples/features/data_parallel/data_parallel_offline.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
  - pytest -v -s tests/v1/distributed/test_dbo.py
@@ -2361,6 +2364,7 @@ steps:
  - tests/distributed/test_utils
  - vllm/platforms/rocm.py
  commands:
+  - export TORCH_NCCL_BLOCKING_WAIT=1
  - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py
  - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
  - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
@@ -2398,7 +2402,7 @@ steps:
  - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
  - DP_EP=1 ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh

- label: Hybrid SSM NixlConnector PD accuracy tests (4 GPUs) # TBD
+- label: Hyrbid SSM NixlConnector PD accuracy tests (4 GPUs) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_4
@@ -2490,6 +2494,7 @@ steps:
  - tests/entrypoints/llm/test_collective_rpc.py
  - vllm/platforms/rocm.py
  commands:
+  - export TORCH_NCCL_BLOCKING_WAIT=1
  - pytest -v -s entrypoints/llm/test_collective_rpc.py
  - pytest -v -s ./compile/fullgraph/test_basic_correctness.py
  - pytest -v -s ./compile/test_wrapper.py
@@ -2514,6 +2519,7 @@ steps:
  - tests/v1/worker/test_worker_memory_snapshot.py
  - vllm/platforms/rocm.py
  commands:
+  - export TORCH_NCCL_BLOCKING_WAIT=1
  - VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
  - VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
  - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
@@ -2534,6 +2540,7 @@ steps:
  - tests/distributed/test_multiproc_executor.py
  - vllm/platforms/rocm.py
  commands:
+  - export TORCH_NCCL_BLOCKING_WAIT=1
  - pytest -v -s compile/fullgraph/test_basic_correctness.py
  - pytest -v -s distributed/test_pynccl.py
  - pytest -v -s distributed/test_events.py
@@ -2600,7 +2607,7 @@ steps:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
  - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm.txt --tp-size=8

-#-----------------------------------------------------  mi325 · models / language  -----------------------------------------------------#
+#----------------------------------------------------------  mi325 · models  -----------------------------------------------------------#

 - label: Language Models Test (Extended Generation) # TBD
  timeout_in_minutes: 180
@@ -2621,7 +2628,6 @@ steps:
  agent_pool: mi325_1
  torch_nightly: true
  parallelism: 2
-  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -2631,8 +2637,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

-#----------------------------------------------------  mi325 · models / multimodal  ----------------------------------------------------#
-
 - label: Multi-Modal Models (Extended Pooling) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
@@ -2649,7 +2653,6 @@ steps:
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
  agent_pool: mi325_1
  torch_nightly: true
-  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -2704,35 +2707,20 @@ steps:
  optional: true
  working_dir: "/vllm-workspace/"
  source_file_dependencies:
-  - csrc/custom_quickreduce.cu
-  - csrc/ops.h
-  - csrc/torch_bindings.cpp
  - vllm/distributed/
-  - vllm/model_executor/layers/
-  - vllm/entrypoints/llm.py
-  - vllm/config/parallel.py
-  - vllm/model_executor/layers/fused_moe/
-  - vllm/v1/engine/
-  - vllm/v1/executor/
-  - vllm/v1/worker/
  - vllm/v1/distributed/
+  - vllm/model_executor/layers/fused_moe/
  - vllm/v1/attention/backends/
  - vllm/v1/attention/selector.py
-  - vllm/_aiter_ops.py
-  - vllm/_custom_ops.py
-  - vllm/platforms/rocm.py
-  - vllm/envs.py
-  - examples/offline_inference/data_parallel.py
  - tests/distributed/test_context_parallel.py
-  - tests/distributed/test_rocm_quick_reduce.py
-  - tests/distributed/test_quick_all_reduce.py
  - tests/v1/distributed/test_dbo.py
-  - tests/utils.py
+  - examples/features/data_parallel/data_parallel_offline.py
+  - vllm/_aiter_ops.py
+  - vllm/platforms/rocm.py
  commands:
+  - export TORCH_NCCL_BLOCKING_WAIT=1
  - pytest -v -s tests/distributed/test_context_parallel.py
  - pytest -v -s tests/v1/distributed/test_dbo.py
-  - pytest -v -s tests/distributed/test_rocm_quick_reduce.py
-  - pytest -v -s tests/distributed/test_quick_all_reduce.py

 #--------------------------------------------------------  mi355 · entrypoints  --------------------------------------------------------#

@@ -2746,11 +2734,14 @@ steps:
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
-  - tests/entrypoints/serve
+  - tests/entrypoints/rpc
+  - tests/entrypoints/serve/instrumentator
+  - tests/tool_use
  commands:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc
-  - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
+  - pytest -v -s entrypoints/serve/instrumentator
+  - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc
+  - pytest -v -s tool_use

 - label: Entrypoints Integration (API Server openai - Part 1) # TBD
  timeout_in_minutes: 180
@@ -2758,7 +2749,6 @@ steps:
  agent_pool: mi355_1
  fast_check: true
  torch_nightly: true
-  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -2766,7 +2756,7 @@ steps:
  - tests/entrypoints/test_chat_utils
  commands:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
+  - pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py

 - label: Entrypoints Integration (API Server openai - Part 2) # TBD
  timeout_in_minutes: 180
@@ -2774,20 +2764,16 @@ steps:
  agent_pool: mi355_1
  fast_check: true
  torch_nightly: true
-  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
  - tests/entrypoints/openai
  - tests/entrypoints/test_chat_utils
-  - tests/entrypoints/generate
-  - tests/tool_use
  commands:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
  - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
+  - pytest -v -s entrypoints/openai/speech_to_text/
  - pytest -v -s entrypoints/test_chat_utils.py
-  - pytest -v -s entrypoints/generate
-  - pytest -v -s tool_use

 - label: Entrypoints Integration (API Server openai - Part 3) # TBD
  timeout_in_minutes: 180
@@ -2802,21 +2788,7 @@ steps:
  - tests/entrypoints/test_chat_utils
  commands:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py
-
- label: Entrypoints Integration (Speech to Text) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi355]
-  agent_pool: mi355_1
-  fast_check: true
-  torch_nightly: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/
-  - tests/entrypoints/speech_to_text
-  commands:
-  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/speech_to_text
+  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/speech_to_text/ --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py

 - label: Entrypoints Integration (Pooling) # TBD
  timeout_in_minutes: 180
@@ -2897,7 +2869,7 @@ steps:
  commands:
  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx-fp8-and-mixed.txt

- label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy (B200-MI355) # TBD
+- label: Qwen3-30B-A3B-FP8-block Accuracy (B200-MI355) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
  agent_pool: mi355_2
@@ -2967,8 +2939,8 @@ steps:
  - python3 pooling/embed/vision_embedding_offline.py --seed 0
  # Features demo
  - python3 features/automatic_prefix_caching/prefix_caching_offline.py
-  - python3 deployment/llm_engine_example.py
-  - python3 features/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 features/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
+  - python3 offline_inference/llm_engine_example.py
+  - python3 others/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 others/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
  - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle --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 2048
  - 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

@@ -3069,7 +3041,7 @@ steps:
  commands:
    - pytest -v -s kernels/moe/test_deepep_moe.py

-#-----------------------------------------------------  mi355 · models / language  -----------------------------------------------------#
+#----------------------------------------------------------  mi355 · models  -----------------------------------------------------------#

 - label: Language Models Test (Extended Generation) # TBD
  timeout_in_minutes: 180
@@ -3080,7 +3052,7 @@ steps:
  - vllm/
  - tests/models/language/generation
  commands:
-  - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@rocm-7.0-v2.3.0'
+  - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr'
  - 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)'

@@ -3088,7 +3060,6 @@ steps:
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
  agent_pool: mi355_1
-  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -3137,8 +3108,6 @@ steps:
  - pip freeze | grep -E 'torch'
  - pytest -v -s models/language -m 'core_model and (not slow_test)'

-#----------------------------------------------------  mi355 · models / multimodal  ----------------------------------------------------#
-
 - label: Multi-Modal Models (Extended Generation 1) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
@@ -3210,8 +3179,6 @@ steps:
  - pytest -v -s models/multimodal/generation/test_memory_leak.py -m core_model
  - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model

-#-----------------------------------------------------  mi355 · models / quantized  -----------------------------------------------------#
-
 - label: Quantized Models Test # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
@@ -3292,7 +3259,6 @@ steps:
  timeout_in_minutes: 60
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
  agent_pool: mi355_1
-  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -3318,7 +3284,6 @@ steps:
  timeout_in_minutes: 60
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
  agent_pool: mi355_1
-  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -3359,7 +3324,7 @@ steps:
  - vllm/platforms/rocm.py
  commands:
  - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
-  - ATTENTION_BACKEND=ROCM_ATTN bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
+  - ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh

 - label: Distributed NixlConnector PD accuracy (4 GPUs) # TBD
  timeout_in_minutes: 180
@@ -7,11 +7,7 @@ steps:
  timeout_in_minutes: 15
  device: h200_18gb
  source_file_dependencies:
-  - vllm/envs.py
-  - vllm/logger.py
-  - vllm/platforms/
-  - vllm/plugins/
-  - vllm/utils/
+  - vllm/
  - tests/cuda
  commands:
    - pytest -v -s cuda/test_cuda_context.py
@@ -27,5 +23,4 @@ steps:
  - vllm/compilation
  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_cudagraph_mode.py
@@ -8,10 +8,10 @@ steps:
  working_dir: "/vllm-workspace/tests"
  num_devices: 4
  source_file_dependencies:
-    - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
    - tests/v1/kv_connector/nixl_integration/
  commands:
-    - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
+    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
    - bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
 - label: Distributed FlashInfer NixlConnector PD accuracy (4 GPUs)
  key: distributed-flashinfer-nixlconnector-pd-accuracy-4-gpus
@@ -19,10 +19,10 @@ steps:
  working_dir: "/vllm-workspace/tests"
  num_devices: 4
  source_file_dependencies:
-    - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
    - tests/v1/kv_connector/nixl_integration/
  commands:
-    - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
+    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
    - FLASHINFER=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh

 - label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs)
@@ -31,10 +31,10 @@ steps:
  working_dir: "/vllm-workspace/tests"
  num_devices: 4
  source_file_dependencies:
-    - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
    - tests/v1/kv_connector/nixl_integration/
  commands:
-    - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
+    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
    - DP_EP=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh

 - label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs)
@@ -43,22 +43,22 @@ steps:
  working_dir: "/vllm-workspace/tests"
  num_devices: 4
  source_file_dependencies:
-    - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
    - tests/v1/kv_connector/nixl_integration/
  commands:
-    - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
+    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
    - CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh

 - label: Hybrid SSM NixlConnector PD accuracy tests (4 GPUs)
  key: hybrid-ssm-nixlconnector-pd-accuracy-tests-4-gpus
-  timeout_in_minutes: 25
+  timeout_in_minutes: 20
  working_dir: "/vllm-workspace/tests"
  num_devices: 4
  source_file_dependencies:
-    - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
    - tests/v1/kv_connector/nixl_integration/
  commands:
-    - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
+    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
    - HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh

 - label: MultiConnector (Nixl+Offloading) PD accuracy (2 GPUs)
@@ -67,13 +67,13 @@ steps:
  working_dir: "/vllm-workspace/tests"
  num_devices: 2
  source_file_dependencies:
-    - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
    - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
    - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
    - vllm/distributed/kv_transfer/kv_connector/v1/offloading/
    - tests/v1/kv_connector/nixl_integration/
  commands:
-    - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
+    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
    - bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh

 - label: NixlConnector PD + Spec Decode acceptance (2 GPUs)
@@ -83,12 +83,12 @@ steps:
  working_dir: "/vllm-workspace/tests"
  num_devices: 2
  source_file_dependencies:
-    - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
    - vllm/v1/worker/kv_connector_model_runner_mixin.py
    - tests/v1/kv_connector/nixl_integration/
  commands:
-    - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
-    - bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh
+    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
+    - bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh

 - label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs)
  key: multiconnector-nixl-offloading-pd-edge-cases-2-gpus
@@ -96,11 +96,11 @@ steps:
  working_dir: "/vllm-workspace/tests"
  num_devices: 2
  source_file_dependencies:
-    - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
    - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
    - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
    - vllm/distributed/kv_transfer/kv_connector/v1/offloading/
    - tests/v1/kv_connector/nixl_integration/
  commands:
-    - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
-    - bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
+    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
+    - bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
@@ -2,8 +2,8 @@ group: E2E Integration
 depends_on: 
  - image-build
 steps:
- label: DeepSeek V2-Lite Sync EPLB Accuracy
-  key: deepseek-v2-lite-sync-eplb-accuracy
+- label: DeepSeek V2-Lite Accuracy
+  key: deepseek-v2-lite-accuracy
  timeout_in_minutes: 60
  device: h100
  optional: true
@@ -12,8 +12,8 @@ steps:
  commands:
  - bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh 0.25 200 8010

- label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy
-  key: qwen3-30b-a3b-fp8-block-sync-eplb-accuracy
+- label: Qwen3-30B-A3B-FP8-block Accuracy
+  key: qwen3-30b-a3b-fp8-block-accuracy
  timeout_in_minutes: 60
  device: h100
  optional: true
@@ -22,8 +22,8 @@ steps:
  commands:
  - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020

- label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy (B200)
-  key: qwen3-30b-a3b-fp8-block-sync-eplb-accuracy-b200
+- label: Qwen3-30B-A3B-FP8-block Accuracy (B200)
+  key: qwen3-30b-a3b-fp8-block-accuracy-b200
  timeout_in_minutes: 60
  device: b200-k8s
  optional: true
@@ -7,17 +7,7 @@ steps:
  timeout_in_minutes: 15
  device: h200_18gb
  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/
+  - vllm/
  - tests/engine
  - tests/test_sequence
  - tests/test_config
@@ -36,12 +26,6 @@ steps:
  commands:
    - pytest -v -s v1/engine/test_preprocess_error_handling.py
    - pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py
-  mirror:
-    amd:
-      device: mi325_1
-      timeout_in_minutes: 40
-      depends_on:
-      - image-build-amd

 - label: e2e Scheduling (1 GPU)
  key: e2e-scheduling-1-gpu
@@ -52,15 +36,8 @@ steps:
    - tests/v1/e2e/general/
  commands:
    - pytest -v -s v1/e2e/general/test_async_scheduling.py
-  mirror:
-    amd:
-      device: mi250_1
-      timeout_in_minutes: 60
-      depends_on:
-      - image-build-amd

 - label: e2e Core (1 GPU)
-  device: h200_35gb
  key: e2e-core-1-gpu
  timeout_in_minutes: 30
  source_file_dependencies:
@@ -75,27 +52,16 @@ steps:
  optional: true
  num_devices: 2
  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
+    - vllm/
+    - tests/v1/e2e
  commands:
    # Only run tests that need exactly 2 GPUs
    - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "tensor_parallelism"
+  mirror:
+    amd:
+      device: mi325_2
+      depends_on:
+      - image-build-amd

 - label: V1 e2e (4 GPUs)
  key: v1-e2e-4-gpus
@@ -103,27 +69,16 @@ steps:
  optional: true
  num_devices: 4
  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
+    - vllm/
+    - tests/v1/e2e
  commands:
    # Only run tests that need 4 GPUs
    - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "eagle_correctness_heavy"
+  mirror:
+    amd:
+      device: mi325_4
+      depends_on:
+      - image-build-amd

 - label: V1 e2e (4xH100)
  key: v1-e2e-4xh100
@@ -11,7 +11,7 @@ steps:
  - tests/entrypoints/
  commands:
  - pytest -v -s entrypoints/openai/tool_parsers
-  - pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/offline_mode --ignore=entrypoints/openai --ignore=entrypoints/serve --ignore=entrypoints/test_chat_utils.py  --ignore=entrypoints/pooling --ignore=entrypoints/speech_to_text --ignore=tests/entrypoints/generate
+  - pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/rpc --ignore=entrypoints/sleep --ignore=entrypoints/serve/instrumentator --ignore=entrypoints/openai --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py  --ignore=entrypoints/pooling

 - label: Entrypoints Integration (LLM)
  key: entrypoints-integration-llm
@@ -26,12 +26,6 @@ steps:
  - pytest -v -s entrypoints/llm --ignore=entrypoints/llm/test_generate.py --ignore=entrypoints/llm/test_collective_rpc.py
  - pytest -v -s entrypoints/llm/test_generate.py # it needs a clean process
  - pytest -v -s entrypoints/offline_mode # Needs to avoid interference with other tests
-  mirror:
-    amd:
-      device: mi325_1
-      soft_fail: true
-      depends_on:
-      - image-build-amd

 - label: Entrypoints Integration (API Server openai - Part 1)
  key: entrypoints-integration-api-server-openai-part-1
@@ -43,15 +37,14 @@ steps:
  - tests/entrypoints/test_chat_utils
  commands:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
+  - pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
  mirror:
    amd:
      device: mi325_1
-      soft_fail: true
-      timeout_in_minutes: 80
      depends_on:
      - image-build-amd

+
 - label: Entrypoints Integration (API Server openai - Part 2)
  key: entrypoints-integration-api-server-openai-part-2
  timeout_in_minutes: 50
@@ -60,18 +53,13 @@ steps:
  - vllm/
  - tests/entrypoints/openai
  - tests/entrypoints/test_chat_utils
-  - tests/entrypoints/generate
-  - tests/tool_use
  commands:
  - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
+  - pytest -v -s entrypoints/openai/speech_to_text/
  - pytest -v -s entrypoints/test_chat_utils.py
-  - pytest -v -s entrypoints/generate
-  - pytest -v -s tool_use
  mirror:
    amd:
      device: mi325_1
-      soft_fail: true
-      timeout_in_minutes: 60
      depends_on:
      - image-build-amd

@@ -86,45 +74,22 @@ steps:
  - tests/entrypoints/test_chat_utils
  commands:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py
-  mirror:
-    amd:
-      device: mi325_1
-      soft_fail: true
-      timeout_in_minutes: 60
-      depends_on:
-      - image-build-amd
+  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/speech_to_text/ --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py

 - label: Entrypoints Integration (API Server 2)
-  device: h200_35gb
  key: entrypoints-integration-api-server-2
  timeout_in_minutes: 130
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
-  - tests/entrypoints/serve
+  - tests/entrypoints/rpc
+  - tests/entrypoints/serve/instrumentator
+  - tests/tool_use
  commands:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc
-  - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
-  mirror:
-    amd:
-      device: mi325_1
-      soft_fail: true
-      depends_on:
-      - image-build-amd
-
- label: Entrypoints Integration (Speech to Text)
-  device: h200_35gb
-  key: entrypoints-integration-speech_to_text
-  timeout_in_minutes: 50
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/
-  - tests/entrypoints/speech_to_text
-  commands:
-  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/speech_to_text
+  - pytest -v -s entrypoints/serve/instrumentator
+  - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc
+  - pytest -v -s tool_use

 - label: Entrypoints Integration (Pooling)
  key: entrypoints-integration-pooling
@@ -154,5 +119,6 @@ steps:
  source_file_dependencies:
  - csrc/
  - vllm/entrypoints/openai/
-  commands: # LMEval
+  - vllm/model_executor/models/whisper.py
+  commands: # LMEval+Transcription WER check
  - pytest -s entrypoints/openai/correctness/
@@ -38,28 +38,6 @@ steps:
  commands:
    - pytest -v -s kernels/core/test_minimax_reduce_rms.py

- label: Deepseek V4 Kernel Test (H100)
-  key: deepseek-v4-kernel-test-h100
-  timeout_in_minutes: 15
-  device: h100
-  source_file_dependencies:
-  - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
-  - vllm/models/deepseek_v4/common/ops/
-  - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
-  commands:
-    - pytest -v -s kernels/test_fused_deepseek_v4_*.py
-
- label: Deepseek V4 Kernel Test (B200)
-  key: deepseek-v4-kernel-test-b200
-  timeout_in_minutes: 15
-  device: b200-k8s
-  source_file_dependencies:
-  - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
-  - vllm/models/deepseek_v4/common/ops/
-  - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
-  commands:
-    - pytest -v -s kernels/test_fused_deepseek_v4_*.py
-
 - label: Kernels Attention Test %N
  key: kernels-attention-test
  timeout_in_minutes: 35
@@ -84,19 +62,6 @@ steps:
  commands:
    - pytest -v -s kernels/quantization --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
  parallelism: 2
-  mirror:
-    amd:
-      device: mi325_1
-      source_file_dependencies:
-      - csrc/quantization/
-      - vllm/model_executor/layers/quantization
-      - tests/kernels/quantization
-      - tests/kernels/quantization/test_rocm_skinny_gemms.py
-      - vllm/_aiter_ops.py
-      - vllm/platforms/rocm.py
-      - vllm/model_executor/kernels/
-      depends_on:
-      - image-build-amd

 - label: Kernels MoE Test %N
  key: kernels-moe-test
@@ -126,7 +91,6 @@ steps:

 - label: Kernels KDA Test
  timeout_in_minutes: 20
-  device: h200_18gb
  source_file_dependencies:
  - vllm/model_executor/layers/fla/ops/kda.py
  - vllm/model_executor/layers/fla/ops/chunk_delta_h.py
@@ -141,11 +105,7 @@ steps:
  device: h100
  num_devices: 1
  source_file_dependencies:
-  - cmake/external_projects/deepgemm.cmake
  - tools/install_deepgemm.sh
-  - tools/build_deepgemm_C.py
-  - tools/setup_deepgemm_pythons.sh
-  - tools/check_wheel_deepgemm.py
  - vllm/utils/deep_gemm.py
  - vllm/model_executor/layers/fused_moe
  - vllm/model_executor/layers/quantization
@@ -155,7 +115,6 @@ steps:
  - tests/kernels/attention/test_deepgemm_attention.py
  - tests/quantization/test_cutlass_w4a16.py
  commands:
-    - python3 ../tools/check_wheel_deepgemm.py
    - pytest -v -s kernels/quantization/test_block_fp8.py
    - pytest -v -s kernels/moe/test_deepgemm.py
    - pytest -v -s kernels/moe/test_batched_deepgemm.py
@@ -172,12 +131,9 @@ steps:
  - csrc/quantization/fp4/
  - csrc/attention/mla/
  - csrc/quantization/cutlass_w8a8/moe/
-  - vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py
-  - vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py
-  - vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py
-  - vllm/model_executor/layers/fused_moe/oracle/nvfp4.py
-  - vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_one_sided.py
-  - vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_two_sided.py
+  - vllm/model_executor/layers/fused_moe/cutlass_moe.py
+  - vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py
+  - vllm/model_executor/layers/fused_moe/flashinfer_a2a_prepare_finalize.py
  - vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
  - vllm/v1/attention/backends/flashinfer.py
  - vllm/v1/attention/backends/mla/cutlass_mla.py
@@ -210,7 +166,6 @@ steps:
    - pytest -v -s tests/kernels/moe/test_ocp_mx_moe.py
    - pytest -v -s tests/kernels/moe/test_flashinfer.py
    - 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
    # e2e
    - pytest -v -s tests/models/quantization/test_nvfp4.py
@@ -287,7 +242,7 @@ steps:
 - label: Kernels FusedMoE Layer Test (2 B200s)
  key: kernels-fusedmoe-layer-test-2-b200s
  timeout_in_minutes: 90
-  device: b200-k8s
+  device: b200
  num_devices: 2
  source_file_dependencies:
  - csrc/quantization/cutlass_w8a8/moe/
@@ -3,7 +3,6 @@ depends_on:
  - image-build
 steps:
 - label: LM Eval Small Models
-  device: h200_35gb
  key: lm-eval-small-models
  timeout_in_minutes: 75
  source_file_dependencies:
@@ -41,7 +40,7 @@ steps:
 - label: LM Eval Small Models (B200)
  key: lm-eval-small-models-b200
  timeout_in_minutes: 120
-  device: b200-k8s
+  device: b200
  optional: true
  source_file_dependencies:
  - csrc/
@@ -49,18 +48,6 @@ steps:
  commands:
  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-blackwell.txt

- label: LM Eval Large Models (B200, EP)
-  key: lm-eval-large-models-b200-ep
-  timeout_in_minutes: 120
-  device: b200-k8s
-  optional: true
-  num_devices: 2
-  source_file_dependencies:
-  - csrc/
-  - vllm/model_executor/layers/quantization
-  commands:
-  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-blackwell-ep.txt
-
 - label: LM Eval Qwen3.5 Models (B200)
  key: lm-eval-qwen3-5-models-b200
  timeout_in_minutes: 120
@@ -105,7 +92,7 @@ steps:

 - label: MoE Refactor Integration Test (B200 DP - TEMPORARY)
  key: moe-refactor-integration-test-b200-dp-temporary
-  device: b200-k8s
+  device: b200
  optional: true
  num_devices: 2
  commands:
@@ -115,7 +102,6 @@ steps:
 - label: LM Eval TurboQuant KV Cache
  key: lm-eval-turboquant-kv-cache
  timeout_in_minutes: 75
-  device: h200_18gb
  source_file_dependencies:
  - vllm/model_executor/layers/quantization/turboquant/
  - vllm/v1/attention/backends/turboquant_attn.py
@@ -141,7 +127,7 @@ steps:
 - label: GPQA Eval (GPT-OSS) (B200)
  key: gpqa-eval-gpt-oss-b200
  timeout_in_minutes: 120
-  device: b200-k8s
+  device: b200
  optional: true
  num_devices: 2
  source_file_dependencies:
@@ -153,7 +139,6 @@ steps:
    - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-b200.txt

 - label: MRCR Eval Small Models
-  device: h200_35gb
  timeout_in_minutes: 30
  source_file_dependencies:
  - tests/evals/mrcr/
@@ -3,14 +3,13 @@ depends_on:
  - image-build
 steps:
 - label: LoRA %N
-  device: h200_35gb
  key: lora
  timeout_in_minutes: 30
  source_file_dependencies:
  - vllm/lora
  - tests/lora
  commands:
-    - pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_qwen3_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py --ignore=lora/test_qwen35_densemodel_lora.py 
+    - pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_llm_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py --ignore=lora/test_qwen35_densemodel_lora.py 
  parallelism: 4


@@ -20,7 +19,6 @@ steps:
  num_devices: 4
  source_file_dependencies:
  - vllm/lora
-  - vllm/model_executor/layers/fused_moe/
  - tests/lora
  commands:
    # FIXIT: find out which code initialize cuda before running the test
@@ -32,7 +30,7 @@ steps:
    # requires multi-GPU testing for validation.
    - pytest -v -s -x lora/test_chatglm3_tp.py
    - pytest -v -s -x lora/test_llama_tp.py
-    - pytest -v -s -x lora/test_qwen3_with_multi_loras.py
+    - pytest -v -s -x lora/test_llm_with_multi_loras.py
    - pytest -v -s -x lora/test_olmoe_tp.py
    - pytest -v -s -x lora/test_gptoss_tp.py
    - pytest -v -s -x lora/test_qwen35_densemodel_lora.py
@@ -3,41 +3,27 @@ depends_on:
  - image-build
 steps:
 - label: V1 Spec Decode
-  device: h200_35gb
  key: v1-spec-decode
  timeout_in_minutes: 30
  source_file_dependencies:
-    - vllm/config/
-    - vllm/distributed/
-    - vllm/inputs/
-    - vllm/model_executor/
-    - vllm/platforms/
-    - vllm/sampling_params.py
-    - vllm/transformers_utils/
-    - vllm/utils/
-    - vllm/v1/
+    - vllm/
    - tests/v1/spec_decode
  commands:
    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
    # TODO: create another `optional` test group for slow tests
    - pytest -v -s -m 'not slow_test' v1/spec_decode
+  mirror:
+    amd:
+      device: mi325_1
+      depends_on:
+      - image-build-amd

 - label: V1 Sample + Logits
  key: v1-sample-logits
  timeout_in_minutes: 30
  device: h200_18gb
  source_file_dependencies:
-    - vllm/config/
-    - vllm/distributed/
-    - vllm/engine/
-    - vllm/inputs/
-    - vllm/logger.py
-    - vllm/model_executor/
-    - vllm/platforms/
-    - vllm/sampling_params.py
-    - vllm/transformers_utils/
-    - vllm/utils/
-    - vllm/v1/
+    - vllm/
    - tests/v1/sample
    - tests/v1/logits_processors
    - tests/v1/test_oracle.py
@@ -60,66 +46,39 @@ steps:
  key: v1-core-kv-metrics
  timeout_in_minutes: 30
  source_file_dependencies:
-    - vllm/config/
-    - vllm/distributed/
-    - vllm/engine/
-    - vllm/entrypoints/pooling/
-    - vllm/inputs/
-    - vllm/lora/
-    - vllm/model_executor/
-    - vllm/multimodal/
-    - vllm/outputs.py
-    - vllm/platforms/
-    - vllm/pooling_params.py
-    - vllm/profiler/
-    - vllm/sampling_params.py
-    - vllm/tokenizers/
-    - vllm/transformers_utils/
-    - vllm/utils/
-    - vllm/v1/
+    - vllm/
    - tests/v1/core
    - tests/v1/executor
    - tests/v1/kv_offload
-    - tests/v1/simple_kv_offload
    - tests/v1/worker
    - tests/v1/kv_connector/unit
    - tests/v1/metrics
    - tests/entrypoints/openai/correctness/test_lmeval.py
  commands:
-    - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
+    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
    # split the test to avoid interference
    - pytest -v -s -m 'not cpu_test' v1/core
    - pytest -v -s v1/executor
    - pytest -v -s v1/kv_offload
-    - pytest -v -s v1/simple_kv_offload
    - pytest -v -s v1/worker
    - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
    - pytest -v -s -m 'not cpu_test' v1/metrics
    # Integration test for streaming correctness (requires special branch).
    - pip install -U git+https://github.com/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api
    - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
+  mirror:
+    amd:
+      device: mi325_1
+      depends_on:
+      - image-build-amd

 - label: V1 Others (CPU)
  key: v1-others-cpu
  depends_on:
    - image-build-cpu
  source_file_dependencies:
-    - vllm/config/
-    - vllm/distributed/
-    - vllm/engine/
-    - vllm/inputs/
-    - vllm/lora/
-    - vllm/multimodal/
-    - vllm/outputs.py
-    - vllm/platforms/
-    - vllm/pooling_params.py
-    - vllm/profiler/
-    - vllm/sampling_params.py
-    - vllm/tokenizers/
-    - vllm/transformers_utils/
-    - vllm/utils/
-    - vllm/v1/
+    - vllm/
    - tests/v1
  device: cpu-small
  commands:
@@ -130,35 +89,12 @@ steps:
    - pytest -v -s -m 'cpu_test' v1/kv_connector/unit
    - pytest -v -s -m 'cpu_test' v1/metrics

- label: Extract Hidden States Integration
-  key: extract-hidden-states-integration
-  timeout_in_minutes: 20
-  device: h200_18gb
-  source_file_dependencies:
-    - vllm/v1/spec_decode/extract_hidden_states.py
-    - vllm/model_executor/models/extract_hidden_states.py
-    - vllm/transformers_utils/configs/extract_hidden_states.py
-    - tests/v1/kv_connector/extract_hidden_states_integration
-  commands:
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - pytest -v -s v1/kv_connector/extract_hidden_states_integration
-
 - label: Regression
  key: regression
  timeout_in_minutes: 20
  device: h200_18gb
  source_file_dependencies:
-  - vllm/config/
-  - vllm/distributed/
-  - vllm/engine/
-  - vllm/inputs/
-  - vllm/model_executor/
-  - vllm/multimodal/
-  - vllm/platforms/
-  - vllm/sampling_params.py
-  - vllm/transformers_utils/
-  - vllm/utils/
-  - vllm/v1/
+  - vllm/
  - tests/test_regression
  commands:
  - pip install modelscope
@@ -166,7 +102,6 @@ steps:
  working_dir: "/vllm-workspace/tests" # optional

 - label: Examples
-  device: h200_35gb
  key: examples
  timeout_in_minutes: 45
  working_dir: "/vllm-workspace/examples"
@@ -192,8 +127,8 @@ steps:
    - python3 pooling/embed/vision_embedding_offline.py --seed 0
    # for features demo
    - python3 features/automatic_prefix_caching/prefix_caching_offline.py
-    - python3 deployment/llm_engine_example.py
-    - python3 features/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 features/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
+    - python3 offline_inference/llm_engine_example.py
+    - python3 others/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 others/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
    - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle --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 2048
    # https://github.com/vllm-project/vllm/pull/26682 uses slightly more memory in PyTorch 2.9+ causing this test to OOM in 1xL4 GPU
    - 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
@@ -203,18 +138,7 @@ steps:
  timeout_in_minutes: 20
  num_devices: 2
  source_file_dependencies:
-  - vllm/config/
-  - vllm/distributed/
-  - vllm/engine/
-  - vllm/inputs/
-  - vllm/model_executor/
-  - vllm/multimodal/
-  - vllm/platforms/
-  - vllm/sampling_params.py
-  - vllm/tracing/
-  - vllm/transformers_utils/
-  - vllm/utils/
-  - vllm/v1/
+  - vllm/
  - tests/v1/tracing
  commands:
  - "pip install \
@@ -227,9 +151,7 @@ steps:
 - label: Python-only Installation
  key: python-only-installation
  depends_on: ~
-  optional: true
  timeout_in_minutes: 20
-  device: h200_18gb
  source_file_dependencies:
  - tests/standalone_tests/python_only_compile.sh
  - setup.py
@@ -237,23 +159,10 @@ steps:
  - bash standalone_tests/python_only_compile.sh

 - label: Async Engine, Inputs, Utils, Worker
-  device: h200_35gb
  key: async-engine-inputs-utils-worker
  timeout_in_minutes: 50
  source_file_dependencies:
-  - vllm/assets/
-  - vllm/config/
-  - vllm/distributed/
-  - vllm/engine/
-  - vllm/inputs/
-  - vllm/model_executor/
-  - vllm/multimodal/
-  - vllm/platforms/
-  - vllm/sampling_params.py
-  - vllm/tokenizers/
-  - vllm/transformers_utils/
-  - vllm/utils/
-  - vllm/v1/
+  - vllm/
  - tests/detokenizer
  - tests/multimodal
  - tests/utils_
@@ -268,30 +177,7 @@ steps:
  - image-build-cpu
  timeout_in_minutes: 30
  source_file_dependencies:
-  - vllm/assets/
-  - vllm/config/
-  - vllm/engine/arg_utils.py
-  - vllm/entrypoints/chat_utils.py
-  - vllm/entrypoints/mcp/
-  - vllm/entrypoints/openai/chat_completion/protocol.py
-  - vllm/entrypoints/openai/engine/protocol.py
-  - vllm/envs.py
-  - vllm/exceptions.py
-  - vllm/inputs/
-  - vllm/model_executor/layers/quantization/quark/
-  - vllm/multimodal/
-  - vllm/outputs.py
-  - vllm/platforms/
-  - vllm/pooling_params.py
-  - vllm/ray/
-  - vllm/reasoning/
-  - vllm/renderers/
-  - vllm/sampling_params.py
-  - vllm/tokenizers/
-  - vllm/tool_parsers/
-  - vllm/transformers_utils/
-  - vllm/utils/
-  - vllm/v1/
+  - vllm/
  - tests/test_inputs.py
  - tests/test_outputs.py
  - tests/test_pooling_params.py
@@ -302,7 +188,6 @@ steps:
  - tests/tokenizers_
  - tests/reasoning
  - tests/tool_parsers
-  - tests/parser
  - tests/transformers_utils
  - tests/config
  device: cpu-small
@@ -315,26 +200,11 @@ steps:
  - pytest -v -s -m 'cpu_test' multimodal
  - pytest -v -s renderers
  - pytest -v -s tokenizers_
-  - pytest -v -s reasoning --ignore=reasoning/test_seedoss_reasoning_parser.py --ignore=reasoning/test_glm4_moe_reasoning_parser.py
+  - pytest -v -s reasoning --ignore=reasoning/test_seedoss_reasoning_parser.py --ignore=reasoning/test_glm4_moe_reasoning_parser.py --ignore=reasoning/test_gemma4_reasoning_parser.py
  - pytest -v -s tool_parsers
-  - pytest -v -s parser
  - pytest -v -s transformers_utils
  - pytest -v -s config

- label: Batch Invariance (A100)
-  key: batch-invariance-a100
-  timeout_in_minutes: 30
-  device: a100
-  source_file_dependencies:
-    - vllm/v1/attention
-    - vllm/model_executor/layers
-    - tests/v1/determinism/
-  commands:
-    - 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[TRITON_MLA]
-
 - label: Batch Invariance (H100)
  key: batch-invariance-h100
  timeout_in_minutes: 30
@@ -367,10 +237,8 @@ steps:
    - 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[FLASH_ATTN]
    - pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py
-    - pytest -v -s v1/determinism/test_nvfp4_batch_invariant_scaled_mm.py
  
 - label: Acceptance Length Test (Large Models) # optional
-  device: h200_35gb
  key: acceptance-length-test-large-models
  timeout_in_minutes: 25
  gpu: h100
@@ -14,12 +14,5 @@ steps:
  commands:
    - apt-get update && apt-get install -y curl libsodium23
    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    # Dump tracebacks of all threads if a test hangs, so a wedged GPU/CUDA
-    # init surfaces a stack instead of silently stalling.
-    - export PYTHONFAULTHANDLER=1
-    # Per-test watchdog: a single hung test (e.g. stuck during engine/CUDA
-    # init) fails fast with a traceback instead of running until the global
-    # build timeout. The `thread` method also handles hangs inside C/CUDA
-    # calls that the signal method cannot interrupt.
-    - pytest -v -s model_executor -m '(not slow_test)' --timeout=900 --timeout-method=thread
-    - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py --timeout=900 --timeout-method=thread
+    - pytest -v -s model_executor -m '(not slow_test)'
+    - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py
@@ -3,7 +3,6 @@ depends_on:
  - image-build
 steps:
 - label: Model Runner V2 Core Tests
-  device: h200_35gb
  key: model-runner-v2-core-tests
  timeout_in_minutes: 45
  source_file_dependencies:
@@ -27,7 +26,6 @@ steps:
  - 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"

 - label: Model Runner V2 Examples
-  device: h200_35gb
  key: model-runner-v2-examples
  timeout_in_minutes: 45
  working_dir: "/vllm-workspace/examples"
@@ -39,7 +37,7 @@ steps:
    - examples/generate/multimodal/
    - examples/features/
    - examples/pooling/embed/vision_embedding_offline.py
-    - examples/features/tensorize_vllm_model.py
+    - examples/others/tensorize_vllm_model.py
  commands:
    - set -x
    - export VLLM_USE_V2_MODEL_RUNNER=1
@@ -57,8 +55,8 @@ steps:
    - python3 pooling/embed/vision_embedding_offline.py --seed 0
    # for features demo
    - python3 features/automatic_prefix_caching/prefix_caching_offline.py
-    - python3 deployment/llm_engine_example.py
-    - python3 features/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 features/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
+    - python3 offline_inference/llm_engine_example.py
+    - python3 others/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 others/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
    - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle --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 2048
    # https://github.com/vllm-project/vllm/pull/26682 uses slightly more memory in PyTorch 2.9+ causing this test to OOM in 1xL4 GPU
    - 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
@@ -101,7 +99,6 @@ steps:
    - pytest -v -s distributed/test_pp_cudagraph.py -k "not ray"

 - label: Model Runner V2 Spec Decode
-  device: h200_35gb
  key: model-runner-v2-spec-decode
  timeout_in_minutes: 30
  working_dir: "/vllm-workspace/tests"
@@ -109,12 +106,13 @@ steps:
  - 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/spec_decode/test_probabilistic_rejection_sampler_utils.py
+  - tests/v1/spec_decode/test_synthetic_rejection_sampler_utils.py
  - tests/v1/e2e/spec_decode/test_spec_decode.py
  commands:
  - set -x
  - export VLLM_USE_V2_MODEL_RUNNER=1
  - pytest -v -s v1/spec_decode/test_max_len.py -k "eagle or mtp"
-  - pytest -v -s v1/spec_decode/test_rejection_sampler_utils.py
+  - pytest -v -s v1/spec_decode/test_probabilistic_rejection_sampler_utils.py
  - pytest -v -s v1/spec_decode/test_synthetic_rejection_sampler_utils.py
  - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "eagle or mtp"
@@ -5,7 +5,6 @@ steps:
 - label: Basic Models Tests (Initialization)
  key: basic-models-tests-initialization
  timeout_in_minutes: 45
-  device: h200_18gb
  torch_nightly: true
  source_file_dependencies:
  - vllm/
@@ -18,7 +17,6 @@ steps:
    torch_nightly: {}

 - label: Basic Models Tests (Extra Initialization) %N
-  device: h200_35gb
  key: basic-models-tests-extra-initialization
  timeout_in_minutes: 45
  source_file_dependencies:
@@ -35,7 +33,6 @@ steps:
    torch_nightly: {}

 - label: Basic Models Tests (Other)
-  device: h200_35gb
  key: basic-models-tests-other
  timeout_in_minutes: 45
  source_file_dependencies:
@@ -45,6 +42,12 @@ steps:
  - tests/models/test_registry.py
  commands:
    - pytest -v -s models/test_terratorch.py models/test_transformers.py models/test_registry.py
+  mirror:
+    amd:
+      device: mi325_1
+      depends_on:
+      - image-build-amd
+

 - label: Basic Models Test (Other CPU) # 5min
  key: basic-models-test-other-cpu
@@ -58,3 +61,35 @@ steps:
  device: cpu-small
  commands:
    - pytest -v -s models/test_utils.py models/test_vision.py
+
+- label: Transformers Nightly Models
+  key: transformers-nightly-models
+  working_dir: "/vllm-workspace/"
+  optional: true
+  soft_fail: true
+  commands:
+    - pip install --upgrade git+https://github.com/huggingface/transformers
+    - pytest -v -s tests/models/test_initialization.py
+    - pytest -v -s tests/models/test_transformers.py
+    - pytest -v -s tests/models/multimodal/processing/
+    - pytest -v -s tests/models/multimodal/test_mapping.py
+    - python3 examples/basic/offline_inference/chat.py
+    - python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl
+    # Whisper needs spawn method to avoid deadlock
+    - VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/generate/multimodal/audio_language_offline.py --model-type whisper
+
+- label: Transformers Backward Compatibility Models Test
+  key: transformers-backward-compatibility-models-test
+  working_dir: "/vllm-workspace/"
+  optional: true
+  soft_fail: true
+  commands:
+    - pip install transformers==4.57.5
+    - pytest -v -s tests/models/test_initialization.py
+    - pytest -v -s tests/models/test_transformers.py
+    - pytest -v -s tests/models/multimodal/processing/
+    - pytest -v -s tests/models/multimodal/test_mapping.py
+    - python3 examples/basic/offline_inference/chat.py
+    - python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl
+    # Whisper needs spawn method to avoid deadlock
+    - VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/generate/multimodal/audio_language_offline.py --model-type whisper
@@ -5,7 +5,6 @@ steps:
 - label: Language Models Tests (Standard)
  key: language-models-tests-standard
  timeout_in_minutes: 25
-  device: h200_18gb
  source_file_dependencies:
  - vllm/
  - tests/models/language
@@ -49,17 +48,8 @@ steps:
  parallelism: 2
  mirror:
    torch_nightly: {}
-    amd:
-      device: mi325_1
-      depends_on:
-      - image-build-amd
-      commands:
-      - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr'
-      - 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

 - label: Language Models Test (Extended Generation) # 80min
-  device: h200_35gb
  key: language-models-test-extended-generation
  timeout_in_minutes: 110
  optional: true
@@ -72,6 +62,15 @@ steps:
    - 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)'
+  mirror:
+    amd:
+      device: mi325_1
+      depends_on:
+      - image-build-amd
+      commands:
+      - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr'
+      - 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)'

 - label: Language Models Test (PPL)
  key: language-models-test-ppl
@@ -85,7 +84,6 @@ steps:
    - pytest -v -s models/language/generation_ppl_test

 - label: Language Models Test (Extended Pooling)  # 36min
-  device: h200_35gb
  key: language-models-test-extended-pooling
  timeout_in_minutes: 50
  optional: true
@@ -97,7 +95,6 @@ steps:
  mirror:
    amd:
      device: mi325_1
-      timeout_in_minutes: 100
      depends_on:
      - image-build-amd

@@ -38,7 +38,6 @@ steps:
      - image-build-amd

 - label: "Multi-Modal Models (Standard) 3: llava + qwen2_vl"
-  device: h200_35gb
  key: multi-modal-models-standard-3-llava-qwen2-vl
  timeout_in_minutes: 45
  source_file_dependencies:
@@ -55,7 +54,6 @@ steps:
      - 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: 45
  source_file_dependencies:
@@ -66,6 +64,11 @@ steps:
    - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py  --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/processing
    - pytest models/multimodal/generation/test_memory_leak.py -m core_model
    - 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:
+      device: mi325_1
+      depends_on:
+      - image-build-amd

 - label: Multi-Modal Processor (CPU)
  key: multi-modal-processor-cpu
@@ -94,7 +97,6 @@ steps:
    - pytest -v -s models/multimodal/processing/test_tensor_schema.py

 - label: Multi-Modal Accuracy Eval (Small Models) # 50min
-  device: h200_35gb
  key: multi-modal-accuracy-eval-small-models
  timeout_in_minutes: 70
  working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
@@ -123,7 +125,6 @@ steps:
      - image-build-amd

 - label: Multi-Modal Models (Extended Generation 2)
-  device: h200_35gb
  key: multi-modal-models-extended-generation-2
  optional: true
  source_file_dependencies:
@@ -134,7 +135,6 @@ steps:
    - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=0) and not core_model'

 - label: Multi-Modal Models (Extended Generation 3)
-  device: h200_35gb
  key: multi-modal-models-extended-generation-3
  optional: true
  source_file_dependencies:
@@ -3,34 +3,10 @@ depends_on:
  - image-build
 steps:
 - label: PyTorch Compilation Unit Tests
-  device: h200_35gb
  key: pytorch-compilation-unit-tests
  timeout_in_minutes: 10
  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/
+    - vllm/
    - tests/compile
  commands:
  # Run unit tests defined directly under compile/,
@@ -48,30 +24,7 @@ steps:
  device: h100
  num_devices: 1
  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/
+    - vllm/
    - tests/compile/h100/
  commands:
  - "find compile/h100/ -name 'test_*.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
@@ -80,30 +33,7 @@ steps:
  key: pytorch-compilation-passes-unit-tests
  timeout_in_minutes: 20
  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/
+    - vllm/
    - tests/compile/passes
  commands:
  - pytest -s -v compile/passes --ignore compile/passes/distributed
@@ -112,30 +42,7 @@ steps:
  key: pytorch-fullgraph-smoke-test
  timeout_in_minutes: 35
  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/
+  - vllm/
  - tests/compile
  commands:
  # Run smoke tests under fullgraph directory, except test_full_graph.py
@@ -149,30 +56,7 @@ steps:
  timeout_in_minutes: 30
  device: h200_18gb
  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/
+  - vllm/
  - tests/compile
  commands:
    # fp8 kv scales not supported on sm89, tested on Blackwell instead
@@ -1,107 +0,0 @@
-group: Rust Frontend E2E
-depends_on:
-  - image-build
-steps:
- label: Rust Frontend OpenAI Coverage
-  timeout_in_minutes: 90
-  device: h200_18gb
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - rust/
-  - vllm/benchmarks/
-  - vllm/entrypoints/openai/
-  - vllm/entrypoints/serve/
-  - vllm/v1/sample/
-  - tests/utils.py
-  - tests/benchmarks/test_serve_cli.py
-  - tests/entrypoints/openai/chat_completion/test_chat_completion.py
-  # - tests/entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py
-
-  # - tests/entrypoints/openai/completion/test_prompt_validation.py
-  - tests/entrypoints/openai/completion/test_shutdown.py
-  # - tests/entrypoints/openai/test_return_token_ids.py
-  # - tests/entrypoints/openai/test_uds.py
-  - tests/v1/sample/test_logprobs_e2e.py
-  commands:
-  - 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
-  # - pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not invalid"
-
-  # - pytest -v -s entrypoints/openai/completion/test_prompt_validation.py -k "not prompt_embeds"
-  - pytest -v -s entrypoints/openai/completion/test_shutdown.py -k "not engine_failure and not test_abort_timeout_exits_quickly"
-  # - pytest -v -s entrypoints/openai/test_return_token_ids.py
-  # - pytest -v -s entrypoints/openai/test_uds.py
-  - pytest -v -s v1/sample/test_logprobs_e2e.py -k "test_prompt_logprobs_e2e_server"
-
- label: Rust Frontend Serve/Admin Coverage
-  timeout_in_minutes: 60
-  device: h200_18gb
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - rust/
-  - vllm/entrypoints/openai/
-  - vllm/entrypoints/serve/
-  - vllm/v1/engine/
-  - tests/utils.py
-  # - tests/entrypoints/serve/dev/rpc/test_collective_rpc.py
-  - tests/entrypoints/serve/disagg/test_serving_tokens.py
-  - tests/entrypoints/serve/instrumentator/test_basic.py
-  - tests/entrypoints/serve/instrumentator/test_metrics.py
-  # - tests/entrypoints/serve/dev/test_sleep.py
-  commands:
-  - export VLLM_USE_RUST_FRONTEND=1
-  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  # - pytest -v -s entrypoints/serve/dev/rpc/test_collective_rpc.py
-  - pytest -v -s entrypoints/serve/instrumentator/test_basic.py -k "not show_version and not server_load"
-  - pytest -v -s entrypoints/serve/disagg/test_serving_tokens.py -k "not stream and not lora and not test_generate_logprobs and not stop_string_workflow"
-  - pytest -v -s entrypoints/serve/instrumentator/test_metrics.py -k "text and not show and not run_batch and not test_metrics_counts and not test_metrics_exist"
-  # - pytest -v -s entrypoints/serve/dev/test_sleep.py
-
- label: Rust Frontend Core Correctness
-  timeout_in_minutes: 30
-  device: h200_18gb
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - rust/
-  - vllm/entrypoints/openai/
-  - tests/utils.py
-  - tests/entrypoints/openai/correctness/test_lmeval.py
-  commands:
-  - export VLLM_USE_RUST_FRONTEND=1
-  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
-
- label: Rust Frontend Tool Use
-  timeout_in_minutes: 60
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - rust/
-  - vllm/entrypoints/openai/
-  - vllm/tool_parsers/
-  - tests/utils.py
-  - tests/tool_use/
-  commands:
-  - export VLLM_USE_RUST_FRONTEND=1
-  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s tool_use --ignore=tool_use/mistral --models llama3.2 -k "not test_response_format_with_tool_choice_required and not test_parallel_tool_calls_false and not test_tool_call_and_choice"
-
- label: Rust Frontend Distributed
-  timeout_in_minutes: 30
-  num_devices: 4
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - rust/
-  - vllm/distributed/
-  - vllm/engine/
-  - vllm/executor/
-  - vllm/v1/engine/
-  - vllm/v1/worker/
-  - tests/utils.py
-  - tests/v1/distributed/test_internal_lb_dp.py
-  commands:
-  - export VLLM_USE_RUST_FRONTEND=1
-  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - export NCCL_CUMEM_HOST_ENABLE=0
-  - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py -k "not 4 and not server_info"
@@ -1,30 +0,0 @@
-group: Rust Frontend Cargo
-depends_on: []
-steps:
- label: Rust Frontend Cargo Style + Clippy
-  key: rust-frontend-cargo-style-clippy
-  depends_on: []
-  timeout_in_minutes: 30
-  device: cpu-medium
-  no_plugin: true
-  source_file_dependencies:
-  - rust/
-  - rust-toolchain.toml
-  - .buildkite/test_areas/rust_frontend_cargo.yaml
-  - .buildkite/scripts/run-rust-frontend-cargo-ci.sh
-  commands:
-  - .buildkite/scripts/run-rust-frontend-cargo-ci.sh style-clippy
-
- label: Rust Frontend Cargo Tests
-  key: rust-frontend-cargo-tests
-  depends_on: []
-  timeout_in_minutes: 30
-  device: cpu-medium
-  no_plugin: true
-  source_file_dependencies:
-  - rust/
-  - rust-toolchain.toml
-  - .buildkite/test_areas/rust_frontend_cargo.yaml
-  - .buildkite/scripts/run-rust-frontend-cargo-ci.sh
-  commands:
-  - .buildkite/scripts/run-rust-frontend-cargo-ci.sh test
@@ -3,7 +3,6 @@ depends_on:
  - image-build
 steps:
 - label: Samplers Test
-  device: h200_35gb
  key: samplers-test
  timeout_in_minutes: 75
  source_file_dependencies:
@@ -11,7 +10,6 @@ steps:
  - vllm/sampling_metadata.py
  - tests/samplers
  - tests/conftest.py
-  - vllm/entrypoints/generate/beam_search
  commands:
    # VLLM_USE_FLASHINFER_SAMPLER defaults to 1 now, so we need to pin both
    # values explicitly to still cover the PyTorch-native (Triton) path.
@@ -19,7 +17,7 @@ steps:
    - VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers
  mirror:
    amd:
-      device: mi250_1
+      device: mi325_1
      depends_on:
      - image-build-amd
      commands:
@@ -16,7 +16,7 @@ steps:
 - label: Spec Decode Eagle Nightly B200
  key: spec-decode-eagle-nightly-b200
  timeout_in_minutes: 30
-  device: b200-k8s
+  device: b200
  optional: true
  source_file_dependencies:
    - vllm/v1/spec_decode/
@@ -32,7 +32,6 @@ steps:
  source_file_dependencies:
    - vllm/v1/spec_decode/
    - vllm/v1/worker/gpu/spec_decode/
-    - vllm/v1/attention/backends/
    - vllm/transformers_utils/configs/speculators/
    - tests/v1/e2e/spec_decode/
  commands:
@@ -41,7 +40,7 @@ steps:
 - label: Spec Decode Speculators + MTP Nightly B200
  key: spec-decode-speculators-mtp-nightly-b200
  timeout_in_minutes: 30
-  device: b200-k8s
+  device: b200
  optional: true
  source_file_dependencies:
    - vllm/v1/spec_decode/
@@ -85,23 +84,23 @@ steps:
  commands:
    - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"

- label: Speculators Correctness
-  key: speculators-correctness
-  timeout_in_minutes: 60
+- label: DFlash Speculators Correctness
+  key: dflash-speculators-correctness
+  timeout_in_minutes: 30
  device: h100
  optional: true
  num_devices: 1
  source_file_dependencies:
    - vllm/v1/spec_decode/
    - vllm/model_executor/models/qwen3_dflash.py
-    - tests/v1/spec_decode/test_speculators_correctness.py
+    - tests/v1/spec_decode/test_speculators_dflash.py
  commands:
    - export VLLM_ALLOW_INSECURE_SERIALIZATION=1
-    - pytest -v -s v1/spec_decode/test_speculators_correctness.py -m slow_test
+    - pytest -v -s v1/spec_decode/test_speculators_dflash.py -m slow_test

 - label: Spec Decode MTP hybrid (B200)
  timeout_in_minutes: 30
-  device: b200-k8s
+  device: b200
  optional: true
  source_file_dependencies:
    - vllm/v1/spec_decode/
@@ -2,7 +2,6 @@
 /build
 dist
 vllm/*.so
-vllm/vllm-rs

 # Byte-compiled / optimized / DLL files
 __pycache__/
@@ -32,4 +31,3 @@ share/python-wheels/
 .installed.cfg
 *.egg
 MANIFEST
-rust/target/
@@ -6,8 +6,8 @@
 /vllm/distributed/kv_transfer @NickLucche @ApostaC @orozery @xuechendi
 /vllm/lora @jeejeelee
 /vllm/model_executor/layers/attention @LucasWilkinson @MatthewBonanni
-/vllm/model_executor/layers/fused_moe @mgoin @pavanimajety @zyongye
-/vllm/model_executor/layers/quantization @mgoin @robertgshaw2-redhat @tlrmchlsmth @yewentao256 @pavanimajety @zyongye
+/vllm/model_executor/layers/fused_moe @mgoin @pavanimajety
+/vllm/model_executor/layers/quantization @mgoin @robertgshaw2-redhat @tlrmchlsmth @yewentao256 @pavanimajety
 /vllm/model_executor/layers/mamba @tdoublep @tomeras91
 /vllm/model_executor/layers/mamba/gdn_linear_attn.py @tdoublep @ZJY0516 @vadiklyutiy
 /vllm/model_executor/layers/rotary_embedding.py @vadiklyutiy
@@ -18,8 +18,7 @@
 /vllm/kernels/helion @ProExpertProg @zou3519
 /vllm/multimodal @DarkLight1337 @ywang96 @NickLucche @tjtanaa
 /vllm/vllm_flash_attn @LucasWilkinson @MatthewBonanni
-/CMakeLists.txt @tlrmchlsmth @LucasWilkinson @Harry-Chen
-/cmake @tlrmchlsmth @LucasWilkinson @Harry-Chen
+CMakeLists.txt @tlrmchlsmth @LucasWilkinson

 # Any change to the VllmConfig changes can have a large user-facing impact,
 # so spam a lot of people
@@ -31,8 +30,8 @@
 /vllm/entrypoints/cli @hmellor @mgoin @DarkLight1337 @russellb
 /vllm/entrypoints/mcp @heheda12345
 /vllm/entrypoints/openai @aarnphm @chaunceyjiang @DarkLight1337 @russellb
-/vllm/entrypoints/speech_to_text/realtime @njhill
-/vllm/entrypoints/speech_to_text @NickLucche
+/vllm/entrypoints/openai/realtime @njhill
+/vllm/entrypoints/openai/speech_to_text @NickLucche
 /vllm/entrypoints/pooling @noooop
 /vllm/entrypoints/sagemaker @DarkLight1337
 /vllm/entrypoints/serve @njhill
@@ -40,12 +39,6 @@
 /vllm/entrypoints/chat_utils.py @DarkLight1337
 /vllm/entrypoints/llm.py @DarkLight1337

-# Rust Frontend
-/rust/ @BugenZhao @njhill
-/build_rust.sh @BugenZhao @njhill
-/rust-toolchain.toml @BugenZhao @njhill
-/.buildkite/test_areas/rust* @BugenZhao @njhill
-
 # Input/Output Processing
 /vllm/sampling_params.py @njhill @NickLucche
 /vllm/pooling_params.py @noooop @DarkLight1337
@@ -77,24 +70,18 @@
 /vllm/v1/worker/gpu @WoosukKwon @njhill
 /vllm/v1/worker/gpu/kv_connector.py @orozery

-# CI & building
-/.buildkite @Harry-Chen @khluu
-/docker/Dockerfile @Harry-Chen @khluu
-/pyproject.toml @khluu
-/setup.py @khluu
-
 # Test ownership
-/.buildkite/lm-eval-harness @mgoin
+/.buildkite/lm-eval-harness @mgoin 
 /tests/distributed/test_multi_node_assignment.py @youkaichao
 /tests/distributed/test_pipeline_parallel.py @youkaichao
 /tests/distributed/test_same_node.py @youkaichao
-/tests/entrypoints @DarkLight1337 @robertgshaw2-redhat @aarnphm @NickLucche @AndreasKaratzas
-/tests/evals @mgoin @vadiklyutiy @AndreasKaratzas
-/tests/kernels @mgoin @tlrmchlsmth @WoosukKwon @yewentao256 @zyongye @AndreasKaratzas
+/tests/entrypoints @DarkLight1337 @robertgshaw2-redhat @aarnphm @NickLucche
+/tests/evals @mgoin @vadiklyutiy
+/tests/kernels @mgoin @tlrmchlsmth @WoosukKwon @yewentao256
 /tests/kernels/ir @ProExpertProg @tjtanaa
-/tests/models @DarkLight1337 @ywang96 @AndreasKaratzas
+/tests/models @DarkLight1337 @ywang96
 /tests/multimodal @DarkLight1337 @ywang96 @NickLucche
-/tests/quantization @mgoin @robertgshaw2-redhat @yewentao256 @pavanimajety @zyongye @AndreasKaratzas
+/tests/quantization @mgoin @robertgshaw2-redhat @yewentao256 @pavanimajety
 /tests/test_inputs.py @DarkLight1337 @ywang96
 /tests/entrypoints/llm/test_struct_output_generate.py @mgoin @russellb @aarnphm
 /tests/v1/structured_output @mgoin @russellb @aarnphm
@@ -160,10 +147,6 @@ mkdocs.yaml @hmellor
 # MTP-specific files
 /vllm/model_executor/models/deepseek_mtp.py @luccafong

-# DeepseekV4-specific files
-/vllm/models/deepseek_v4 @zyongye
-/vllm/model_executor/layers/sparse_attn_indexer.py @zyongye
-
 # Mistral-specific files
 /vllm/model_executor/models/mistral*.py @patrickvonplaten
 /vllm/model_executor/models/mixtral*.py @patrickvonplaten
@@ -178,21 +161,21 @@ mkdocs.yaml @hmellor
 /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
-/docker/Dockerfile.rocm* @tjtanaa @dllehr-amd @AndreasKaratzas
-/vllm/v1/attention/backends/rocm*.py @tjtanaa @dllehr-amd
-/vllm/v1/attention/backends/mla/rocm*.py @tjtanaa @dllehr-amd
-/vllm/v1/attention/ops/rocm*.py @tjtanaa @dllehr-amd
-/vllm/model_executor/layers/fused_moe/rocm*.py @tjtanaa @dllehr-amd
-/csrc/rocm @tjtanaa @dllehr-amd
-/requirements/*rocm* @tjtanaa @AndreasKaratzas
-/tests/**/*rocm* @tjtanaa @AndreasKaratzas
+/vllm/**/*rocm* @tjtanaa
+/docker/Dockerfile.rocm* @gshtras @tjtanaa
+/vllm/v1/attention/backends/rocm*.py @gshtras @tjtanaa
+/vllm/v1/attention/backends/mla/rocm*.py @gshtras @tjtanaa
+/vllm/v1/attention/ops/rocm*.py @gshtras @tjtanaa
+/vllm/model_executor/layers/fused_moe/rocm*.py @gshtras @tjtanaa
+/csrc/rocm @gshtras @tjtanaa
+/requirements/*rocm* @tjtanaa
+/tests/**/*rocm* @tjtanaa
 /docs/**/*rocm* @tjtanaa
 /vllm/**/*quark* @tjtanaa
-/tests/**/*quark* @tjtanaa @AndreasKaratzas
+/tests/**/*quark* @tjtanaa
 /docs/**/*quark* @tjtanaa
-/vllm/**/*aiter* @tjtanaa @AndreasKaratzas
-/tests/**/*aiter* @tjtanaa @AndreasKaratzas
+/vllm/**/*aiter* @tjtanaa
+/tests/**/*aiter* @tjtanaa

 # TPU
 /vllm/v1/worker/tpu* @NickLucche
@@ -103,19 +103,6 @@ pull_request_rules:
      add:
        - frontend

- name: label-rust
-  description: Automatically apply rust label
-  conditions:
-    - label != stale
-    - or:
-      - files~=(?i)rust
-      - title~=(?i)rust
-      - title~=(?i)vllm-rs
-  actions:
-    label:
-      add:
-        - rust
-
 - name: label-llama
  description: Automatically apply llama label
  conditions:
@@ -490,7 +477,9 @@ pull_request_rules:
  conditions:
    - label != stale
    - or:
-      - files~=^examples/disaggregated/
+      - files~=^examples/online_serving/disaggregated[^/]*/.*
+      - files~=^examples/offline_inference/disaggregated[^/]*/.*
+      - files~=^examples/others/lmcache/
      - files~=^tests/v1/kv_connector/
      - files~=^vllm/distributed/kv_transfer/
      - title~=(?i)\bP/?D\b
@@ -10,7 +10,7 @@ jobs:
        runs-on: ubuntu-latest
        steps:
            -   name: Add label
-                uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
+                uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
                with:
                    script: |
                        github.rest.issues.addLabels({
@@ -14,7 +14,7 @@ jobs:
    steps:
      - name: Label issues based on keywords
        id: label-step
-        uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3  # v9.0.0
+        uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd  # v8.0.0
        with:
          script: |
            // Configuration: Add new labels and keywords here
@@ -315,7 +315,7 @@ jobs:

      - name: CC users for labeled issues
        if: steps.label-step.outputs.labels_added != '[]'
-        uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3  # v9.0.0
+        uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd  # v8.0.0
        with:
          script: |
            // Configuration: Map labels to GitHub users to CC
@@ -392,7 +392,7 @@ jobs:

      - name: Request missing ROCm info from issue author
        if: contains(steps.label-step.outputs.labels_added, 'rocm') && contains(toJSON(github.event.issue.labels.*.name), 'bug')
-        uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3  # v9.0.0
+        uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd  # v8.0.0
        with:
          script: |
            const body = (context.payload.issue.body || '').toLowerCase();
@@ -12,7 +12,7 @@ jobs:
    runs-on: ubuntu-latest
    steps:
      - name: Update PR description
-        uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
+        uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
        with:
          script: |
            const { owner, repo } = context.repo;
@@ -55,7 +55,7 @@ jobs:
    runs-on: ubuntu-latest
    steps:
      - name: Post welcome comment for first-time contributors
-        uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
+        uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
        with:
          script: |
            const { owner, repo } = context.repo;
@@ -20,7 +20,7 @@ jobs:
    runs-on: ubuntu-latest
    steps:
    - name: Check PR label and author merge count
-      uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
+      uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
      with:
        script: |
          const { data: pr } = await github.rest.pulls.get({
@@ -49,7 +49,7 @@ jobs:
    runs-on: ubuntu-latest
    steps:
    - uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1
-    - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
+    - uses: actions/setup-python@83679a892e2d95755f2dac6acb0bfd1e9ac5d548 # v6.1.0
      with:
        python-version: "3.12"
    - run: echo "::add-matcher::.github/workflows/matchers/actionlint.json"
@@ -9,9 +9,6 @@ PATH=${cuda_home}/bin:$PATH
 LD_LIBRARY_PATH=${cuda_home}/lib64:$LD_LIBRARY_PATH

 # Install requirements
-if [ "$(echo $2 | cut -d. -f1)" = "12" ]; then
-    sed -i 's/^nvidia-cutlass-dsl\[cu13\]>=/nvidia-cutlass-dsl>=/' requirements/cuda.txt
-fi
 $python_executable -m pip install -r requirements/build/cuda.txt -r requirements/cuda.txt

 # Limit the number of parallel jobs to avoid OOM
@@ -26,9 +26,6 @@ __pycache__/
 # C extensions
 *.so

-# Rust binaries
-vllm/vllm-rs
-
 # Distribution / packaging
 .Python
 build/
@@ -21,7 +21,7 @@ repos:
  rev: v21.1.2
  hooks:
  - id: clang-format
-    exclude: 'csrc/(moe/topk_softmax_kernels.cu|libtorch_stable/quantization/gguf/(ggml-common.h|dequantize.cuh|vecdotq.cuh|mmq.cuh|mmvq.cuh))|vllm/third_party/.*'
+    exclude: 'csrc/(moe/topk_softmax_kernels.cu|quantization/gguf/(ggml-common.h|dequantize.cuh|vecdotq.cuh|mmq.cuh|mmvq.cuh))|vllm/third_party/.*'
    types_or: [c++, cuda]
    args: [--style=file, --verbose]
 - repo: https://github.com/DavidAnson/markdownlint-cli2
@@ -131,16 +131,6 @@ repos:
        --python-version, "3.12",
      ]
      files: ^requirements/(common|xpu|test/xpu)\.(in|txt)$
-    - id: pip-compile
-      alias: pip-compile-docs
-      name: pip-compile-docs
-      args: [
-        requirements/docs.in,
-        -o, requirements/docs.txt,
-        --python-platform, x86_64-manylinux_2_28,
-        --python-version, "3.12",
-      ]
-      files: ^requirements/docs\.(in|txt)$
 - repo: local
  hooks:
  - id: format-torch-nightly-test
@@ -222,12 +212,6 @@ repos:
    name: Update Dockerfile dependency graph
    entry: tools/pre_commit/update-dockerfile-graph.sh
    language: script
-  - id: test-nonroot-entrypoint
-    name: Test non-root entrypoint wrapper
-    entry: bash docker/entrypoints/test_vllm_nonroot_entrypoint.sh
-    language: system
-    pass_filenames: false
-    files: ^docker/entrypoints/(vllm-nonroot-entrypoint|test_vllm_nonroot_entrypoint)\.sh$
  - id: check-forbidden-imports
    name: Check for forbidden imports
    entry: python tools/pre_commit/check_forbidden_imports.py
@@ -262,32 +246,6 @@ repos:
    entry: python tools/pre_commit/check_boolean_context_manager.py
    language: python
    types: [python]
-  # Rust hooks. These shell out to `cargo`; tools/pre_commit/rust-check.sh
-  # skips with a warning when cargo is not installed.
-  - id: rust-cargo-autoinherit
-    name: Rust - Normalize Cargo manifests with autoinherit
-    entry: tools/pre_commit/rust-check.sh autoinherit --prefer-simple-dotted
-    language: script
-    pass_filenames: false
-    require_serial: true
-    stages: [pre-commit] # Only run locally as Buildkite will cover this
-    files: ^rust/(Cargo\.toml|src/.*/Cargo\.toml)$
-  - id: rust-cargo-sort
-    name: Rust - Sort Cargo manifest sections
-    entry: tools/pre_commit/rust-check.sh sort --workspace
-    language: script
-    pass_filenames: false
-    require_serial: true
-    stages: [pre-commit] # Only run locally as Buildkite will cover this
-    files: ^rust/(Cargo\.toml|src/.*/Cargo\.toml)$
-  - id: rust-cargo-fmt
-    name: Rust - Format code
-    entry: tools/pre_commit/rust-check.sh fmt
-    language: script
-    pass_filenames: false
-    require_serial: true
-    stages: [pre-commit] # Only run locally as Buildkite will cover this
-    files: ^rust/.*(\.rs|Cargo\.toml|rustfmt\.toml)$
  # Keep `suggestion` last
  - id: suggestion
    name: Suggestion
@@ -9,8 +9,8 @@ build:
    python: "3.12"
  jobs:
    post_checkout:
+      # - bash docs/maybe_skip_pr_build.sh
      - git fetch origin main --unshallow --no-tags --filter=blob:none || true
-      - bash docs/pre_run_check.sh
    pre_create_environment:
      - pip install uv
    create_environment:
@@ -101,8 +101,6 @@ pre-commit run ruff-check --all-files
 pre-commit run mypy-3.10 --all-files --hook-stage manual
 ```

-The line length limit for Python code is 88 characters. If you are not sure, use pre-commit to check.
-
 ### Commit messages

 Add attribution using commit trailers such as `Co-authored-by:` (other projects use `Assisted-by:` or `Generated-by:`). For example:
@@ -13,12 +13,8 @@ cmake_minimum_required(VERSION 3.26)
 # cmake --install . --component _C
 project(vllm_extensions LANGUAGES CXX)

-set(CMAKE_CXX_STANDARD 20)
+set(CMAKE_CXX_STANDARD 17)
 set(CMAKE_CXX_STANDARD_REQUIRED ON)
-set(CMAKE_CUDA_STANDARD 20)
-set(CMAKE_CUDA_STANDARD_REQUIRED ON)
-set(CMAKE_HIP_STANDARD 20)
-set(CMAKE_HIP_STANDARD_REQUIRED ON)


 # CUDA by default, can be overridden by using -DVLLM_TARGET_DEVICE=... (used by setup.py)
@@ -109,24 +105,6 @@ else()
  set(CUDA_SUPPORTED_ARCHS "7.0;7.5;8.0;8.6;8.7;8.9;9.0")
 endif()

-#
-# spinloop extension (pure CXX; must stay above the non-CUDA device branch so
-# CPU builds define the target before the early return)
-#
-set(VLLM_SPINLOOP_EXT_SRC "csrc/spinloop.cpp")
-set(SPINLOOP_COMPILE_FLAGS "")
-if(CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|amd64")
-  list(APPEND SPINLOOP_COMPILE_FLAGS "-mmwaitx")
-endif()
-define_extension_target(
-  spinloop
-  DESTINATION vllm
-  LANGUAGE CXX
-  SOURCES ${VLLM_SPINLOOP_EXT_SRC}
-  COMPILE_FLAGS ${SPINLOOP_COMPILE_FLAGS}
-  USE_SABI 3.11
-  WITH_SOABI)
-
 #
 # Forward the non-CUDA device extensions to external CMake scripts.
 #
@@ -144,14 +122,14 @@ endif()
 # Set up GPU language and check the torch version and warn if it isn't
 # what is expected.
 #
-if (NOT HIP_FOUND AND NOT PYTORCH_FOUND_HIP AND CUDA_FOUND)
+if (NOT HIP_FOUND AND CUDA_FOUND)
  set(VLLM_GPU_LANG "CUDA")

  if (NOT Torch_VERSION VERSION_EQUAL ${TORCH_SUPPORTED_VERSION_CUDA})
    message(WARNING "Pytorch version ${TORCH_SUPPORTED_VERSION_CUDA} "
      "expected for CUDA build, saw ${Torch_VERSION} instead.")
  endif()
-elseif(HIP_FOUND OR PYTORCH_FOUND_HIP)
+elseif(HIP_FOUND)
  set(VLLM_GPU_LANG "HIP")

  # Importing torch recognizes and sets up some HIP/ROCm configuration but does
@@ -305,8 +283,27 @@ endif()
 #

 set(VLLM_EXT_SRC
+  "csrc/mamba/mamba_ssm/selective_scan_fwd.cu"
+  "csrc/cache_kernels.cu"
+  "csrc/cache_kernels_fused.cu"
+  "csrc/attention/paged_attention_v1.cu"
+  "csrc/attention/paged_attention_v2.cu"
+  "csrc/attention/merge_attn_states.cu"
+  "csrc/attention/vertical_slash_index.cu"
+  "csrc/pos_encoding_kernels.cu"
+  "csrc/activation_kernels.cu"
+  "csrc/layernorm_kernels.cu"
+  "csrc/fused_qknorm_rope_kernel.cu"
+  "csrc/layernorm_quant_kernels.cu"
+  "csrc/sampler.cu"
+  "csrc/topk.cu"
  "csrc/cuda_view.cu"
+  "csrc/quantization/gptq/q_gemm.cu"
+  "csrc/quantization/w8a8/int8/scaled_quant.cu"
+  "csrc/quantization/w8a8/fp8/common.cu"
+  "csrc/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu"
  "csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu"
+  "csrc/quantization/gguf/gguf_kernel.cu"
  "csrc/quantization/activation_kernels.cu"
  "csrc/cuda_utils_kernels.cu"
  "csrc/custom_all_reduce.cu"
@@ -350,6 +347,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
  FetchContent_MakeAvailable(cutlass)

  list(APPEND VLLM_EXT_SRC
+    "csrc/quantization/awq/gemm_kernels.cu"
    "csrc/cutlass_extensions/common.cpp")

  set_gencode_flags_for_srcs(
@@ -361,30 +359,16 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
  # are not supported by Machete yet.

  # marlin arches for fp16 output
-  # Family-conditional 12.0f (one cubin for SM12x family) requires CUDA >= 13.0;
-  # fall back to architecture-specific 12.0a;12.1a on CUDA < 13.0 (e.g. 12.8).
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(MARLIN_ARCHS "8.0+PTX;12.0f" "${CUDA_ARCHS}")
-  else()
-    cuda_archs_loose_intersection(MARLIN_ARCHS "8.0+PTX;12.0a;12.1a" "${CUDA_ARCHS}")
-  endif()
+  cuda_archs_loose_intersection(MARLIN_ARCHS "8.0+PTX" "${CUDA_ARCHS}")
  # marlin has limited support for turing
  cuda_archs_loose_intersection(MARLIN_SM75_ARCHS "7.5" "${CUDA_ARCHS}")
  # marlin arches for bf16 output (we need 9.0 for bf16 atomicAdd PTX)
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(MARLIN_BF16_ARCHS "8.0+PTX;9.0+PTX;12.0f" "${CUDA_ARCHS}")
-  else()
-    cuda_archs_loose_intersection(MARLIN_BF16_ARCHS "8.0+PTX;9.0+PTX;12.0a;12.1a" "${CUDA_ARCHS}")
-  endif()
+  cuda_archs_loose_intersection(MARLIN_BF16_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}")
  # marlin arches for fp8 input
  # - sm80 doesn't support fp8 computation
  # - sm90 and sm100 don't support QMMA.16832.F32.E4M3.E4M3 SAAS instruction
  # so we only enable fp8 computation for SM89 (e.g. RTX 40x0)  and 12.0 (e.g. RTX 50x0)
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(MARLIN_FP8_ARCHS "8.9;12.0f" "${CUDA_ARCHS}")
-  else()
-    cuda_archs_loose_intersection(MARLIN_FP8_ARCHS "8.9;12.0a;12.1a" "${CUDA_ARCHS}")
-  endif()
+  cuda_archs_loose_intersection(MARLIN_FP8_ARCHS "8.9;12.0;12.1" "${CUDA_ARCHS}")
  # marlin arches for other files
  cuda_archs_loose_intersection(MARLIN_OTHER_ARCHS "7.5;8.0+PTX" "${CUDA_ARCHS}")

@@ -495,6 +479,46 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
                   " in CUDA target architectures")
  endif()

+  # Only build AllSpark kernels if we are building for at least some compatible archs.
+  cuda_archs_loose_intersection(ALLSPARK_ARCHS "8.0;8.6;8.7;8.9" "${CUDA_ARCHS}")
+  if (ALLSPARK_ARCHS)
+    set(ALLSPARK_SRCS
+       "csrc/quantization/gptq_allspark/allspark_repack.cu"
+       "csrc/quantization/gptq_allspark/allspark_qgemm_w8a16.cu")
+    set_gencode_flags_for_srcs(
+      SRCS "${ALLSPARK_SRCS}"
+      CUDA_ARCHS "${ALLSPARK_ARCHS}")
+    list(APPEND VLLM_EXT_SRC "${ALLSPARK_SRCS}")
+    message(STATUS "Building AllSpark kernels for archs: ${ALLSPARK_ARCHS}")
+  else()
+    message(STATUS "Not building AllSpark kernels as no compatible archs found"
+                   " in CUDA target architectures")
+  endif()
+
+  # CUTLASS MLA Archs and flags
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
+    cuda_archs_loose_intersection(MLA_ARCHS "10.0f;11.0f;12.0f" "${CUDA_ARCHS}")
+  else()
+    cuda_archs_loose_intersection(MLA_ARCHS "10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
+  endif()
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND MLA_ARCHS)
+    set(SRCS
+      "csrc/attention/mla/sm100_cutlass_mla_kernel.cu")
+    set_gencode_flags_for_srcs(
+      SRCS "${SRCS}"
+      CUDA_ARCHS "${MLA_ARCHS}")
+    list(APPEND VLLM_EXT_SRC "${SRCS}")
+    list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MLA=1")
+    # Add MLA-specific include directories only to MLA source files
+    set_source_files_properties(${SRCS}
+      PROPERTIES INCLUDE_DIRECTORIES "${CUTLASS_DIR}/examples/77_blackwell_fmha;${CUTLASS_DIR}/examples/common")
+    message(STATUS "Building CUTLASS MLA for archs: ${MLA_ARCHS}")
+  else()
+    message(STATUS "Not building CUTLASS MLA as no compatible archs were found.")
+    # clear MLA_ARCHS
+    set(MLA_ARCHS)
+  endif()
+
  # Expert-specialization MXFP8 blockscaled grouped kernels (SM100+).
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
    cuda_archs_loose_intersection(ES_MXFP8_GROUPED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
@@ -522,6 +546,24 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
    endif()
  endif()

+  # 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;11.0f" "${CUDA_ARCHS}")
+  else()
+    cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
+  endif()
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND DSV3_FUSED_A_GEMM_ARCHS)
+    set(DSV3_FUSED_A_GEMM_SRC "csrc/dsv3_fused_a_gemm.cu")
+    set_gencode_flags_for_srcs(
+      SRCS "${DSV3_FUSED_A_GEMM_SRC}"
+      CUDA_ARCHS "${DSV3_FUSED_A_GEMM_ARCHS}")
+    list(APPEND VLLM_EXT_SRC ${DSV3_FUSED_A_GEMM_SRC})
+    message(STATUS "Building dsv3_fused_a_gemm for archs: ${DSV3_FUSED_A_GEMM_ARCHS}")
+  else()
+    message(STATUS "Not building dsv3_fused_a_gemm as no compatible archs found "
+                   "in CUDA target architectures.")
+  endif()
+
  #
  # Machete kernels

@@ -593,6 +635,16 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
  endif()


+  # Hadacore kernels
+  cuda_archs_loose_intersection(HADACORE_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}")
+  if(HADACORE_ARCHS)
+    set(SRCS "csrc/quantization/hadamard/hadacore/hadamard_transform_cuda.cu")
+    set_gencode_flags_for_srcs(
+      SRCS "${SRCS}"
+      CUDA_ARCHS "${HADACORE_ARCHS}")
+    list(APPEND VLLM_EXT_SRC "${SRCS}")
+    message(STATUS "Building hadacore")
+  endif()

 # if CUDA endif
 endif()
@@ -624,95 +676,31 @@ define_extension_target(
 # Setting this variable sidesteps the issue by calling the driver directly.
 target_compile_definitions(_C PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)

-if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
+# add OR VLLM_GPU_LANG STREQUAL "HIP" here once
+# https://github.com/vllm-project/vllm/issues/35163 is resolved
+if(VLLM_GPU_LANG STREQUAL "CUDA")
  #
  # _C_stable_libtorch extension (ops registered via STABLE_TORCH_LIBRARY)
  #
  set(VLLM_STABLE_EXT_SRC
    "csrc/libtorch_stable/torch_bindings.cpp"
-    "csrc/libtorch_stable/activation_kernels.cu"
-    "csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu"
-    "csrc/libtorch_stable/quantization/w8a8/fp8/common.cu"
-    "csrc/libtorch_stable/quantization/gptq/q_gemm.cu"
-    "csrc/libtorch_stable/quantization/gguf/gguf_kernel.cu"
-    "csrc/libtorch_stable/pos_encoding_kernels.cu"
-    "csrc/libtorch_stable/fused_qknorm_rope_kernel.cu"
-    "csrc/libtorch_stable/layernorm_kernels.cu"
-    "csrc/libtorch_stable/layernorm_quant_kernels.cu"
-    "csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu"
-    "csrc/libtorch_stable/attention/merge_attn_states.cu"
-    "csrc/libtorch_stable/sampler.cu"
-    "csrc/libtorch_stable/topk.cu"
-    "csrc/libtorch_stable/mamba/selective_scan_fwd.cu"
-    "csrc/libtorch_stable/attention/paged_attention_v1.cu"
-    "csrc/libtorch_stable/attention/paged_attention_v2.cu"
-    "csrc/libtorch_stable/cache_kernels.cu"
-    "csrc/libtorch_stable/cache_kernels_fused.cu")
+    "csrc/cutlass_extensions/common.cpp"
+    "csrc/cuda_utils_kernels.cu"
+    "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu"
+    "csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu"
+    "csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu")

  if(VLLM_GPU_LANG STREQUAL "CUDA")
    list(APPEND VLLM_STABLE_EXT_SRC
-      "csrc/cuda_utils_kernels.cu"
-      "csrc/cutlass_extensions/common.cpp"
-      "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu"
-      "csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu"
-      "csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu"
      "csrc/libtorch_stable/permute_cols.cu"
      "csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu"
-      "csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu"
-      "csrc/libtorch_stable/quantization/awq/gemm_kernels.cu")
+      "csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu")
+  endif()

+  if(VLLM_GPU_LANG STREQUAL "CUDA")
    set_gencode_flags_for_srcs(
      SRCS "${VLLM_STABLE_EXT_SRC}"
      CUDA_ARCHS "${CUDA_ARCHS}")
-
-  # 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;11.0f" "${CUDA_ARCHS}")
-  else()
-    cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
-  endif()
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND DSV3_FUSED_A_GEMM_ARCHS)
-    set(SRCS "csrc/libtorch_stable/dsv3_fused_a_gemm.cu")
-    set_gencode_flags_for_srcs(
-      SRCS "${SRCS}"
-      CUDA_ARCHS "${DSV3_FUSED_A_GEMM_ARCHS}")
-    list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
-    message(STATUS "Building dsv3_fused_a_gemm for archs: ${DSV3_FUSED_A_GEMM_ARCHS}")
-  else()
-    message(STATUS "Not building dsv3_fused_a_gemm as no compatible archs found "
-                   "in CUDA target architectures.")
-  endif()
-
-  # FP32 router GEMM (H=3072, E=256, M<=32). Requires SM90+ and CUDA >= 12.0.
-  cuda_archs_sm90plus(FP32_ROUTER_GEMM_ARCHS "${CUDA_ARCHS}")
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND FP32_ROUTER_GEMM_ARCHS)
-    set(SRCS
-      "csrc/libtorch_stable/fp32_router_gemm_entry.cu"
-      "csrc/libtorch_stable/fp32_router_gemm.cu")
-    set_gencode_flags_for_srcs(
-      SRCS "${SRCS}"
-      CUDA_ARCHS "${FP32_ROUTER_GEMM_ARCHS}")
-    list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
-    message(STATUS "Building fp32_router_gemm for archs: ${FP32_ROUTER_GEMM_ARCHS}")
-  else()
-    message(STATUS "Not building fp32_router_gemm as no compatible archs found "
-                   "(requires SM90+ and CUDA >= 12.0).")
-  endif()
-
-  # Only build AllSpark kernels if we are building for at least some compatible archs.
-  cuda_archs_loose_intersection(ALLSPARK_ARCHS "8.0;8.6;8.7;8.9" "${CUDA_ARCHS}")
-  if (ALLSPARK_ARCHS)
-    set(SRCS
-       "csrc/libtorch_stable/quantization/gptq_allspark/allspark_repack.cu"
-       "csrc/libtorch_stable/quantization/gptq_allspark/allspark_qgemm_w8a16.cu")
-    set_gencode_flags_for_srcs(
-      SRCS "${SRCS}"
-      CUDA_ARCHS "${ALLSPARK_ARCHS}")
-    list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
-    message(STATUS "Building AllSpark kernels for archs: ${ALLSPARK_ARCHS}")
-  else()
-    message(STATUS "Not building AllSpark kernels as no compatible archs found"
-                   " in CUDA target architectures")
  endif()

  #
@@ -940,11 +928,13 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
      SRCS "${SRCS}"
      CUDA_ARCHS "${FP4_ARCHS}")
    list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
-    set(NVFP4_KV_SRC "csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu")
+    # nvfp4_kv_cache_kernels uses non-stable torch API and is called directly
+    # from cache_kernels.cu, so it belongs in _C rather than _C_stable.
+    set(NVFP4_KV_SRC "csrc/nvfp4_kv_cache_kernels.cu")
    set_gencode_flags_for_srcs(
      SRCS "${NVFP4_KV_SRC}"
      CUDA_ARCHS "${FP4_ARCHS}")
-    list(APPEND VLLM_STABLE_EXT_SRC "${NVFP4_KV_SRC}")
+    target_sources(_C PRIVATE ${NVFP4_KV_SRC})
    target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM120=1)
    list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM120=1")
    list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM120=1")
@@ -974,11 +964,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
      SRCS "${SRCS}"
      CUDA_ARCHS "${FP4_ARCHS}")
    list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
-    set(NVFP4_KV_SRC "csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu")
+    set(NVFP4_KV_SRC "csrc/nvfp4_kv_cache_kernels.cu")
    set_gencode_flags_for_srcs(
      SRCS "${NVFP4_KV_SRC}"
      CUDA_ARCHS "${FP4_ARCHS}")
-    list(APPEND VLLM_STABLE_EXT_SRC "${NVFP4_KV_SRC}")
+    target_sources(_C PRIVATE ${NVFP4_KV_SRC})
    target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM100=1)
    list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM100=1")
    list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
@@ -1022,44 +1012,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
    endif()
  endif()

-  # CUTLASS MLA Archs and flags
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(MLA_ARCHS "10.0f;11.0f;12.0f" "${CUDA_ARCHS}")
-  else()
-    cuda_archs_loose_intersection(MLA_ARCHS "10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
-  endif()
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND MLA_ARCHS)
-    set(SRCS
-      "csrc/libtorch_stable/attention/mla/sm100_cutlass_mla_kernel.cu")
-    set_gencode_flags_for_srcs(
-      SRCS "${SRCS}"
-      CUDA_ARCHS "${MLA_ARCHS}")
-    list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
-    list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MLA=1")
-    # Add MLA-specific include directories only to MLA source files
-    set_source_files_properties(${SRCS}
-      PROPERTIES INCLUDE_DIRECTORIES "${CUTLASS_DIR}/examples/77_blackwell_fmha;${CUTLASS_DIR}/examples/common")
-    message(STATUS "Building CUTLASS MLA for archs: ${MLA_ARCHS}")
-  else()
-    message(STATUS "Not building CUTLASS MLA as no compatible archs were found.")
-    # clear MLA_ARCHS
-    set(MLA_ARCHS)
-  endif()
-
-  # Hadacore kernels
-  cuda_archs_loose_intersection(HADACORE_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}")
-  if(HADACORE_ARCHS)
-    set(SRCS "csrc/libtorch_stable/quantization/hadamard/hadacore/hadamard_transform_cuda.cu")
-    set_gencode_flags_for_srcs(
-      SRCS "${SRCS}"
-      CUDA_ARCHS "${HADACORE_ARCHS}")
-    list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
-    message(STATUS "Building hadacore")
-  endif()
-
-  # if CUDA endif
-  endif()
-
  message(STATUS "Enabling C_stable extension.")
  define_extension_target(
    _C_stable_libtorch
@@ -1079,34 +1031,13 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
  target_compile_definitions(_C_stable_libtorch PRIVATE
    TORCH_TARGET_VERSION=0x020A000000000000ULL)

-  # Needed to use cuda/hip APIs from C-shim
-  if(VLLM_GPU_LANG STREQUAL "CUDA")
-    target_compile_definitions(_C_stable_libtorch PRIVATE USE_CUDA)
-    # Needed by CUTLASS kernels
-    target_compile_definitions(_C_stable_libtorch PRIVATE
-      CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
-  elseif(VLLM_GPU_LANG STREQUAL "HIP")
-    target_compile_definitions(_C_stable_libtorch PRIVATE USE_ROCM)
-  endif()
+  # Needed to use cuda APIs from C-shim
+  target_compile_definitions(_C_stable_libtorch PRIVATE
+    USE_CUDA)

-  # On ROCm, _C_stable_libtorch calls raw HIP APIs (e.g. hipGetDevice in
-  # get_device_prop()) which must resolve to the same libamdhip64.so that
-  # PyTorch uses.  When PyTorch bundles its own copy (pip/conda wheels),
-  # the raw HIP calls would otherwise resolve to the system ROCm copy,
-  # initializing a second HIP runtime that corrupts device state (wrong
-  # device on DeviceGuard, core dumps on multi-GPU tests).
-  #
-  # If PyTorch doesn't bundle libamdhip64 (built from source against system
-  # ROCm), there is only one copy in the process and no action is needed —
-  # the HIP compiler already links the system libamdhip64 automatically.
-  if(VLLM_GPU_LANG STREQUAL "HIP")
-    find_library(_STABLE_TORCH_AMDHIP64 amdhip64
-      PATHS "${TORCH_INSTALL_PREFIX}/lib" NO_DEFAULT_PATH)
-    if(_STABLE_TORCH_AMDHIP64)
-      message(STATUS "Found PyTorch-bundled libamdhip64 at ${_STABLE_TORCH_AMDHIP64}")
-      target_link_libraries(_C_stable_libtorch PRIVATE ${_STABLE_TORCH_AMDHIP64})
-    endif()
-  endif()
+  # Needed by CUTLASS kernels
+  target_compile_definitions(_C_stable_libtorch PRIVATE
+    CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
 endif()

 #
@@ -1149,11 +1080,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
  # moe marlin arches
  # note that we always set `use_atomic_add=False` for moe marlin now,
  # so we don't need 9.0 for bf16 atomicAdd PTX
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(MARLIN_MOE_ARCHS "8.0+PTX;12.0f" "${CUDA_ARCHS}")
-  else()
-    cuda_archs_loose_intersection(MARLIN_MOE_ARCHS "8.0+PTX;12.0a;12.1a" "${CUDA_ARCHS}")
-  endif()
+  cuda_archs_loose_intersection(MARLIN_MOE_ARCHS "8.0+PTX" "${CUDA_ARCHS}")
  # moe marlin has limited support for turing
  cuda_archs_loose_intersection(MARLIN_MOE_SM75_ARCHS "7.5" "${CUDA_ARCHS}")
  # moe marlin arches for fp8 input
@@ -1256,22 +1183,24 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
                   " in CUDA target architectures")
  endif()

-  # DeepSeek V3 router GEMM kernel requires SM90+ and CUDA >= 12.0.
-  # (fp32_router_gemm has been migrated to _C_stable_libtorch above.)
-  cuda_archs_sm90plus(SM90PLUS_ROUTER_GEMM_ARCHS "${CUDA_ARCHS}")
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND SM90PLUS_ROUTER_GEMM_ARCHS)
+  # DeepSeek V3 router GEMM kernel - requires SM90+
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
+    cuda_archs_loose_intersection(DSV3_ROUTER_GEMM_ARCHS "9.0a;10.0f;11.0f" "${CUDA_ARCHS}")
+  else()
+    cuda_archs_loose_intersection(DSV3_ROUTER_GEMM_ARCHS "9.0a;10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
+  endif()
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND DSV3_ROUTER_GEMM_ARCHS)
    set(DSV3_ROUTER_GEMM_SRC
      "csrc/moe/dsv3_router_gemm_entry.cu"
      "csrc/moe/dsv3_router_gemm_float_out.cu"
      "csrc/moe/dsv3_router_gemm_bf16_out.cu")
    set_gencode_flags_for_srcs(
      SRCS "${DSV3_ROUTER_GEMM_SRC}"
-      CUDA_ARCHS "${SM90PLUS_ROUTER_GEMM_ARCHS}")
+      CUDA_ARCHS "${DSV3_ROUTER_GEMM_ARCHS}")
    list(APPEND VLLM_MOE_EXT_SRC "${DSV3_ROUTER_GEMM_SRC}")
-
-    message(STATUS "Building DSV3 router GEMM kernels for archs: ${SM90PLUS_ROUTER_GEMM_ARCHS}")
+    message(STATUS "Building DSV3 router GEMM kernel for archs: ${DSV3_ROUTER_GEMM_ARCHS}")
  else()
-    message(STATUS "Not building DSV3 router GEMM kernels as no compatible archs found"
+    message(STATUS "Not building DSV3 router GEMM kernel as no compatible archs found"
                   " (requires SM90+ and CUDA >= 12.0)")
  endif()
 endif()
@@ -1298,14 +1227,6 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
    "csrc/rocm/skinny_gemms.cu"
    "csrc/rocm/attention.cu")

-  set(VLLM_ROCM_HAS_GFX1100 OFF)
-  if(VLLM_GPU_ARCHES MATCHES "gfx1100")
-    set(VLLM_ROCM_HAS_GFX1100 ON)
-    list(APPEND VLLM_ROCM_EXT_SRC
-      "csrc/rocm/q_gemm_rdna3.cu"
-      "csrc/rocm/q_gemm_rdna3_wmma.cu")
-  endif()
-
  define_extension_target(
    _rocm_C
    DESTINATION vllm
@@ -1315,16 +1236,6 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
    ARCHITECTURES ${VLLM_GPU_ARCHES}
    USE_SABI 3
    WITH_SOABI)
-
-  if(VLLM_ROCM_HAS_GFX1100)
-    target_compile_definitions(_rocm_C PRIVATE VLLM_ROCM_GFX1100)
-  endif()
-endif()
-
-# Must run after the last HIP `define_extension_target` so every extension
-# has registered its sources.
-if (VLLM_GPU_LANG STREQUAL "HIP")
-  vllm_finalize_hipify_target()
 endif()

 # For CUDA and HIP builds also build the triton_kernels external package.
@@ -53,7 +53,6 @@ backends:
  - FLASHINFER_MLA
  - FLASH_ATTN_MLA  # Hopper only
  - FLASHMLA        # Hopper only
-  - TOKENSPEED_MLA  # Blackwell + R1 dims + FP8 KV (use --kv-cache-dtype fp8)

 device: "cuda:0"
 repeats: 100
@@ -3,7 +3,6 @@
 # Compares all available MLA prefill backends:
 #   FA backends:  fa2, fa3, fa4 (FlashAttention versions)
 #   Non-FA:       flashinfer, cudnn, trtllm (Blackwell-only, require flashinfer)
-#   CuTe DSL:     tokenspeed (Blackwell + R1 dims, requires tokenspeed_mla)
 #
 # Uses cutlass_mla as the decode backend for impl construction
 # (only the prefill path is exercised).
@@ -121,7 +120,6 @@ prefill_backends:
  - flashinfer
  - cudnn
  - trtllm
-  - tokenspeed

 device: "cuda:0"
 repeats: 20
@@ -29,7 +29,6 @@ from vllm.config import (
    VllmConfig,
    set_current_vllm_config,
 )
-from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum

 # ============================================================================
 # VllmConfig Creation
@@ -80,8 +79,8 @@ def create_minimal_vllm_config(
        index_topk: Optional topk value for sparse MLA backends. If provided,
                    the config will include index_topk for sparse attention.
        prefill_backend: Prefill backend name (e.g., "fa3", "fa4", "flashinfer",
-                        "trtllm"). Configures the attention config to force
-                        the specified prefill backend.
+                        "cudnn", "trtllm"). Configures the attention config to
+                        force the specified prefill backend.

    Returns:
        VllmConfig for benchmarking
@@ -180,13 +179,19 @@ def create_minimal_vllm_config(

    if prefill_backend is not None:
        prefill_cfg = get_prefill_backend_config(prefill_backend)
-        vllm_config.attention_config.mla_prefill_backend = prefill_cfg[
-            "mla_prefill_backend"
-        ]
        if prefill_cfg["flash_attn_version"] is not None:
            vllm_config.attention_config.flash_attn_version = prefill_cfg[
                "flash_attn_version"
            ]
+        vllm_config.attention_config.disable_flashinfer_prefill = prefill_cfg[
+            "disable_flashinfer_prefill"
+        ]
+        vllm_config.attention_config.use_cudnn_prefill = prefill_cfg[
+            "use_cudnn_prefill"
+        ]
+        vllm_config.attention_config.use_trtllm_ragged_deepseek_prefill = prefill_cfg[
+            "use_trtllm_ragged_deepseek_prefill"
+        ]

    return vllm_config

@@ -201,27 +206,39 @@ def create_minimal_vllm_config(
 _PREFILL_BACKEND_CONFIG: dict[str, dict] = {
    "fa2": {
        "flash_attn_version": 2,
-        "mla_prefill_backend": MLAPrefillBackendEnum.FLASH_ATTN,
+        "disable_flashinfer_prefill": True,
+        "use_cudnn_prefill": False,
+        "use_trtllm_ragged_deepseek_prefill": False,
    },
    "fa3": {
        "flash_attn_version": 3,
-        "mla_prefill_backend": MLAPrefillBackendEnum.FLASH_ATTN,
+        "disable_flashinfer_prefill": True,
+        "use_cudnn_prefill": False,
+        "use_trtllm_ragged_deepseek_prefill": False,
    },
    "fa4": {
        "flash_attn_version": 4,
-        "mla_prefill_backend": MLAPrefillBackendEnum.FLASH_ATTN,
+        "disable_flashinfer_prefill": True,
+        "use_cudnn_prefill": False,
+        "use_trtllm_ragged_deepseek_prefill": False,
    },
    "flashinfer": {
        "flash_attn_version": None,
-        "mla_prefill_backend": MLAPrefillBackendEnum.FLASHINFER,
+        "disable_flashinfer_prefill": False,
+        "use_cudnn_prefill": False,
+        "use_trtllm_ragged_deepseek_prefill": False,
+    },
+    "cudnn": {
+        "flash_attn_version": None,
+        "disable_flashinfer_prefill": True,
+        "use_cudnn_prefill": True,
+        "use_trtllm_ragged_deepseek_prefill": False,
    },
    "trtllm": {
        "flash_attn_version": None,
-        "mla_prefill_backend": MLAPrefillBackendEnum.TRTLLM_RAGGED,
-    },
-    "tokenspeed": {
-        "flash_attn_version": None,
-        "mla_prefill_backend": MLAPrefillBackendEnum.TOKENSPEED_MLA,
+        "disable_flashinfer_prefill": True,
+        "use_cudnn_prefill": False,
+        "use_trtllm_ragged_deepseek_prefill": True,
    },
 }

@@ -608,21 +625,6 @@ def _create_backend_impl(
    # Create mock layer
    layer = MockLayer(device, impl=impl, kv_cache_spec=kv_cache_spec)

-    # Attach a prefill backend (MLAAttention does this in __init__; the metadata
-    # builder reads layer.prefill_backend from static_forward_context).
-    from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend
-
-    prefill_backend_cls = get_mla_prefill_backend(vllm_config)
-    layer.prefill_backend = prefill_backend_cls(
-        num_heads=mla_dims["num_q_heads"],
-        scale=(mla_dims["qk_nope_head_dim"] + mla_dims["qk_rope_head_dim"]) ** -0.5,
-        kv_lora_rank=mla_dims["kv_lora_rank"],
-        qk_nope_head_dim=mla_dims["qk_nope_head_dim"],
-        qk_rope_head_dim=mla_dims["qk_rope_head_dim"],
-        v_head_dim=mla_dims["v_head_dim"],
-        vllm_config=vllm_config,
-    )
-
    # Create builder instance if needed
    builder_instance = None
    if builder_class:
@@ -959,6 +961,19 @@ def _run_mla_benchmark_batched(
    results = []

    with set_current_vllm_config(vllm_config):
+        # Clear cached prefill backend detection functions so they re-evaluate
+        # with the current VllmConfig. These are @functools.cache decorated and
+        # would otherwise return stale results from a previous backend's config.
+        from vllm.model_executor.layers.attention.mla_attention import (
+            use_cudnn_prefill,
+            use_flashinfer_prefill,
+            use_trtllm_ragged_deepseek_prefill,
+        )
+
+        use_flashinfer_prefill.cache_clear()
+        use_cudnn_prefill.cache_clear()
+        use_trtllm_ragged_deepseek_prefill.cache_clear()
+
        # Create backend impl, layer, builder, and indexer (reused across benchmarks)
        impl, layer, builder_instance, indexer = _create_backend_impl(
            backend_cfg,
@@ -970,36 +985,38 @@ def _run_mla_benchmark_batched(
            kv_cache_dtype=kv_cache_dtype,
        )

-        # Verify the actual prefill backend matches what was requested. The
-        # selector + impl construction already raise on misuse; here we just
-        # check the resolved class against the requested name as a sanity guard.
+        # Verify the actual prefill backend matches what was requested
        if prefill_backend is not None:
-            expected_class = {
-                "fa2": "FlashAttnPrefillBackend",
-                "fa3": "FlashAttnPrefillBackend",
-                "fa4": "FlashAttnPrefillBackend",
-                "flashinfer": "FlashInferPrefillBackend",
-                "trtllm": "TrtllmRaggedPrefillBackend",
-                "tokenspeed": "TokenspeedMLAPrefillBackend",
-            }.get(prefill_backend)
-            actual_class = type(getattr(layer, "prefill_backend", None)).__name__
-            if expected_class and actual_class != expected_class:
-                raise RuntimeError(
-                    f"Prefill backend '{prefill_backend}' requested "
-                    f"{expected_class}, got {actual_class}. Check "
-                    f"attention_config plumbing or installed deps."
-                )
-            if prefill_backend in {"fa2", "fa3", "fa4"}:
-                fa_version = int(prefill_backend[2:])
-                actual_fa_version = getattr(
-                    layer.prefill_backend, "vllm_flash_attn_version", None
-                )
+            prefill_cfg = get_prefill_backend_config(prefill_backend)
+            fa_version = prefill_cfg["flash_attn_version"]
+
+            if fa_version is not None:
+                # FA backend: verify the impl's FA version
+                actual_fa_version = getattr(impl, "vllm_flash_attn_version", None)
                if actual_fa_version != fa_version:
                    raise RuntimeError(
                        f"Prefill backend '{prefill_backend}' requested FA "
-                        f"version {fa_version}, got "
-                        f"{actual_fa_version} on {actual_class}."
+                        f"version {fa_version}, but the impl is using FA "
+                        f"version {actual_fa_version}. Check "
+                        f"vllm/v1/attention/backends/fa_utils.py."
                    )
+            else:
+                # Non-FA backend: verify the builder picked the right path
+                expected_flags = {
+                    "flashinfer": "_use_fi_prefill",
+                    "cudnn": "_use_cudnn_prefill",
+                    "trtllm": "_use_trtllm_ragged_prefill",
+                }
+                flag_name = expected_flags.get(prefill_backend)
+                if flag_name and not getattr(builder_instance, flag_name, False):
+                    raise RuntimeError(
+                        f"Prefill backend '{prefill_backend}' was requested "
+                        f"but the metadata builder did not enable it. This "
+                        f"usually means a dependency is missing (e.g., "
+                        f"flashinfer not installed) or the platform doesn't "
+                        f"support it."
+                    )
+
        # Run each benchmark with the shared impl
        for config, threshold, num_splits in configs_with_params:
            # Set threshold for this benchmark (FlashAttn/FlashMLA only)
@@ -1,415 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Benchmark hidden state extraction throughput.
-
-Measures two modes:
-  1. Baseline: bulk inference with max_tokens=1, no extraction.
-  2. Extract:  async hidden state extraction via ExampleHiddenStatesConnector
-               with N concurrent clients, each consuming hidden states as
-               soon as their request finishes (overlapping I/O with generation).
-
-Reports tokens/s and prompts/s for each mode.
-
-Usage:
-  python benchmarks/benchmark_hidden_state_extraction.py \
-      --model Qwen/Qwen3-0.6B \
-      --num-prompts 64 \
-      --num-clients 8 \
-      --prompt-len 8192 \
-      --layers 1 2 3 4
-"""
-
-import argparse
-import asyncio
-import time
-from concurrent.futures import ThreadPoolExecutor
-
-import torch
-from transformers import AutoConfig
-
-from vllm import LLM, SamplingParams
-from vllm.config.kv_transfer import KVTransferConfig
-from vllm.distributed.kv_transfer.kv_connector.v1 import (
-    example_hidden_states_connector,
-)
-from vllm.engine.arg_utils import AsyncEngineArgs
-from vllm.sampling_params import RequestOutputKind
-from vllm.v1.engine.async_llm import AsyncLLM
-
-
-def _make_profiler_config(profile_dir: str) -> dict:
-    """Build a profiler_config dict for torch profiling."""
-    return {
-        "profiler": "torch",
-        "torch_profiler_dir": profile_dir,
-        "torch_profiler_with_stack": True,
-    }
-
-
-def make_random_prompts(
-    num_prompts: int, prompt_len: int, vocab_size: int, seed: int = 42
-) -> list[list[int]]:
-    """Generate lists of random token IDs."""
-    # Set seed for reproducibility
-    torch.manual_seed(seed)
-    return [
-        torch.randint(0, vocab_size, (prompt_len,)).tolist() for _ in range(num_prompts)
-    ]
-
-
-def consume_hidden_states(path: str) -> float:
-    """Load hidden states from disk and compute per-position mean.
-
-    Returns a single float: the grand mean of all hidden state values.
-    This forces the benchmark to actually read and reduce the data.
-
-    Uses :func:`load_hidden_states` which acquires a shared flock,
-    blocking (without polling) until the async writer releases its
-    exclusive lock.
-    """
-    obj = example_hidden_states_connector.load_hidden_states(path)
-    hs = obj["hidden_states"]
-    total = hs.mean().item()
-
-    example_hidden_states_connector.cleanup_hidden_states(path)
-
-    return total
-
-
-def run_baseline(
-    model: str,
-    prompts: list[list[int]],
-    extra_args: dict,
-    profile_dir: str | None = None,
-) -> dict:
-    """Baseline: bulk inference, no hidden state extraction."""
-    if profile_dir:
-        extra_args = {
-            **extra_args,
-            "profiler_config": _make_profiler_config(profile_dir),
-        }
-    llm = LLM(
-        model=model,
-        enable_prefix_caching=False,
-        enable_chunked_prefill=False,
-        **extra_args,
-    )
-    sampling_params = SamplingParams(max_tokens=1)
-    prompt_inputs = [{"prompt_token_ids": p} for p in prompts]
-
-    # Warmup
-    llm.generate(prompt_inputs[:4], sampling_params, use_tqdm=False)
-
-    if profile_dir:
-        llm.start_profile()
-
-    t0 = time.perf_counter()
-    outputs = llm.generate(prompt_inputs, sampling_params, use_tqdm=True)
-    elapsed = time.perf_counter() - t0
-
-    if profile_dir:
-        llm.stop_profile()
-
-    total_prompt_tokens = sum(len(o.prompt_token_ids) for o in outputs)
-    num_prompts = len(outputs)
-
-    del llm
-    torch.accelerator.empty_cache()
-
-    return {
-        "mode": "baseline",
-        "elapsed_s": elapsed,
-        "num_prompts": num_prompts,
-        "total_prompt_tokens": total_prompt_tokens,
-        "tokens_per_s": total_prompt_tokens / elapsed,
-        "prompts_per_s": num_prompts / elapsed,
-    }
-
-
-# ---- Async extraction benchmark ----
-
-
-async def _client_loop(
-    engine: AsyncLLM,
-    prompt_queue: asyncio.Queue,
-    consume_pool: ThreadPoolExecutor,
-    results: list[dict],
-    client_id: int,
-):
-    """A single async client: pulls prompts, submits to engine, consumes
-    hidden states as soon as each request finishes."""
-    loop = asyncio.get_event_loop()
-    while True:
-        item = await prompt_queue.get()
-        if item is None:
-            prompt_queue.task_done()
-            break
-        idx, token_ids = item
-
-        request_id = f"req-{idx}"
-        sampling_params = SamplingParams(
-            max_tokens=1,
-            output_kind=RequestOutputKind.FINAL_ONLY,
-        )
-
-        final_output = None
-        async for output in engine.generate(
-            request_id=request_id,
-            prompt={"prompt_token_ids": token_ids},
-            sampling_params=sampling_params,
-        ):
-            if output.finished:
-                final_output = output
-
-        # Consume hidden states on a thread (disk I/O)
-        path = final_output.kv_transfer_params["hidden_states_path"]
-        mean_val = await loop.run_in_executor(consume_pool, consume_hidden_states, path)
-        num_tokens = len(final_output.prompt_token_ids)
-
-        results.append(
-            {
-                "request_id": request_id,
-                "num_prompt_tokens": num_tokens,
-                "mean_hidden_value": mean_val,
-            }
-        )
-        prompt_queue.task_done()
-
-
-async def _run_extraction_async(
-    model: str,
-    prompts: list[list[int]],
-    num_clients: int,
-    layers: list[int],
-    tmpdir: str,
-    extra_args: dict,
-    profile_dir: str | None = None,
-) -> dict:
-    if profile_dir:
-        extra_args = {
-            **extra_args,
-            "profiler_config": _make_profiler_config(profile_dir),
-        }
-    engine_args = AsyncEngineArgs(
-        model=model,
-        enable_prefix_caching=False,
-        enable_chunked_prefill=False,
-        max_num_batched_tokens=40960,
-        max_model_len=40960,
-        speculative_config={
-            "method": "extract_hidden_states",
-            "num_speculative_tokens": 1,
-            "draft_model_config": {
-                "hf_config": {
-                    "eagle_aux_hidden_state_layer_ids": layers,
-                },
-            },
-        },
-        kv_transfer_config=KVTransferConfig(
-            kv_connector="ExampleHiddenStatesConnector",
-            kv_role="kv_producer",
-            kv_connector_extra_config={
-                "shared_storage_path": tmpdir,
-            },
-        ),
-        **extra_args,
-    )
-    engine = AsyncLLM.from_engine_args(engine_args)
-
-    try:
-        # Warmup: run a few prompts sequentially, cleaning up generated files
-        for i in range(min(4, len(prompts))):
-            sp = SamplingParams(max_tokens=1, output_kind=RequestOutputKind.FINAL_ONLY)
-            final_output = None
-            async for output in engine.generate(
-                request_id=f"warmup-{i}",
-                prompt={"prompt_token_ids": prompts[i]},
-                sampling_params=sp,
-            ):
-                if output.finished:
-                    final_output = output
-            if final_output and final_output.kv_transfer_params:
-                path = final_output.kv_transfer_params.get("hidden_states_path")
-                if path:
-                    example_hidden_states_connector.cleanup_hidden_states(path)
-
-        if profile_dir:
-            await engine.start_profile()
-
-        # Fill prompt queue
-        prompt_queue: asyncio.Queue = asyncio.Queue()
-        for idx, token_ids in enumerate(prompts):
-            prompt_queue.put_nowait((idx, token_ids))
-        # Sentinel per client
-        for _ in range(num_clients):
-            prompt_queue.put_nowait(None)
-
-        results: list[dict] = []
-        consume_pool = ThreadPoolExecutor(max_workers=num_clients)
-
-        t0 = time.perf_counter()
-        tasks = [
-            asyncio.create_task(
-                _client_loop(engine, prompt_queue, consume_pool, results, i)
-            )
-            for i in range(num_clients)
-        ]
-        await asyncio.gather(*tasks)
-        elapsed = time.perf_counter() - t0
-
-        consume_pool.shutdown(wait=True)
-
-        if profile_dir:
-            await engine.stop_profile()
-
-        total_prompt_tokens = sum(r["num_prompt_tokens"] for r in results)
-        num_prompts = len(results)
-        mean_hidden = sum(r["mean_hidden_value"] for r in results) / max(
-            len(results), 1
-        )
-
-        return {
-            "mode": "extract",
-            "elapsed_s": elapsed,
-            "num_prompts": num_prompts,
-            "total_prompt_tokens": total_prompt_tokens,
-            "tokens_per_s": total_prompt_tokens / elapsed,
-            "prompts_per_s": num_prompts / elapsed,
-            "mean_hidden_value": mean_hidden,
-        }
-    finally:
-        engine.shutdown()
-
-
-def run_extraction(
-    model: str,
-    prompts: list[list[int]],
-    num_clients: int,
-    layers: list[int],
-    extra_args: dict,
-    profile_dir: str | None = None,
-) -> dict:
-    return asyncio.run(
-        _run_extraction_async(
-            model,
-            prompts,
-            num_clients,
-            layers,
-            "/dev/shm",
-            extra_args,
-            profile_dir=profile_dir,
-        )
-    )
-
-
-def print_results(results: dict):
-    mode = results["mode"]
-    print(f"\n{'=' * 60}")
-    print(f"  {mode.upper()} RESULTS")
-    print(f"{'=' * 60}")
-    print(f"  Prompts:             {results['num_prompts']}")
-    print(f"  Total prompt tokens: {results['total_prompt_tokens']:,}")
-    print(f"  Wall time:           {results['elapsed_s']:.2f}s")
-    print(f"  Tokens/s:            {results['tokens_per_s']:,.0f}")
-    print(f"  Prompts/s:           {results['prompts_per_s']:.2f}")
-    if mode == "extract":
-        print(f"  Mean hidden value:   {results['mean_hidden_value']:.6f}")
-    print(f"{'=' * 60}\n")
-
-
-def main():
-    parser = argparse.ArgumentParser(
-        description="Benchmark hidden state extraction throughput"
-    )
-    parser.add_argument("--model", type=str, required=True)
-    parser.add_argument("--num-prompts", type=int, default=64)
-    parser.add_argument("--num-clients", type=int, default=8)
-    parser.add_argument("--prompt-len", type=int, default=8192)
-    parser.add_argument("--layers", type=int, nargs="+", default=[1, 2, 3, 4])
-    parser.add_argument("--skip-baseline", action="store_true")
-    parser.add_argument("--skip-extract", action="store_true")
-    parser.add_argument("--gpu-memory-utilization", type=float, default=0.9)
-    parser.add_argument("--max-num-batched-tokens", type=int, default=None)
-    parser.add_argument("--max-cudagraph-capture-size", type=int, default=None)
-    parser.add_argument("--max-model-len", type=int, default=None)
-    parser.add_argument("--enforce-eager", action="store_true")
-    parser.add_argument("--load-format", type=str, default=None)
-    parser.add_argument(
-        "--profile",
-        action="store_true",
-        help="Enable torch profiler for both baseline and extraction runs.",
-    )
-    parser.add_argument(
-        "--torch-profiler-dir",
-        type=str,
-        default="./vllm_profile",
-        help="Directory to save torch profiler traces (default: ./vllm_profile).",
-    )
-    parser.add_argument(
-        "--enable-flashinfer-autotune",
-        action="store_true",
-        default=False,
-        help="Enable FlashInfer autotuning (can be slow).",
-    )
-    args = parser.parse_args()
-
-    extra_args = {
-        "gpu_memory_utilization": args.gpu_memory_utilization,
-    }
-    if args.max_model_len is not None:
-        extra_args["max_model_len"] = args.max_model_len
-    if args.max_num_batched_tokens is not None:
-        extra_args["max_num_batched_tokens"] = args.max_num_batched_tokens
-        if args.max_model_len and args.max_num_batched_tokens < args.max_model_len:
-            raise ValueError(
-                "max_num_batched_tokens must be >= max_model_len since chunked prefill"
-                " is not supported by hidden state extraction."
-            )
-    if args.enforce_eager:
-        extra_args["enforce_eager"] = True
-    if args.load_format is not None:
-        extra_args["load_format"] = args.load_format
-    if args.max_cudagraph_capture_size is not None:
-        extra_args["max_cudagraph_capture_size"] = args.max_cudagraph_capture_size
-    extra_args["enable_flashinfer_autotune"] = args.enable_flashinfer_autotune
-
-    # Get vocab size from HF config without loading the full model
-    hf_config = AutoConfig.from_pretrained(args.model, trust_remote_code=True)
-    vocab_size = hf_config.vocab_size
-    prompts = make_random_prompts(args.num_prompts, args.prompt_len, vocab_size)
-    print(
-        f"Generated {args.num_prompts} prompts, "
-        f"{args.prompt_len} tokens each (vocab {vocab_size})"
-    )
-
-    profile_dir = args.torch_profiler_dir if args.profile else None
-    if profile_dir:
-        print(f"Torch profiler enabled, traces will be saved to {profile_dir}/")
-
-    if not args.skip_baseline:
-        baseline_profile_dir = f"{profile_dir}/baseline" if profile_dir else None
-        baseline = run_baseline(
-            args.model, prompts, extra_args, profile_dir=baseline_profile_dir
-        )
-        print_results(baseline)
-
-    if not args.skip_extract:
-        extract_profile_dir = f"{profile_dir}/extract" if profile_dir else None
-        extract = run_extraction(
-            args.model,
-            prompts,
-            args.num_clients,
-            args.layers,
-            extra_args,
-            profile_dir=extract_profile_dir,
-        )
-        print_results(extract)
-
-    if not args.skip_baseline and not args.skip_extract:
-        slowdown = baseline["tokens_per_s"] / extract["tokens_per_s"]
-        print("Extraction slowdown factor: {:.2f}x".format(slowdown))
-
-
-if __name__ == "__main__":
-    main()
@@ -115,39 +115,6 @@ class SampleRequest:
 def sample_requests(
    tokenizer: PreTrainedTokenizerBase, args: argparse.Namespace
 ) -> list[SampleRequest]:
-    def _apply_random_prefix(
-        tokenizer: PreTrainedTokenizerBase,
-        requests: list[SampleRequest],
-        prefix_len: int,
-        seed: int,
-    ) -> list[SampleRequest]:
-        if prefix_len <= 0:
-            return requests
-        rng = np.random.default_rng(seed)
-        vocab_size = tokenizer.vocab_size
-        prohibited = getattr(tokenizer, "all_special_ids", None) or []
-        allowed = np.array([i for i in range(vocab_size) if i not in prohibited])
-        if len(allowed) == 0:
-            return requests
-        prefix_ids = rng.integers(0, len(allowed), size=prefix_len)
-        prefix_token_ids = allowed[prefix_ids].tolist()
-        out = []
-        for req in requests:
-            prompt_ids = tokenizer(req.prompt, add_special_tokens=False).input_ids
-            full_ids = prefix_token_ids + prompt_ids
-            full_prompt = tokenizer.decode(full_ids, skip_special_tokens=False)
-            out.append(
-                SampleRequest(
-                    prompt=full_prompt,
-                    prompt_len=len(tokenizer(full_prompt).input_ids),
-                    expected_output_len=req.expected_output_len,
-                    schema=req.schema,
-                    structure_type=req.structure_type,
-                    completion=req.completion,
-                )
-            )
-        return out
-
    if args.dataset == "json" or args.dataset == "json-unique":
        if args.json_schema_path is None:
            dir_path = os.path.dirname(os.path.realpath(__file__))
@@ -294,9 +261,6 @@ def sample_requests(
                )
            )

-    requests = _apply_random_prefix(
-        tokenizer, requests, args.random_prefix_len, args.seed
-    )
    return requests


@@ -981,15 +945,6 @@ def create_argument_parser():
        "results in a more uniform arrival of requests.",
    )
    parser.add_argument("--seed", type=int, default=0)
-    parser.add_argument(
-        "--random-prefix-len",
-        type=int,
-        default=0,
-        help=(
-            "Number of prefix tokens to prepend to every prompt. "
-            "The same prefix is used for all prompts to enable prefix caching."
-        ),
-    )
    parser.add_argument(
        "--trust-remote-code",
        action="store_true",
@@ -1,465 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Benchmark the fused MoE-LoRA fast path (one-shot) vs two-kernel baseline.
-
-The "one_shot" provider goes through `vllm.lora.ops.triton_ops.fused_moe_lora`
-which dispatches to the single-kernel one-shot implementation when
-fully_sharded=False (the prefill default).
-
-The "two_kernel" provider drives `fused_moe_lora_shrink` + `fused_moe_lora_expand`
-directly, bypassing the dispatch and matching the legacy two-kernel path's
-work distribution. This isolates the win from kernel fusion.
-
-Run:
-    .venv/bin/python -m benchmarks.kernels.benchmark_fused_moe_lora_one_shot
-    .venv/bin/python -m benchmarks.kernels.benchmark_fused_moe_lora_one_shot \\
-        --model qwen3moe
-"""
-
-from __future__ import annotations
-
-import argparse
-import os
-import random
-
-import torch
-
-from vllm import _custom_ops as ops
-from vllm.lora.ops.triton_ops import (
-    fused_moe_lora,
-    fused_moe_lora_expand,
-    fused_moe_lora_shrink,
-)
-from vllm.triton_utils import triton
-
-DTYPE = torch.bfloat16
-DEVICE = "cuda"
-
-
-# ----- input fabrication -----------------------------------------------------
-
-
-def _round_up(x: int, base: int) -> int:
-    return ((x + base - 1) // base) * base
-
-
-def _ceildiv(x: int, y: int) -> int:
-    return (x + y - 1) // y
-
-
-def _assign_loras(num_tokens: int, num_sequences: int, max_loras: int) -> torch.Tensor:
-    tokens_per_seq = num_tokens // num_sequences
-    rem = num_tokens % num_sequences
-    out = torch.empty(num_tokens, dtype=torch.int32)
-    start = 0
-    for i in range(num_sequences):
-        end = start + tokens_per_seq + (1 if i < rem else 0)
-        out[start:end] = random.randint(0, max_loras - 1)
-        start = end
-    return out
-
-
-def _assign_experts(num_tokens: int, num_experts: int, top_k: int):
-    expert_indices = torch.empty((num_tokens, top_k), dtype=torch.int32)
-    for i in range(num_tokens):
-        expert_indices[i] = torch.randperm(num_experts)[:top_k]
-    weights = torch.rand((num_tokens, top_k), dtype=torch.float32)
-    weights = weights / weights.sum(dim=1, keepdim=True)
-    return expert_indices, weights
-
-
-def _make_inputs(
-    M: int,
-    K: int,
-    N_per_slice: int,
-    rank: int,
-    num_experts: int,
-    top_k: int,
-    max_loras: int,
-    num_slices: int,
-    block_size_m: int,
-):
-    """Mirrors the production caller's tensor layout."""
-    torch.manual_seed(0)
-    random.seed(0)
-
-    num_sequences = max(1, min(M, 8))
-    topk_ids_cpu, topk_weights_cpu = _assign_experts(M, num_experts, top_k)
-    token_lora_cpu = _assign_loras(M, num_sequences, max_loras)
-    lora_ids_cpu = torch.full((max_loras + 1,), -1, dtype=torch.int32)
-    uniq = torch.unique(token_lora_cpu, sorted=True)
-    lora_ids_cpu[: uniq.size(0)].copy_(uniq)
-
-    topk_ids = topk_ids_cpu.to(DEVICE)
-    topk_weights = topk_weights_cpu.to(device=DEVICE, dtype=DTYPE)
-    token_lora_mapping = token_lora_cpu.to(DEVICE)
-    lora_ids = lora_ids_cpu.to(DEVICE)
-    adapter_enabled = torch.ones(max_loras + 1, dtype=torch.int32, device=DEVICE)
-
-    lora_a = [
-        torch.randn((max_loras, num_experts, rank, K), dtype=DTYPE, device=DEVICE)
-        / max(K, 1) ** 0.5
-        for _ in range(num_slices)
-    ]
-    lora_b = [
-        torch.randn(
-            (max_loras, num_experts, N_per_slice, rank),
-            dtype=DTYPE,
-            device=DEVICE,
-        )
-        / max(rank, 1) ** 0.5
-        for _ in range(num_slices)
-    ]
-    hidden = torch.randn((M, K), dtype=DTYPE, device=DEVICE)
-    out_template = torch.zeros(
-        (M, top_k, num_slices * N_per_slice), dtype=DTYPE, device=DEVICE
-    )
-
-    # Sorted-path metadata (the prefill default).
-    max_pad = topk_ids.numel() + num_experts * (block_size_m - 1)
-    max_pad = _round_up(max_pad, block_size_m)
-    max_blocks = _ceildiv(max_pad, block_size_m)
-    sorted_token_ids = torch.empty(
-        (max_loras * max_pad,), dtype=torch.int32, device=DEVICE
-    )
-    expert_ids = torch.empty(
-        (max_loras * max_blocks,), dtype=torch.int32, device=DEVICE
-    )
-    num_post = torch.empty((max_loras,), dtype=torch.int32, device=DEVICE)
-    ops.moe_lora_align_block_size(
-        topk_ids,
-        token_lora_mapping,
-        num_experts,
-        block_size_m,
-        max_loras,
-        max_pad,
-        max_blocks,
-        sorted_token_ids,
-        expert_ids,
-        num_post,
-        adapter_enabled,
-        lora_ids,
-    )
-    expert_ids = expert_ids.view(max_loras, -1).contiguous()
-    sorted_token_ids = sorted_token_ids.view(max_loras, -1).contiguous()
-    num_active = torch.tensor([max_loras + 1], dtype=torch.int32, device="cpu")
-
-    return dict(
-        hidden=hidden,
-        lora_a=lora_a,
-        lora_b=lora_b,
-        topk_weights=topk_weights,
-        sorted_token_ids=sorted_token_ids,
-        expert_ids=expert_ids,
-        num_post=num_post,
-        token_lora_mapping=token_lora_mapping,
-        lora_ids=lora_ids,
-        num_active=num_active,
-        adapter_enabled=adapter_enabled,
-        out_template=out_template,
-        # bookkeeping
-        M=M,
-        K=K,
-        N_per_slice=N_per_slice,
-        rank=rank,
-        num_experts=num_experts,
-        top_k=top_k,
-        max_loras=max_loras,
-        num_slices=num_slices,
-        block_size_m=block_size_m,
-    )
-
-
-# ----- providers -------------------------------------------------------------
-
-
-def _run_one_shot(inp: dict):
-    """Drive `fused_moe_lora` with fully_sharded=False -> one-shot fast path."""
-    out = inp["out_template"].clone()
-    fused_moe_lora(
-        out,
-        inp["hidden"],
-        inp["lora_a"],
-        inp["lora_b"],
-        inp["topk_weights"],
-        inp["sorted_token_ids"],
-        inp["expert_ids"],
-        inp["num_post"],
-        inp["token_lora_mapping"],
-        inp["rank"],
-        inp["top_k"],
-        inp["lora_ids"],
-        inp["num_active"],
-        inp["adapter_enabled"],
-        inp["block_size_m"],
-        64,
-        32,
-        8,
-        4,
-        3,
-        1,
-        inp["block_size_m"],
-        64,
-        32,
-        8,
-        4,
-        3,
-        1,
-        False,
-        False,
-        0,
-    )
-    return out
-
-
-def _run_two_kernel(inp: dict):
-    """Drive `fused_moe_lora_shrink` + `fused_moe_lora_expand` directly,
-    bypassing the dispatch. Matches the legacy two-kernel work distribution.
-    """
-    M = inp["M"]
-    top_k = inp["top_k"]
-    rank = inp["rank"]
-    num_slices = inp["num_slices"]
-    N_per_slice = inp["N_per_slice"]
-    K = inp["K"]
-    num_experts = inp["num_experts"]
-    block_m = inp["block_size_m"]
-
-    intermediate = torch.zeros((num_slices, M, top_k, rank), dtype=DTYPE, device=DEVICE)
-    out = inp["out_template"].clone()
-    EM = inp["sorted_token_ids"].shape[1]
-    num_tokens = M * top_k
-
-    fused_moe_lora_shrink(
-        intermediate,
-        inp["hidden"],
-        inp["lora_a"],
-        inp["topk_weights"],
-        inp["sorted_token_ids"],
-        inp["expert_ids"],
-        inp["num_post"],
-        inp["token_lora_mapping"],
-        top_k,
-        inp["lora_ids"],
-        inp["adapter_enabled"],
-        torch.device(DEVICE),
-        rank,
-        M,
-        EM,
-        K,
-        num_tokens,
-        num_experts,
-        num_slices,
-        block_m,
-        64,
-        32,
-        8,
-        4,
-        3,
-        1,
-        inp["num_active"],
-        False,
-    )
-    fused_moe_lora_expand(
-        out,
-        intermediate,
-        inp["lora_b"],
-        inp["topk_weights"],
-        inp["sorted_token_ids"],
-        inp["expert_ids"],
-        inp["num_post"],
-        inp["token_lora_mapping"],
-        top_k,
-        inp["lora_ids"],
-        inp["adapter_enabled"],
-        torch.device(DEVICE),
-        rank,
-        M,
-        EM,
-        K,
-        num_tokens,
-        num_experts,
-        num_slices,
-        rank,
-        N_per_slice,
-        block_m,
-        64,
-        32,
-        8,
-        4,
-        3,
-        1,
-        inp["num_active"],
-        False,
-        0,
-    )
-    return out
-
-
-PROVIDER_FNS = {
-    "one_shot": _run_one_shot,
-    "two_kernel": _run_two_kernel,
-}
-
-
-# ----- model presets ---------------------------------------------------------
-
-
-MODEL_PRESETS: dict[str, dict] = {
-    # Mixtral-8x7B style: E=8, top_k=2, hidden=4096, intermediate=14336
-    "mixtral": dict(
-        K=4096,
-        N_per_slice=7168,
-        num_experts=8,
-        top_k=2,
-        max_loras=4,
-        num_slices=2,
-        block_size_m=64,
-    ),
-    # Qwen3-MoE / DeepSeek-V2 style: E=64, top_k=8, hidden=2048, inter=1408
-    "qwen3moe": dict(
-        K=2048,
-        N_per_slice=1408,
-        num_experts=64,
-        top_k=8,
-        max_loras=4,
-        num_slices=2,
-        block_size_m=64,
-    ),
-    # GLM-5.1 (zai-org/GLM-5.1-FP8): E=256, top_k=8, hidden=6144,
-    # moe_intermediate=2048
-    "glm5_1": dict(
-        K=6144,
-        N_per_slice=2048,
-        num_experts=256,
-        top_k=8,
-        max_loras=4,
-        num_slices=2,
-        block_size_m=64,
-    ),
-}
-
-
-M_RANGE = [16, 64, 256, 1024, 4096, 16384]
-RANK_RANGE = [8, 16, 32, 64]
-
-
-def get_benchmark(model: str, max_loras: int | None = None):
-    preset = dict(MODEL_PRESETS[model])
-    if max_loras is not None:
-        preset["max_loras"] = max_loras
-
-    @triton.testing.perf_report(
-        triton.testing.Benchmark(
-            x_names=["M", "rank"],
-            x_vals=[(M, R) for M in M_RANGE for R in RANK_RANGE],
-            line_arg="provider",
-            line_vals=list(PROVIDER_FNS.keys()),
-            line_names=["one_shot (fused)", "two_kernel (legacy)"],
-            styles=[("red", "-"), ("blue", "-")],
-            ylabel="ms",
-            plot_name=f"fused_moe_lora-{model}-loras{preset['max_loras']}",
-            args={"preset": preset},
-        )
-    )
-    def benchmark(M, rank, provider, preset):
-        inp = _make_inputs(
-            M=M,
-            K=preset["K"],
-            N_per_slice=preset["N_per_slice"],
-            rank=rank,
-            num_experts=preset["num_experts"],
-            top_k=preset["top_k"],
-            max_loras=preset["max_loras"],
-            num_slices=preset["num_slices"],
-            block_size_m=preset["block_size_m"],
-        )
-        fn = PROVIDER_FNS[provider]
-        quantiles = [0.5, 0.2, 0.8]
-        ms, min_ms, max_ms = triton.testing.do_bench(
-            lambda: fn(inp), quantiles=quantiles
-        )
-        return ms, max_ms, min_ms
-
-    return benchmark
-
-
-# ----- correctness sanity ---------------------------------------------------
-
-
-def calculate_diff(model: str, M: int, rank: int, max_loras: int | None = None):
-    preset = dict(MODEL_PRESETS[model])
-    if max_loras is not None:
-        preset["max_loras"] = max_loras
-    inp = _make_inputs(
-        M=M,
-        K=preset["K"],
-        N_per_slice=preset["N_per_slice"],
-        rank=rank,
-        num_experts=preset["num_experts"],
-        top_k=preset["top_k"],
-        max_loras=preset["max_loras"],
-        num_slices=preset["num_slices"],
-        block_size_m=preset["block_size_m"],
-    )
-    out_one = _run_one_shot(inp)
-    out_two = _run_two_kernel(inp)
-    max_abs = (out_one.float() - out_two.float()).abs().max().item()
-    print(
-        f"  model={model:<9} M={M:<6} rank={rank:<3}  "
-        f"max|one_shot - two_kernel|={max_abs:.4g}  "
-        f"ref|max|={out_two.float().abs().max().item():.3g}"
-    )
-    if max_abs <= 5e-2:
-        print("  ✅ outputs match within bf16 tolerance")
-    else:
-        print("  ❌ outputs differ beyond expected bf16 noise")
-
-
-# ----- main ------------------------------------------------------------------
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "--model",
-        type=str,
-        default="mixtral",
-        choices=list(MODEL_PRESETS.keys()),
-        help="Model preset to sweep",
-    )
-    parser.add_argument(
-        "--save-path",
-        type=str,
-        default="./configs/fused_moe_lora_one_shot/",
-        help="Directory to save benchmark results",
-    )
-    parser.add_argument(
-        "--check-only",
-        action="store_true",
-        help="Run correctness sanity check only, no perf sweep",
-    )
-    parser.add_argument(
-        "--max-loras",
-        type=int,
-        default=None,
-        help="Override max_loras in the model preset (number of LoRA adapters "
-        "active in the batch). Defaults to the preset's value.",
-    )
-    args = parser.parse_args()
-
-    print(f"Correctness check ({args.model}):")
-    calculate_diff(args.model, M=256, rank=32, max_loras=args.max_loras)
-    if args.check_only:
-        raise SystemExit(0)
-
-    effective_max_loras = (
-        args.max_loras
-        if args.max_loras is not None
-        else MODEL_PRESETS[args.model]["max_loras"]
-    )
-    print(f"\nGPU: {torch.cuda.get_device_name()}")
-    print(f"Model preset: {args.model}  max_loras={effective_max_loras}\n")
-    benchmark = get_benchmark(args.model, max_loras=args.max_loras)
-    os.makedirs(args.save_path, exist_ok=True)
-    benchmark.run(print_data=True, save_path=args.save_path)
@@ -27,10 +27,10 @@ from vllm.model_executor.layers.fused_moe.config import (
    RoutingMethodType,
    _get_config_dtype_str,
 )
-from vllm.model_executor.layers.fused_moe.experts.triton_deep_gemm_moe import (
+from vllm.model_executor.layers.fused_moe.fused_moe import *
+from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (
    TritonOrDeepGemmExperts,
 )
-from vllm.model_executor.layers.fused_moe.fused_moe import *
 from vllm.transformers_utils.config import get_config
 from vllm.triton_utils import triton
 from vllm.utils.argparse_utils import FlexibleArgumentParser
@@ -10,7 +10,6 @@ from transformers import AutoConfig

 from vllm.model_executor.layers.fused_moe import fused_topk
 from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import (
-    MoEPermuteScratch,
    moe_permute,
    moe_unpermute,
 )
@@ -55,15 +54,6 @@ def benchmark_permute(
    topk_weights, topk_ids, token_expert_indices = fused_topk(
        qhidden_states, input_gating, topk, False
    )
-    scratch = MoEPermuteScratch(
-        max_num_tokens=num_tokens,
-        topk=topk,
-        num_experts=num_experts,
-        num_local_experts=num_experts,
-        device=qhidden_states.device,
-        hidden_size=hidden_size,
-        hidden_dtype=qhidden_states.dtype,
-    )

    def prepare(i: int):
        input_gating.copy_(gating_output[i])
@@ -75,7 +65,6 @@ def benchmark_permute(
            topk_ids=topk_ids,
            n_expert=num_experts,
            expert_map=None,
-            scratch=scratch,
        )

    # JIT compilation & warmup
@@ -134,15 +123,6 @@ def benchmark_unpermute(
    topk_weights, topk_ids, token_expert_indices = fused_topk(
        qhidden_states, input_gating, topk, False
    )
-    scratch = MoEPermuteScratch(
-        max_num_tokens=num_tokens,
-        topk=topk,
-        num_experts=num_experts,
-        num_local_experts=num_experts,
-        device=qhidden_states.device,
-        hidden_size=hidden_size,
-        hidden_dtype=qhidden_states.dtype,
-    )

    def prepare():
        (
@@ -157,7 +137,6 @@ def benchmark_unpermute(
            topk_ids=topk_ids,
            n_expert=num_experts,
            expert_map=None,
-            scratch=scratch,
        )
        # convert to fp16/bf16 as gemm output
        return (
@@ -1,154 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import torch
-import torch.nn.functional as F
-
-from vllm import _custom_ops as ops
-from vllm.platforms import current_platform
-from vllm.transformers_utils.config import get_config
-from vllm.triton_utils import triton
-from vllm.utils.argparse_utils import FlexibleArgumentParser
-
-# Dimensions supported by the DSV3 specialized kernel
-DSV3_SUPPORTED_NUM_EXPERTS = [256, 384]
-DSV3_SUPPORTED_HIDDEN_SIZES = [7168]
-
-# Dimensions supported by the gpt-oss specialized kernel
-GPT_OSS_SUPPORTED_NUM_EXPERTS = [32, 128]
-GPT_OSS_SUPPORTED_HIDDEN_SIZES = [2880]
-
-# Dimensions supported by the fp32 specialized kernel (MiniMax-M2)
-FP32_SUPPORTED_NUM_EXPERTS = [256]
-FP32_SUPPORTED_HIDDEN_SIZES = [3072]
-FP32_MAX_TOKENS = 32
-
-
-def get_batch_size_range(max_batch_size):
-    return [2**x for x in range(14) if 2**x <= max_batch_size]
-
-
-def get_model_params(config):
-    if config.architectures[0] in (
-        "DeepseekV2ForCausalLM",
-        "DeepseekV3ForCausalLM",
-        "DeepseekV32ForCausalLM",
-    ):
-        num_experts = config.n_routed_experts
-        hidden_size = config.hidden_size
-    elif config.architectures[0] in ("GptOssForCausalLM",) or config.architectures[
-        0
-    ] in ("MiniMaxM2ForCausalLM",):
-        num_experts = config.num_local_experts
-        hidden_size = config.hidden_size
-    else:
-        raise ValueError(f"Unsupported architecture: {config.architectures}")
-    return num_experts, hidden_size
-
-
-def get_benchmark(model, max_batch_size, trust_remote_code):
-    @triton.testing.perf_report(
-        triton.testing.Benchmark(
-            x_names=["batch_size"],
-            x_vals=get_batch_size_range(max_batch_size),
-            x_log=False,
-            line_arg="provider",
-            line_vals=[
-                "torch",
-                "vllm",
-            ],
-            line_names=["PyTorch", "vLLM"],
-            styles=([("blue", "-"), ("red", "-")]),
-            ylabel="TFLOPs",
-            plot_name=f"{model} router gemm throughput",
-            args={},
-        )
-    )
-    def benchmark(batch_size, provider):
-        config = get_config(model=model, trust_remote_code=trust_remote_code)
-        num_experts, hidden_size = get_model_params(config)
-
-        is_hopper_or_blackwell = current_platform.is_device_capability(
-            90
-        ) or current_platform.is_device_capability_family(100)
-        allow_dsv3_router_gemm = (
-            is_hopper_or_blackwell
-            and num_experts in DSV3_SUPPORTED_NUM_EXPERTS
-            and hidden_size in DSV3_SUPPORTED_HIDDEN_SIZES
-        )
-        allow_gpt_oss_router_gemm = (
-            is_hopper_or_blackwell
-            and num_experts in GPT_OSS_SUPPORTED_NUM_EXPERTS
-            and hidden_size in GPT_OSS_SUPPORTED_HIDDEN_SIZES
-        )
-        is_fp32_router_model = (
-            is_hopper_or_blackwell
-            and num_experts in FP32_SUPPORTED_NUM_EXPERTS
-            and hidden_size in FP32_SUPPORTED_HIDDEN_SIZES
-        )
-        allow_fp32_router_gemm = is_fp32_router_model and batch_size <= FP32_MAX_TOKENS
-
-        # Weight dtype: fp32 kernel requires fp32 weights; others use bf16.
-        weight_dtype = torch.float32 if is_fp32_router_model else torch.bfloat16
-        mat_a = torch.randn(
-            (batch_size, hidden_size), dtype=torch.bfloat16, device="cuda"
-        ).contiguous()
-        mat_b = torch.randn(
-            (num_experts, hidden_size), dtype=weight_dtype, device="cuda"
-        ).contiguous()
-        bias = torch.randn(
-            num_experts, dtype=torch.bfloat16, device="cuda"
-        ).contiguous()
-
-        has_bias = allow_gpt_oss_router_gemm
-
-        quantiles = [0.5, 0.2, 0.8]
-
-        if provider == "torch":
-
-            def runner():
-                if allow_fp32_router_gemm:
-                    F.linear(mat_a.float(), mat_b)
-                elif has_bias:
-                    F.linear(mat_a, mat_b, bias)
-                else:
-                    F.linear(mat_a, mat_b)
-        elif provider == "vllm":
-
-            def runner():
-                if allow_dsv3_router_gemm:
-                    ops.dsv3_router_gemm(mat_a, mat_b, torch.bfloat16)
-                elif allow_fp32_router_gemm:
-                    ops.fp32_router_gemm(mat_a, mat_b)
-                elif allow_gpt_oss_router_gemm:
-                    ops.gpt_oss_router_gemm(mat_a, mat_b, bias)
-                elif is_fp32_router_model:
-                    # batch_size > FP32_MAX_TOKENS: fall back to F.linear
-                    F.linear(mat_a.float(), mat_b)
-                else:
-                    F.linear(mat_a, mat_b)
-
-        ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
-            runner, quantiles=quantiles
-        )
-
-        def tflops(t_ms):
-            flops = 2 * batch_size * hidden_size * num_experts
-            return flops / (t_ms * 1e-3) / 1e12
-
-        return tflops(ms), tflops(max_ms), tflops(min_ms)
-
-    return benchmark
-
-
-if __name__ == "__main__":
-    parser = FlexibleArgumentParser()
-    parser.add_argument("--model", type=str, default="openai/gpt-oss-20b")
-    parser.add_argument("--max-batch-size", default=16, type=int)
-    parser.add_argument("--trust-remote-code", action="store_true")
-    args = parser.parse_args()
-
-    # Get the benchmark function
-    benchmark = get_benchmark(args.model, args.max_batch_size, args.trust_remote_code)
-    # Run performance benchmark
-    benchmark.run(print_data=True)
@@ -1,774 +0,0 @@
-#!/usr/bin/env python3
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Benchmark and tuning script for the Mamba selective_state_update kernel.
-
-Mirrors the fused MoE tuning workflow: sweeps (BLOCK_SIZE_M, num_warps) across
-an effective_batch grid for a given (headdim, dstate, ngroups, cache_dtype) and
-saves the best config per effective_batch to JSON. Generated configs are picked
-up by selective_state_update at runtime.
-
-Usage:
-    python -m benchmarks.kernels.benchmark_selective_state_update \
-        --all-dstates --save-configs --compare
-"""
-
-import argparse
-import json
-import os
-import sys
-from io import StringIO
-from itertools import product
-from typing import Any
-
-import torch
-
-from tests.kernels.mamba.utils import selective_state_update_ref
-from vllm.model_executor.layers.mamba.ops.mamba_ssm import (
-    _CONFIGS_DIR,
-    _canonical_cache_dtype,
-    _get_default_ssm_launch_config,
-    get_ssm_config_file_name,
-    get_ssm_device_name,
-    override_ssm_config,
-    selective_state_update,
-)
-from vllm.triton_utils import triton
-
-# bf16 shares configs with fp16 - same bit width.
-_SSM_CACHE_DTYPE_MAP: dict[str, torch.dtype] = {
-    "float32": torch.float32,
-    "float16": torch.float16,
-    "bfloat16": torch.float16,
-}
-
-_RESULTS_DIR = os.path.dirname(os.path.realpath(__file__))
-
-# ---------------------------------------------------------------------------
-# Tuning search space
-# ---------------------------------------------------------------------------
-
-_BSM_CHOICES_ALL = [4, 8, 16, 32, 64, 128, 256]
-
-NUM_WARPS_CHOICES = [1, 2, 4, 8]
-
-
-def _block_size_m_choices(headdim: int) -> list[int]:
-    """BLOCK_SIZE_M candidates worth sweeping for a given headdim.
-
-    BLOCK_SIZE_M > next_pow2(headdim) wastes >=50% of each tile via masking
-    (offs_m >= dim rows are zeroed out), so we cap the sweep there.
-    """
-    ceiling = 1
-    while ceiling < headdim:
-        ceiling <<= 1
-    return [b for b in _BSM_CHOICES_ALL if b <= ceiling]
-
-
-# Default deployment shapes. effective_batch = batch * nheads scales the
-# kernel grid, so configs transfer across (model, TP) combos sharing
-# (headdim, dstate, cache_dtype).
-DEFAULT_BATCH_SIZES = [1, 8, 16, 32, 64, 128, 256, 512, 1024, 1536, 2048]
-DEFAULT_NHEADS = [128, 256]
-
-ALL_DSTATES = [16, 32, 64, 128, 256]
-
-# Default tuning shape — matches Nemotron-3-Super and Nemotron-3-Nano Mamba layers.
-# Override with CLI flags for other architectures.
-DEFAULT_HEADDIM = 64
-DEFAULT_NGROUPS = 8
-
-
-# ---------------------------------------------------------------------------
-# Benchmark helper
-# ---------------------------------------------------------------------------
-
-
-def _make_inputs(
-    batch: int,
-    nheads: int,
-    dim: int,
-    dstate: int,
-    ngroups: int,
-    dtype: torch.dtype,
-    state_dtype: torch.dtype | None = None,
-    device: str = "cuda",
-):
-    if state_dtype is None:
-        state_dtype = dtype
-    state = torch.randn(batch, nheads, dim, dstate, dtype=state_dtype, device=device)
-    x = torch.randn(batch, nheads, dim, dtype=dtype, device=device)
-    dt = torch.randn(batch, nheads, dim, dtype=dtype, device=device)
-    A = -torch.rand(nheads, dim, dstate, dtype=torch.float32, device=device)
-    B = torch.randn(batch, ngroups, dstate, dtype=dtype, device=device)
-    C = torch.randn(batch, ngroups, dstate, dtype=dtype, device=device)
-    D = torch.randn(nheads, dim, dtype=dtype, device=device)
-    dt_bias = torch.randn(nheads, dim, dtype=dtype, device=device)
-    out = torch.zeros(batch, nheads, dim, dtype=dtype, device=device)
-    return state, x, dt, A, B, C, D, dt_bias, out
-
-
-def benchmark_config(
-    batch: int,
-    nheads: int,
-    dim: int,
-    dstate: int,
-    ngroups: int,
-    block_size_m: int,
-    num_warps_val: int,
-    dtype: torch.dtype,
-    state_dtype: torch.dtype | None = None,
-    num_iters: int = 100,
-    num_warmup: int = 20,
-    graph_batch_size: int = 10,
-) -> float | None:
-    """
-    Time one (BLOCK_SIZE_M, num_warps) config for selective_state_update.
-    Returns elapsed time in microseconds, or None on error.
-
-    Uses CUDA graph capture-and-replay to isolate kernel time from Python
-    eager-mode dispatch / kwarg-resolution overhead, mirroring the timing
-    methodology in benchmarks/kernels/benchmark_moe.py.
-    """
-    state, x, dt, A, B, C, D, dt_bias, out = _make_inputs(
-        batch, nheads, dim, dstate, ngroups, dtype, state_dtype=state_dtype
-    )
-
-    def _call_kernel() -> None:
-        selective_state_update(
-            state,
-            x,
-            dt,
-            A,
-            B,
-            C,
-            D=D,
-            z=None,
-            dt_bias=dt_bias,
-            dt_softplus=True,
-            out=out,
-        )
-
-    try:
-        with override_ssm_config((block_size_m, num_warps_val)):
-            # Eager-mode warmup: triggers Triton autotune / JIT, primes caches.
-            for _ in range(num_warmup):
-                _call_kernel()
-            torch.accelerator.synchronize()
-
-            # Capture graph_batch_size invocations into a CUDA graph so the
-            # timed region runs without Python dispatch overhead per call.
-            graph = torch.cuda.CUDAGraph()
-            with torch.cuda.graph(graph):
-                for _ in range(graph_batch_size):
-                    _call_kernel()
-            torch.accelerator.synchronize()
-
-            # Warmup graph replays (let the runtime stabilize).
-            for _ in range(5):
-                graph.replay()
-            torch.accelerator.synchronize()
-
-            start = torch.cuda.Event(enable_timing=True)
-            end = torch.cuda.Event(enable_timing=True)
-            latencies: list[float] = []
-            for _ in range(num_iters):
-                start.record()
-                graph.replay()
-                end.record()
-                end.synchronize()
-                latencies.append(start.elapsed_time(end))
-            graph.reset()
-        # elapsed_time returns ms; each replay runs graph_batch_size kernels,
-        # so divide by (num_iters * graph_batch_size) and convert ms -> us.
-        return sum(latencies) / (num_iters * graph_batch_size) * 1000
-    except Exception as e:
-        if "OutOfResources" not in str(e):
-            print(
-                f"    Warning: config M={block_size_m},w={num_warps_val} "
-                f"raised {type(e).__name__}: {e}"
-            )
-        return None
-
-
-# ---------------------------------------------------------------------------
-# Tuning loop
-# ---------------------------------------------------------------------------
-
-
-# CUDA grid Y/Z dim limit — both `batch` and `nheads` must fit individually.
-_CUDA_MAX_GRID_DIM = 65535
-
-# Above this, kernel state-offset arithmetic (batch * nheads * headdim * dstate)
-# overflows int32 and the launch raises cudaErrorIllegalAddress.
-# 262144 covers Nemotron Super TP1 BS=2048.
-_MAX_EFFECTIVE_BATCH = 262144
-
-
-def expand_batch_x_nheads(
-    batch_sizes: list[int],
-    nheads_list: list[int],
-    ngroups: int,
-) -> list[tuple[int, int, int]]:
-    """Cross-product batch_sizes × nheads_list → sorted [(effective_batch,
-    batch, nheads)], deduped by effective_batch. Filters pairs that exceed
-    the CUDA grid dim limit, the effective_batch ceiling, or where nheads is
-    not a positive multiple of ngroups.
-    """
-    seen: dict[int, tuple[int, int]] = {}
-    skipped_grid: list[tuple[int, int]] = []
-    skipped_ngroups: list[tuple[int, int]] = []
-    skipped_eb: list[tuple[int, int]] = []
-    for b, n in product(batch_sizes, nheads_list):
-        if b <= 0 or n <= 0:
-            continue
-        if b > _CUDA_MAX_GRID_DIM or n > _CUDA_MAX_GRID_DIM:
-            skipped_grid.append((b, n))
-            continue
-        if n % ngroups != 0:
-            skipped_ngroups.append((b, n))
-            continue
-        if b * n > _MAX_EFFECTIVE_BATCH:
-            skipped_eb.append((b, n))
-            continue
-        seen.setdefault(b * n, (b, n))
-    if skipped_grid:
-        print(
-            f"  Note: skipping (batch, nheads) pairs exceeding CUDA grid dim "
-            f"{_CUDA_MAX_GRID_DIM}: {skipped_grid}"
-        )
-    if skipped_ngroups:
-        print(
-            f"  Note: skipping (batch, nheads) pairs where nheads % ngroups != 0 "
-            f"for ngroups={ngroups}: {skipped_ngroups}"
-        )
-    if skipped_eb:
-        print(
-            f"  Note: skipping (batch, nheads) pairs whose effective_batch "
-            f"exceeds {_MAX_EFFECTIVE_BATCH}: {skipped_eb}"
-        )
-    return sorted((eb, b, n) for eb, (b, n) in seen.items())
-
-
-def tune_dstate(
-    dstate: int,
-    headdim: int,
-    ngroups: int,
-    dtype: torch.dtype,
-    num_iters: int,
-    verbose: bool,
-    active: list[tuple[int, int, int]],
-    state_dtype: torch.dtype | None = None,
-) -> tuple[dict[int, dict], dict[int, dict[tuple[int, int], float]]]:
-    """For each (effective_batch, batch, nheads) in *active*, sweep
-    (BLOCK_SIZE_M, num_warps) and return
-    ({effective_batch: best_config}, {effective_batch: {(bsm, nw): us}}).
-    The second map is the full timing grid, used downstream so we don't
-    re-measure the same config in the comparison phase.
-    """
-    best_per_eb: dict[int, dict] = {}
-    timings: dict[int, dict[tuple[int, int], float]] = {}
-
-    print(f"\n{'=' * 74}")
-    effective_state_dtype = state_dtype if state_dtype is not None else dtype
-    print(
-        f"Tuning  headdim={headdim}  dstate={dstate}  ngroups={ngroups}  "
-        f"dtype={dtype}  ssm_cache_dtype={effective_state_dtype}"
-    )
-    print(f"{'=' * 74}")
-
-    bsm_choices = _block_size_m_choices(headdim)
-    print(f"BSM candidates (capped at next_pow2(headdim={headdim})): {bsm_choices}")
-
-    hdr = f"{'EffBatch':>8} | {'BLOCK_M':>7} | {'warps':>5} | {'us':>10} | note"
-    print(hdr)
-    print("-" * 52)
-
-    for eb, batch, nheads in active:
-        best_time = float("inf")
-        best_cfg: dict = {}
-        eb_timings: dict[tuple[int, int], float] = {}
-
-        for bsm, nw in product(bsm_choices, NUM_WARPS_CHOICES):
-            t = benchmark_config(
-                batch=batch,
-                nheads=nheads,
-                dim=headdim,
-                dstate=dstate,
-                ngroups=ngroups,
-                block_size_m=bsm,
-                num_warps_val=nw,
-                dtype=dtype,
-                state_dtype=state_dtype,
-                num_iters=num_iters,
-            )
-            if t is None:
-                continue
-            eb_timings[(bsm, nw)] = t
-            is_best = t < best_time
-            if is_best:
-                best_time = t
-                best_cfg = {"BLOCK_SIZE_M": bsm, "num_warps": nw}
-            if verbose:
-                marker = " <-- best" if is_best else ""
-                print(f"{eb:>8} | {bsm:>7} | {nw:>5} | {t:>10.2f} |{marker}")
-
-        timings[eb] = eb_timings
-
-        if not best_cfg:
-            print(
-                f"{eb:>8} | {'-':>7} | {'-':>5} | {'-':>10} | "
-                f"no working config (skipped)"
-            )
-            continue
-
-        if not verbose:
-            print(
-                f"{eb:>8} | {best_cfg['BLOCK_SIZE_M']:>7} | "
-                f"{best_cfg['num_warps']:>5} | {best_time:>10.2f} | best"
-            )
-
-        best_per_eb[eb] = best_cfg
-
-    return best_per_eb, timings
-
-
-# ---------------------------------------------------------------------------
-# Correctness validation
-# ---------------------------------------------------------------------------
-
-
-def validate_configs(
-    dstate: int,
-    headdim: int,
-    ngroups: int,
-    tuned: dict[int, dict],
-    active: list[tuple[int, int, int]],
-    dtype: torch.dtype,
-    atol: float = 1e-2,
-    rtol: float = 1e-2,
-    state_dtype: torch.dtype | None = None,
-) -> dict[int, bool]:
-    """
-    For every (effective_batch, batch, nheads) in *active* that has a tuned
-    config, run the kernel with that config and compare against the reference.
-    Returns {effective_batch: passed}.
-    """
-    # Disable TF32 so the reference's matmul matches the Triton kernel's
-    # fp32 accumulation; otherwise large ebs show bf16 rounding mismatches.
-    torch.set_float32_matmul_precision("highest")
-
-    print(f"\n{'=' * 74}")
-    effective_state_dtype = state_dtype if state_dtype is not None else dtype
-    print(
-        f"Validation  headdim={headdim}  dstate={dstate}  ngroups={ngroups}  "
-        f"dtype={dtype}  ssm_cache_dtype={effective_state_dtype}  atol={atol}"
-    )
-    print(f"{'=' * 74}")
-    print(f"{'EffBatch':>8} | {'MaxAbsErr':>12} | {'Status':>8}")
-    print("-" * 36)
-
-    results: dict[int, bool] = {}
-
-    for eb, batch, nheads in active:
-        cfg = tuned.get(eb)
-        if cfg is None:
-            continue
-        state, x, dt, A, B, C, D, dt_bias, out = _make_inputs(
-            batch=batch,
-            nheads=nheads,
-            dim=headdim,
-            dstate=dstate,
-            ngroups=ngroups,
-            dtype=dtype,
-            state_dtype=state_dtype,
-        )
-        # Clone state before GPU kernel modifies it in-place
-        state_ref = state.clone()
-
-        with override_ssm_config((cfg["BLOCK_SIZE_M"], cfg["num_warps"])):
-            selective_state_update(
-                state,
-                x,
-                dt,
-                A,
-                B,
-                C,
-                D=D,
-                z=None,
-                dt_bias=dt_bias,
-                dt_softplus=True,
-                out=out,
-            )
-        torch.accelerator.synchronize()
-        gpu_out = out.detach().cpu()
-
-        # Reference uses the original (unmodified) state
-        # Upcast to fp32 so the reference sums in fp32 (matches the Triton
-        # kernel); summing in bf16 over `dstate` blows up the error.
-        ref_out = (
-            selective_state_update_ref(
-                state_ref.float(),
-                x.float(),
-                dt.float(),
-                A.float(),
-                B.float(),
-                C.float(),
-                D=D.float(),
-                dt_bias=dt_bias.float(),
-                dt_softplus=True,
-            )
-            .to(out.dtype)
-            .cpu()
-        )
-
-        passed = torch.allclose(gpu_out.float(), ref_out.float(), atol=atol, rtol=rtol)
-        max_err = (gpu_out.float() - ref_out.float()).abs().max().item()
-        status = "PASS" if passed else "FAIL"
-        results[eb] = passed
-        print(f"{eb:>8} | {max_err:>12.6f} | {status:>8}")
-
-    n_pass = sum(results.values())
-    n_total = len(results)
-    print(f"\n  {n_pass}/{n_total} configs passed validation for dstate={dstate}")
-    return results
-
-
-# ---------------------------------------------------------------------------
-# Save configs
-# ---------------------------------------------------------------------------
-
-
-def save_configs(
-    headdim: int,
-    dstate: int,
-    cache_dtype: str,
-    configs: dict[int, dict],
-    save_dir: str | None = None,
-) -> str:
-    # bf16 shares configs with fp16, use common filename for both
-    cache_dtype = _canonical_cache_dtype(cache_dtype)
-
-    base_dir = save_dir if save_dir else _CONFIGS_DIR
-    os.makedirs(base_dir, exist_ok=True)
-    file_path = os.path.join(
-        base_dir,
-        get_ssm_config_file_name(headdim, dstate, cache_dtype, get_ssm_device_name()),
-    )
-    # triton_version is informational only, the loader ignores it
-    payload: dict[str, Any] = {
-        "triton_version": triton.__version__,
-        **{str(k): v for k, v in sorted(configs.items())},
-    }
-    with open(file_path, "w") as f:
-        json.dump(payload, f, indent=4)
-    return file_path
-
-
-# ---------------------------------------------------------------------------
-# Comparison table
-# ---------------------------------------------------------------------------
-
-
-def current_heuristic(dstate: int, is_blackwell: bool = False) -> dict:
-    """Return the current hard-coded BLOCK_SIZE_M / num_warps for dstate."""
-    bsm, nw = _get_default_ssm_launch_config(dstate, is_blackwell)
-    return {"BLOCK_SIZE_M": bsm, "num_warps": nw}
-
-
-def compare_heuristic_vs_tuned(
-    dstate: int,
-    headdim: int,
-    ngroups: int,
-    tuned: dict[int, dict],
-    timings: dict[int, dict[tuple[int, int], float]],
-    active: list[tuple[int, int, int]],
-    dtype: torch.dtype,
-    num_iters: int,
-    is_blackwell: bool,
-    state_dtype: torch.dtype | None = None,
-):
-    heur_cfg = current_heuristic(dstate, is_blackwell)
-    heur_key = (heur_cfg["BLOCK_SIZE_M"], heur_cfg["num_warps"])
-
-    print(f"\n{'=' * 74}")
-    print(
-        f"Comparison  headdim={headdim}  dstate={dstate}  "
-        f"ngroups={ngroups}  —  heuristic vs tuned"
-    )
-    print(
-        f"Heuristic: BLOCK_SIZE_M={heur_cfg['BLOCK_SIZE_M']}, "
-        f"num_warps={heur_cfg['num_warps']}"
-    )
-    print(f"{'=' * 74}")
-    hdr = (
-        f"{'EffBatch':>8} | {'Heur(us)':>10} | {'Tuned(us)':>10} | "
-        f"{'Speedup':>8} | Best config"
-    )
-    print(hdr)
-    print("-" * len(hdr))
-
-    for eb, batch, nheads in active:
-        eb_timings = timings.get(eb, {})
-
-        # Heuristic timing: reuse the tuning measurement if the heuristic
-        # config was in the swept grid; otherwise measure it once.
-        t_h = eb_timings.get(heur_key)
-        if t_h is None:
-            t_h = benchmark_config(
-                batch=batch,
-                nheads=nheads,
-                dim=headdim,
-                dstate=dstate,
-                ngroups=ngroups,
-                block_size_m=heur_cfg["BLOCK_SIZE_M"],
-                num_warps_val=heur_cfg["num_warps"],
-                dtype=dtype,
-                state_dtype=state_dtype,
-                num_iters=num_iters,
-            )
-
-        # `tuned[eb]` may be missing if all configs failed in tune_dstate;
-        # in that case fall back to the heuristic so the table still prints.
-        best = tuned.get(eb) or heur_cfg
-        t_t = eb_timings.get((best["BLOCK_SIZE_M"], best["num_warps"]))
-
-        if t_h is None or t_t is None:
-            print(f"{eb:>8} | {'N/A':>10} | {'N/A':>10} | {'N/A':>8} |")
-            continue
-        speedup = t_h / t_t
-        marker = " <--" if speedup > 1.05 else ""
-        print(
-            f"{eb:>8} | {t_h:>10.2f} | {t_t:>10.2f} | "
-            f"{speedup:>7.2f}x | "
-            f"M={best['BLOCK_SIZE_M']},w={best['num_warps']}{marker}"
-        )
-
-
-# ---------------------------------------------------------------------------
-# CLI
-# ---------------------------------------------------------------------------
-
-
-def save_results(device_name: str, output: str, results_file: str | None = None) -> str:
-    """Save the full benchmark output to a results text file."""
-    if results_file is None:
-        results_file = os.path.join(
-            _RESULTS_DIR, f"ssm_benchmark_results_{device_name}.txt"
-        )
-    with open(results_file, "w") as f:
-        f.write(output)
-    return results_file
-
-
-def main():
-    parser = argparse.ArgumentParser(
-        description="Tune selective_state_update kernel for Mamba SSM"
-    )
-    parser.add_argument(
-        "--dstate",
-        type=int,
-        default=128,
-        help="SSM state size to tune for (default: 128)",
-    )
-    parser.add_argument(
-        "--all-dstates",
-        action="store_true",
-        help="Tune all common dstate values: " + str(ALL_DSTATES),
-    )
-    parser.add_argument(
-        "--dtype",
-        type=str,
-        default="bfloat16",
-        choices=["float16", "bfloat16"],
-        help="Activation / input data type (default: bfloat16)",
-    )
-    parser.add_argument(
-        "--mamba-ssm-cache-dtype",
-        type=str,
-        default="float32",
-        choices=list(_SSM_CACHE_DTYPE_MAP.keys()),
-        help="SSM state cache dtype (default: float32)",
-    )
-    parser.add_argument(
-        "--num-iters",
-        type=int,
-        default=100,
-        help="Number of timing iterations (default: 100)",
-    )
-    parser.add_argument(
-        "--save-configs",
-        action="store_true",
-        help=f"Save best configs to JSON in {_CONFIGS_DIR}",
-    )
-    parser.add_argument(
-        "--compare",
-        action="store_true",
-        help="Show comparison table: heuristic vs tuned",
-    )
-    parser.add_argument(
-        "--verbose",
-        action="store_true",
-        help="Print every (BLOCK_SIZE_M, num_warps) result, not just best",
-    )
-    parser.add_argument(
-        "--results-file",
-        type=str,
-        default=None,
-        help="Path to save the benchmark results text file "
-        "(default: ssm_benchmark_results_<device>.txt alongside this script)",
-    )
-    parser.add_argument(
-        "--save-dir",
-        type=str,
-        default=None,
-        help=f"Directory to save JSON configs (default: {_CONFIGS_DIR})",
-    )
-    parser.add_argument(
-        "--headdim",
-        type=int,
-        default=DEFAULT_HEADDIM,
-        help=f"Per-head feature dim (default: {DEFAULT_HEADDIM})",
-    )
-    parser.add_argument(
-        "--ngroups",
-        type=int,
-        default=DEFAULT_NGROUPS,
-        help=f"Number of B/C groups (default: {DEFAULT_NGROUPS})",
-    )
-    parser.add_argument(
-        "--batch-sizes",
-        type=int,
-        nargs="+",
-        default=DEFAULT_BATCH_SIZES,
-        metavar="B",
-        help=f"Decoder batch sizes to sweep (default: {DEFAULT_BATCH_SIZES})",
-    )
-    parser.add_argument(
-        "--nheads",
-        type=int,
-        nargs="+",
-        default=DEFAULT_NHEADS,
-        metavar="N",
-        help=f"Number of heads per rank to sweep (default: {DEFAULT_NHEADS}). "
-        "effective_batch = batch * nheads; cross-product is deduped by eb.",
-    )
-    parser.add_argument(
-        "--validate",
-        action="store_true",
-        help="After tuning, verify each best config against a CPU reference "
-        "implementation. Configs that fail are flagged in the output.",
-    )
-    parser.add_argument(
-        "--atol",
-        type=float,
-        default=1e-2,
-        help="Absolute tolerance for --validate (default: 1e-2)",
-    )
-    args = parser.parse_args()
-
-    dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
-    state_dtype = _SSM_CACHE_DTYPE_MAP[args.mamba_ssm_cache_dtype]
-    device_name = get_ssm_device_name()
-    cap = torch.cuda.get_device_capability()
-    is_blackwell = cap[0] >= 10
-
-    # Mirror all output to a results file (like Unix tee).
-    buf = StringIO()
-
-    class _Tee:
-        """Writes to both the original stdout and an in-memory buffer."""
-
-        def write(self, s):
-            buf.write(s)
-            sys.__stdout__.write(s)
-
-        def flush(self):
-            sys.__stdout__.flush()
-
-    sys.stdout = _Tee()  # type: ignore[assignment]
-
-    try:
-        print(f"Device : {device_name}  (sm_{cap[0]}{cap[1]})")
-        print(f"Blackwell: {is_blackwell}")
-        print(f"dtype  : {args.dtype}")
-        print(f"ssm_cache_dtype: {args.mamba_ssm_cache_dtype}")
-        print(f"headdim: {args.headdim}")
-        print(f"ngroups: {args.ngroups}")
-        print(f"triton : {triton.__version__}")
-
-        dstates = ALL_DSTATES if args.all_dstates else [args.dstate]
-        active = expand_batch_x_nheads(args.batch_sizes, args.nheads, args.ngroups)
-
-        for dstate in dstates:
-            tuned, timings = tune_dstate(
-                dstate=dstate,
-                headdim=args.headdim,
-                ngroups=args.ngroups,
-                dtype=dtype,
-                num_iters=args.num_iters,
-                verbose=args.verbose,
-                active=active,
-                state_dtype=state_dtype,
-            )
-
-            if args.compare:
-                compare_heuristic_vs_tuned(
-                    dstate=dstate,
-                    headdim=args.headdim,
-                    ngroups=args.ngroups,
-                    tuned=tuned,
-                    timings=timings,
-                    active=active,
-                    dtype=dtype,
-                    num_iters=args.num_iters,
-                    is_blackwell=is_blackwell,
-                    state_dtype=state_dtype,
-                )
-
-            if args.validate:
-                validity = validate_configs(
-                    dstate=dstate,
-                    headdim=args.headdim,
-                    ngroups=args.ngroups,
-                    tuned=tuned,
-                    active=active,
-                    dtype=dtype,
-                    atol=args.atol,
-                    state_dtype=state_dtype,
-                )
-                # Filter out any configs that failed correctness check
-                failed = [eb for eb, ok in validity.items() if not ok]
-                if failed:
-                    print(
-                        f"\n  WARNING: {len(failed)} config(s) failed validation "
-                        f"for dstate={dstate}: effective_batches {failed}"
-                    )
-                    print("  These will NOT be saved even with --save-configs.")
-                    tuned = {
-                        eb: cfg for eb, cfg in tuned.items() if validity.get(eb, True)
-                    }
-
-            if args.save_configs:
-                path = save_configs(
-                    headdim=args.headdim,
-                    dstate=dstate,
-                    cache_dtype=args.mamba_ssm_cache_dtype,
-                    configs=tuned,
-                    save_dir=args.save_dir,
-                )
-                print(f"\nSaved: {path}")
-            else:
-                print(f"\nBest configs for dstate={dstate}:")
-                for eb, cfg in sorted(tuned.items()):
-                    print(f"  effective_batch={eb:>6}: {cfg}")
-                print("\n(Re-run with --save-configs to persist to JSON)")
-    finally:
-        sys.stdout = sys.__stdout__
-        results_path = save_results(device_name, buf.getvalue(), args.results_file)
-        print(f"\nResults saved to: {results_path}")
-
-
-if __name__ == "__main__":
-    main()
@@ -12,6 +12,7 @@ from vllm._custom_ops import (
    cpu_attn_get_scheduler_metadata,
    cpu_attn_reshape_and_cache,
 )
+from vllm.platforms import CpuArchEnum, current_platform
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, set_random_seed
 from vllm.v1.attention.backends.cpu_attn import CPUAttentionBackend, _get_attn_isa
@@ -21,14 +22,15 @@ def get_attn_isa(
    block_size: int | None = None,
    dtype: torch.dtype | None = None,
 ):
-    # Delegate to _get_attn_isa so the fallback path applies the same arch
-    # gating (e.g. RISC-V RVV is only chosen when the build's hardcoded
-    # VLEN=128 kernel is actually present; on VLEN=256 / scalar hosts it
-    # correctly falls through to vec/vec16).
-    return _get_attn_isa(
-        dtype if dtype is not None else torch.bfloat16,
-        block_size if block_size else 32,
-    )
+    if block_size and dtype:
+        return _get_attn_isa(dtype, block_size)
+    else:
+        if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
+            return "neon"
+        elif torch.cpu._is_amx_tile_supported():
+            return "amx"
+        else:
+            return "vec"


 # rand number generation takes too much time, cache rand tensors
@@ -233,7 +235,7 @@ if __name__ == "__main__":
    )
    parser.add_argument("--use-sink", action="store_true")
    parser.add_argument(
-        "--isa", type=str, choices=["vec", "neon", "amx", "vec16", "rvv"], default=None
+        "--isa", type=str, choices=["vec", "neon", "amx", "vec16"], default=None
    )
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--iters", type=int, default=20)
@@ -1,44 +0,0 @@
-#!/bin/bash
-# Build the vllm-rs Rust frontend binary and install it into the vllm package.
-# Usage: ./build_rust.sh [--debug]
-#
-# By default builds in release mode. Pass --debug for faster compile times
-# during development.
-
-set -euo pipefail
-
-REPO_ROOT="$(cd "$(dirname "$0")" && pwd)"
-RUST_DIR="$REPO_ROOT/rust"
-TARGET_PATH="${VLLM_RS_TARGET_PATH:-$REPO_ROOT/vllm/vllm-rs}"
-
-# Read the required toolchain from rust-toolchain.toml.
-TOOLCHAIN=$(grep '^channel' "$REPO_ROOT/rust-toolchain.toml" | sed 's/.*= *"\(.*\)"/\1/')
-
-# Ensure rustup and the required toolchain are available.
-if ! command -v rustup &>/dev/null; then
-    echo "rustup not found, installing..."
-    curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain none
-    source "$HOME/.cargo/env"
-fi
-
-if ! rustup run "$TOOLCHAIN" rustc --version &>/dev/null; then
-    echo "Installing Rust toolchain: $TOOLCHAIN"
-    rustup toolchain install "$TOOLCHAIN"
-fi
-
-if [[ "${1:-}" == "--debug" ]]; then
-    PROFILE_ARGS=()
-    PROFILE_DIR="debug"
-else
-    PROFILE_ARGS=(--release)
-    PROFILE_DIR="release"
-fi
-
-cargo +"$TOOLCHAIN" build "${PROFILE_ARGS[@]}" \
-    --manifest-path "$RUST_DIR/Cargo.toml" \
-    --bin vllm-rs \
-    --features native-tls-vendored
-
-mkdir -p "$(dirname "$TARGET_PATH")"
-cp "$RUST_DIR/target/$PROFILE_DIR/vllm-rs" "$TARGET_PATH"
-echo "Installed vllm-rs to $TARGET_PATH"
@@ -1,7 +1,7 @@
 include(FetchContent)

 set(CMAKE_CXX_STANDARD_REQUIRED ON)
-set(CMAKE_CXX_STANDARD 20)
+set(CMAKE_CXX_STANDARD 17)
 set(CMAKE_CXX_EXTENSIONS ON)
 set(CMAKE_EXPORT_COMPILE_COMMANDS ON)

@@ -32,23 +32,18 @@ else()
        "-DVLLM_CPU_EXTENSION")

    # locate PyTorch's libgomp (e.g. site-packages/torch.libs/libgomp-947d5fa1.so.1.0.0)
-    # and create a local shim dir with it. When PyTorch is built from source or packaged
-    # by a distro (common on RISC-V, s390x, Fedora/RHEL aarch64), no vendored libgomp
-    # exists and the shim dir is empty; fall back to the system libgomp in that case.
+    # and create a local shim dir with it
    vllm_prepare_torch_gomp_shim(VLLM_TORCH_GOMP_SHIM_DIR)

-    if(VLLM_TORCH_GOMP_SHIM_DIR)
-        find_library(OPEN_MP
-            NAMES gomp
-            PATHS "${VLLM_TORCH_GOMP_SHIM_DIR}"
-            NO_DEFAULT_PATH
-            REQUIRED
-        )
-        # Use the same libgomp as PyTorch at runtime
+    find_library(OPEN_MP
+        NAMES gomp
+        PATHS ${VLLM_TORCH_GOMP_SHIM_DIR}
+        NO_DEFAULT_PATH
+        REQUIRED
+    )
+    # Set LD_LIBRARY_PATH to include the shim dir at build time to use the same libgomp as PyTorch
+    if (OPEN_MP)
        set(ENV{LD_LIBRARY_PATH} "${VLLM_TORCH_GOMP_SHIM_DIR}:$ENV{LD_LIBRARY_PATH}")
-    else()
-        # Fall back to system / toolchain libgomp
-        find_library(OPEN_MP NAMES gomp REQUIRED)
    endif()
 endif()

@@ -195,12 +190,10 @@ elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
    endif()
    if(VLLM_RVV_VLEN AND VLLM_RVV_VLEN GREATER 0)
        message(STATUS "RISC-V RVV VLEN=${VLLM_RVV_VLEN}")
-        # Sources gate FP16/BF16 paths on the compiler-provided
-        # __riscv_zvfh / __riscv_zvfbfmin macros, which GCC and clang
-        # define automatically when those extensions appear in -march.
        if(RVV_BF16_FOUND)
            message(STATUS "BF16 extension detected")
            set(MARCH_FLAGS -march=rv64gcv_zvfh_zfbfmin_zvfbfmin_zvl${VLLM_RVV_VLEN}b -mrvv-vector-bits=zvl -mabi=lp64d)
+            add_compile_definitions(RISCV_BF16_SUPPORT)
        elseif(RVV_FP16_FOUND)
            message(WARNING "BF16 functionality is not available")
            set(MARCH_FLAGS -march=rv64gcv_zvfh_zvl${VLLM_RVV_VLEN}b -mrvv-vector-bits=zvl -mabi=lp64d)
@@ -328,6 +321,14 @@ if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND
    set(ONEDNN_VERBOSE "ON")
    set(CMAKE_POLICY_DEFAULT_CMP0077 NEW)

+    # TODO: Refactor this
+    if (ENABLE_X86_ISA)
+        # Note: only enable oneDNN for AVX512
+        list(APPEND DNNL_COMPILE_FLAGS ${CXX_COMPILE_FLAGS_AVX512})
+    else()
+        list(APPEND DNNL_COMPILE_FLAGS ${CXX_COMPILE_FLAGS})
+    endif()
+
    set(VLLM_BUILD_TYPE ${CMAKE_BUILD_TYPE})
    set(CMAKE_BUILD_TYPE "Release") # remove oneDNN debug symbols to reduce size
    FetchContent_MakeAvailable(oneDNN)
@@ -340,14 +341,8 @@ if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND
        PRIVATE ${oneDNN_SOURCE_DIR}/src
    )
    target_link_libraries(dnnl_ext dnnl torch)
-    if (ENABLE_X86_ISA)
-        target_compile_options(dnnl_ext PRIVATE ${CXX_COMPILE_FLAGS_AVX2} -fPIC)
-    else()
-        target_compile_options(dnnl_ext PRIVATE ${CXX_COMPILE_FLAGS} -fPIC)
-    endif()
+    target_compile_options(dnnl_ext PRIVATE ${DNNL_COMPILE_FLAGS} -fPIC)
    list(APPEND LIBS dnnl_ext)
-
-
    set(USE_ONEDNN ON)
 else()
    set(USE_ONEDNN OFF)
@@ -369,18 +364,6 @@ else()
    add_compile_definitions(-DVLLM_NUMA_DISABLED)
 endif()

-# check if the pytorch wheel ships libopenblas.so.
-set(VLLM_OPENBLAS_LIB "")
-if (NOT ENABLE_X86_ISA)
-    file(GLOB _VLLM_TORCH_OPENBLAS_LIBS
-        "${TORCH_INSTALL_PREFIX}/lib/libopenblas*.so*")
-    # Note: we don't link openblas directly to _C extension, as it's available through libtorch.so 
-    if (_VLLM_TORCH_OPENBLAS_LIBS)
-        list(GET _VLLM_TORCH_OPENBLAS_LIBS 0 VLLM_OPENBLAS_LIB)
-        message(STATUS "CPU OpenBLAS library: ${VLLM_OPENBLAS_LIB}")
-    endif()
-endif()
-
 #
 # Generate CPU attention dispatch header
 #
@@ -399,7 +382,6 @@ endif()
 #
 set(VLLM_EXT_SRC
    "csrc/cpu/activation.cpp"
-    "csrc/cpu/sgl-kernels/fla.cpp"
    "csrc/cpu/utils.cpp"
    "csrc/cpu/spec_decode_utils.cpp"
    "csrc/cpu/layernorm.cpp"
@@ -409,13 +391,6 @@ set(VLLM_EXT_SRC
    "csrc/cpu/cpu_attn.cpp"
    "csrc/cpu/torch_bindings.cpp")

-if (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64" AND VLLM_RVV_VLEN AND
-        VLLM_RVV_VLEN GREATER 0 AND (RVV_FP16_FOUND OR RVV_BF16_FOUND))
-    set(VLLM_EXT_SRC
-        "csrc/cpu/cpu_wna16.cpp"
-        ${VLLM_EXT_SRC})
-endif()
-
 if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
    set(VLLM_EXT_SRC
        "csrc/cpu/shm.cpp"
@@ -423,12 +398,6 @@ if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
        ${VLLM_EXT_SRC})
 endif()

-if (POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND)	
-    set(VLLM_EXT_SRC
-        "csrc/cpu/shm.cpp"
-        ${VLLM_EXT_SRC})
-endif()
-
 if(USE_ONEDNN)
    set(VLLM_EXT_SRC
        "csrc/cpu/dnnl_kernels.cpp"
@@ -437,18 +406,15 @@ endif()

 if (ENABLE_X86_ISA)
    set(VLLM_EXT_SRC_SGL
-        "csrc/cpu/sgl-kernels/conv.cpp"
        "csrc/cpu/sgl-kernels/gemm.cpp"
        "csrc/cpu/sgl-kernels/gemm_int8.cpp"
        "csrc/cpu/sgl-kernels/gemm_fp8.cpp"
        "csrc/cpu/sgl-kernels/gemm_int4.cpp"
        "csrc/cpu/sgl-kernels/moe.cpp"
        "csrc/cpu/sgl-kernels/moe_int8.cpp"
-        "csrc/cpu/sgl-kernels/moe_int4.cpp"
        "csrc/cpu/sgl-kernels/moe_fp8.cpp")

    set(VLLM_EXT_SRC_AVX512
-        "csrc/cpu/sgl-kernels/fla.cpp"
        "csrc/cpu/shm.cpp"
        "csrc/cpu/cpu_wna16.cpp"
        "csrc/cpu/cpu_fused_moe.cpp"
@@ -464,12 +430,10 @@ if (ENABLE_X86_ISA)
        "csrc/cpu/pos_encoding.cpp"
        "csrc/moe/dynamic_4bit_int_moe_cpu.cpp") 

-    set(VLLM_EXT_SRC_AVX2
-        "csrc/cpu/sgl-kernels/fla.cpp"
+    set(VLLM_EXT_SRC_AVX2 
        "csrc/cpu/utils.cpp"
        "csrc/cpu/spec_decode_utils.cpp"
        "csrc/cpu/cpu_attn.cpp"
-        "csrc/cpu/dnnl_kernels.cpp"
        "csrc/cpu/torch_bindings.cpp"
        # TODO: Remove these files
        "csrc/cpu/activation.cpp"
@@ -484,7 +448,7 @@ if (ENABLE_X86_ISA)

    set(_C_LIBS numa dnnl_ext)
    set(_C_AVX512_LIBS numa dnnl_ext)
-    set(_C_AVX2_LIBS numa dnnl_ext)
+    set(_C_AVX2_LIBS numa)

    # AMX + AVX512F + AVX512BF16 + AVX512VNNI
    define_extension_target(
@@ -539,9 +503,6 @@ else()
        USE_SABI 3
        WITH_SOABI
    )
-    if (VLLM_OPENBLAS_LIB)
-        target_compile_definitions(_C PRIVATE VLLM_HAS_OPENBLAS)
-    endif()
 endif()

 message(STATUS "Enabling C extension.")
@@ -53,80 +53,49 @@ cuda_archs_loose_intersection(DEEPGEMM_ARCHS
 if(DEEPGEMM_ARCHS)
  message(STATUS "DeepGEMM CUDA architectures: ${DEEPGEMM_ARCHS}")

-  #
-  # DeepGEMM integration notes
-  # --------------------------
-  # We vendor DeepGEMM into vllm/third_party/deep_gemm/ and bundle a
-  # `_C.cpython-X.Y-*.so` for every CPython in `requires-python`. The
-  # per-Python build is delegated to tools/build_deepgemm_C.py.
-  #
-  # Why per-Python: DeepGEMM's binding uses PYBIND11_MODULE, which links
-  # private CPython symbols — a single `_C.abi3.so` is not viable today
-  # (see #41476 / #41512 for the failed attempt).
-  #
-  # TODOs (tracked in vllm-project/vllm#42431):
-  #   - Replace DeepGEMM's pybind11 binding with a TORCH_LIBRARY + shim
-  #     binding (cf. vllm-flash-attention/csrc/common/pytorch_shim.h) to
-  #     collapse to one `_C.abi3.so`. Needs either an upstream change or
-  #     a maintained binding fork in vLLM.
-  #   - AOT-compile DeepGEMM's CUDA kernels instead of runtime JIT to drop
-  #     the vendored CUTLASS/CCCL headers and the CUDA-toolkit-at-runtime
-  #     requirement.
-  #
+  find_package(CUDAToolkit REQUIRED)

-  # DEEPGEMM_PYTHON_INTERPRETERS: ":"-separated target Python paths.
-  # Empty/unset → fall back to the build interpreter (editable installs).
-  # (Empty-but-set env vars test as DEFINED in cmake — treat as unset.)
-  if(NOT "$ENV{DEEPGEMM_PYTHON_INTERPRETERS}" STREQUAL "")
-    string(REPLACE ":" ";" _dg_pythons "$ENV{DEEPGEMM_PYTHON_INTERPRETERS}")
-  else()
-    set(_dg_pythons "${Python_EXECUTABLE}")
-  endif()
-  message(STATUS "DeepGEMM _C will be built for: ${_dg_pythons}")
+  #
+  # Build the _C pybind11 extension from DeepGEMM's C++ source.
+  # This is a CXX-only module — CUDA kernels are JIT-compiled at runtime.
+  #
+  Python_add_library(_deep_gemm_C MODULE WITH_SOABI
+    "${deepgemm_SOURCE_DIR}/csrc/python_api.cpp")

-  # add_custom_command does no implicit header scanning; glob explicitly so
-  # header-only edits in DeepGEMM/cutlass/fmt re-trigger the rebuild.
-  file(GLOB_RECURSE _dg_headers
-    "${deepgemm_SOURCE_DIR}/csrc/*.h"
-    "${deepgemm_SOURCE_DIR}/csrc/*.hpp"
-    "${deepgemm_SOURCE_DIR}/deep_gemm/include/*.h"
-    "${deepgemm_SOURCE_DIR}/deep_gemm/include/*.hpp"
-    "${deepgemm_SOURCE_DIR}/deep_gemm/include/*.cuh")
+  # The pybind11 module name must be _C to match DeepGEMM's Python imports.
+  set_target_properties(_deep_gemm_C PROPERTIES OUTPUT_NAME "_C")

-  set(_dg_markers)
-  set(_dg_seen_soabis)
-  foreach(_pybin IN LISTS _dg_pythons)
-    execute_process(
-      COMMAND "${_pybin}" -c
-        "import sysconfig; print(sysconfig.get_config_var('SOABI'))"
-      OUTPUT_VARIABLE _dg_soabi
-      OUTPUT_STRIP_TRAILING_WHITESPACE
-      COMMAND_ERROR_IS_FATAL ANY)
-    # Dedup interpreters that resolve to the same CPython.
-    if(_dg_soabi IN_LIST _dg_seen_soabis)
-      continue()
-    endif()
-    list(APPEND _dg_seen_soabis "${_dg_soabi}")
-    set(_dg_dir "${CMAKE_CURRENT_BINARY_DIR}/deepgemm_C_${_dg_soabi}")
-    set(_dg_marker "${_dg_dir}/.built")
-    add_custom_command(
-      OUTPUT "${_dg_marker}"
-      COMMAND "${Python_EXECUTABLE}"
-              "${CMAKE_SOURCE_DIR}/tools/build_deepgemm_C.py"
-              "${deepgemm_SOURCE_DIR}" "${_dg_dir}" "${_pybin}"
-      COMMAND "${CMAKE_COMMAND}" -E touch "${_dg_marker}"
-      DEPENDS "${CMAKE_SOURCE_DIR}/tools/build_deepgemm_C.py"
-              "${deepgemm_SOURCE_DIR}/csrc/python_api.cpp"
-              ${_dg_headers}
-      COMMENT "Building DeepGEMM _C for ${_pybin}"
-      VERBATIM)
-    list(APPEND _dg_markers "${_dg_marker}")
-    install(DIRECTORY "${_dg_dir}/"
-      DESTINATION vllm/third_party/deep_gemm
-      COMPONENT _deep_gemm_C
-      FILES_MATCHING PATTERN "_C.cpython-*.so")
-  endforeach()
-  add_custom_target(_deep_gemm_C ALL DEPENDS ${_dg_markers})
+  target_compile_definitions(_deep_gemm_C PRIVATE
+    "-DTORCH_EXTENSION_NAME=_C")
+
+  target_include_directories(_deep_gemm_C PRIVATE
+    "${deepgemm_SOURCE_DIR}/csrc"
+    "${deepgemm_SOURCE_DIR}/deep_gemm/include"
+    "${deepgemm_SOURCE_DIR}/third-party/cutlass/include"
+    "${deepgemm_SOURCE_DIR}/third-party/cutlass/tools/util/include"
+    "${deepgemm_SOURCE_DIR}/third-party/fmt/include")
+
+  target_compile_options(_deep_gemm_C PRIVATE
+    $<$<COMPILE_LANGUAGE:CXX>:-std=c++17>
+    $<$<COMPILE_LANGUAGE:CXX>:-O3>
+    $<$<COMPILE_LANGUAGE:CXX>:-Wno-psabi>
+    $<$<COMPILE_LANGUAGE:CXX>:-Wno-deprecated-declarations>)
+
+  # torch_python is required because DeepGEMM uses pybind11 type casters
+  # for at::Tensor (via PYBIND11_MODULE), unlike vLLM's own extensions which
+  # use torch::Library custom ops.
+  find_library(TORCH_PYTHON_LIBRARY torch_python
+    PATHS "${TORCH_INSTALL_PREFIX}/lib"
+    REQUIRED)
+
+  target_link_libraries(_deep_gemm_C PRIVATE
+    torch ${TORCH_LIBRARIES} "${TORCH_PYTHON_LIBRARY}"
+    CUDA::cudart CUDA::nvrtc)
+
+  # Install the shared library into the vendored package directory
+  install(TARGETS _deep_gemm_C
+    LIBRARY DESTINATION vllm/third_party/deep_gemm
+    COMPONENT _deep_gemm_C)

  #
  # Vendor DeepGEMM Python package files
@@ -1,6 +1,6 @@
 # Install OpenAI triton_kernels from https://github.com/triton-lang/triton/tree/main/python/triton_kernels

-set(DEFAULT_TRITON_KERNELS_TAG "v3.5.1")
+set(DEFAULT_TRITON_KERNELS_TAG "v3.6.0")

 # Set TRITON_KERNELS_SRC_DIR for use with local development with vLLM. We expect TRITON_KERNELS_SRC_DIR to
 # be directly set to the triton_kernels python directory.
@@ -31,7 +31,7 @@ endif()

 if(VLLM_FLASH_ATTN_SRC_DIR)
  FetchContent_Declare(
-          vllm-flash-attn SOURCE_DIR
+          vllm-flash-attn SOURCE_DIR 
          ${VLLM_FLASH_ATTN_SRC_DIR}
          BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
  )
@@ -39,7 +39,7 @@ else()
  FetchContent_Declare(
          vllm-flash-attn
          GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
-          GIT_TAG dd62dac706b1cf7895bd99b18c6cb7e7e117ee25
+          GIT_TAG f5bc33cfc02c744d24a2e9d50e6db656de40611c
          GIT_PROGRESS TRUE
          # Don't share the vllm-flash-attn build between build types
          BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
@@ -47,17 +47,12 @@ macro (append_cmake_prefix_path PKG EXPR)
  list(APPEND CMAKE_PREFIX_PATH ${_PREFIX_PATH})
 endmacro()

-# Resolve hipified output paths for `NAME` into `OUT_SRCS` and register the
-# `.cu` sources with the shared `hipify_all` target. Per-extension hipify
-# targets are unsafe to run in parallel against a shared csrc/ output dir, so
-# accumulation here is paired with a single finalize step.
+#
+# Add a target named `hipify${NAME}` that runs the hipify preprocessor on a set
+# of CUDA source files. The names of the corresponding "hipified" sources are
+# stored in `OUT_SRCS`.
+#
 function (hipify_sources_target OUT_SRCS NAME ORIG_SRCS)
-  if (TARGET hipify_all)
-    message(FATAL_ERROR
-      "hipify_sources_target(${NAME}) called after vllm_finalize_hipify_target. "
-      "Add the new HIP extension before the finalizer call in CMakeLists.txt.")
-  endif()
-
  #
  # Split into C++ and non-C++ (i.e. CUDA) sources.
  #
@@ -78,41 +73,19 @@ function (hipify_sources_target OUT_SRCS NAME ORIG_SRCS)
    list(APPEND HIP_SRCS "${CMAKE_CURRENT_BINARY_DIR}/${SRC}")
  endforeach()

-  set_property(GLOBAL APPEND PROPERTY VLLM_HIPIFY_ALL_SRCS ${SRCS})
-  set_property(GLOBAL APPEND PROPERTY VLLM_HIPIFY_ALL_BYPRODUCTS ${HIP_SRCS})
+  set(CSRC_BUILD_DIR ${CMAKE_CURRENT_BINARY_DIR}/csrc)
+  add_custom_target(
+    hipify${NAME}
+    COMMAND ${Python_EXECUTABLE} ${CMAKE_SOURCE_DIR}/cmake/hipify.py -p ${CMAKE_SOURCE_DIR}/csrc -o ${CSRC_BUILD_DIR} ${SRCS}
+    DEPENDS ${CMAKE_SOURCE_DIR}/cmake/hipify.py ${SRCS}
+    BYPRODUCTS ${HIP_SRCS}
+    COMMENT "Running hipify on ${NAME} extension source files.")

  # Swap out original extension sources with hipified sources.
  list(APPEND HIP_SRCS ${CXX_SRCS})
  set(${OUT_SRCS} ${HIP_SRCS} PARENT_SCOPE)
 endfunction()

-# Define the single shared `hipify_all` custom target that runs hipify once
-# on the union of every HIP extension's sources. Call after the last HIP
-# `define_extension_target`.
-function (vllm_finalize_hipify_target)
-  if (TARGET hipify_all)
-    return()
-  endif()
-
-  get_property(ALL_SRCS GLOBAL PROPERTY VLLM_HIPIFY_ALL_SRCS)
-  get_property(ALL_BYPRODUCTS GLOBAL PROPERTY VLLM_HIPIFY_ALL_BYPRODUCTS)
-
-  if (NOT ALL_SRCS)
-    return()
-  endif()
-
-  list(REMOVE_DUPLICATES ALL_SRCS)
-  list(REMOVE_DUPLICATES ALL_BYPRODUCTS)
-
-  set(CSRC_BUILD_DIR ${CMAKE_CURRENT_BINARY_DIR}/csrc)
-  add_custom_target(
-    hipify_all
-    COMMAND ${Python_EXECUTABLE} ${CMAKE_SOURCE_DIR}/cmake/hipify.py -p ${CMAKE_SOURCE_DIR}/csrc -o ${CSRC_BUILD_DIR} ${ALL_SRCS}
-    DEPENDS ${CMAKE_SOURCE_DIR}/cmake/hipify.py ${ALL_SRCS}
-    BYPRODUCTS ${ALL_BYPRODUCTS}
-    COMMENT "Running hipify on all extension source files.")
-endfunction()
-
 #
 # Get additional GPU compiler flags from torch.
 #
@@ -476,16 +449,6 @@ function(cuda_archs_loose_intersection OUT_CUDA_ARCHS SRC_CUDA_ARCHS TGT_CUDA_AR
  set(${OUT_CUDA_ARCHS} ${_CUDA_ARCHS} PARENT_SCOPE)
 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;11.0f" "${TGT_CUDA_ARCHS}")
-  else()
-    cuda_archs_loose_intersection(_archs "9.0a;10.0a;10.1a;10.3a" "${TGT_CUDA_ARCHS}")
-  endif()
-  set(${OUT_CUDA_ARCHS} ${_archs} PARENT_SCOPE)
-endfunction()
-
 #
 # Override the GPU architectures detected by cmake/torch and filter them by
 # `GPU_SUPPORTED_ARCHES`. Sets the final set of architectures in
@@ -588,7 +551,7 @@ function (define_extension_target MOD_NAME)

  if (ARG_LANGUAGE STREQUAL "HIP")
    # Make this target dependent on the hipify preprocessor step.
-    add_dependencies(${MOD_NAME} hipify_all)
+    add_dependencies(${MOD_NAME} hipify${MOD_NAME})
    # Make sure we include the hipified versions of the headers, and avoid conflicts with the ones in the original source folder
    target_include_directories(${MOD_NAME} PRIVATE ${CMAKE_CURRENT_BINARY_DIR}/csrc
      ${ARG_INCLUDE_DIRECTORIES})
@@ -1,12 +1,12 @@
-#include <cuda.h>
-#include <torch/csrc/stable/tensor.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <torch/all.h>
+#include <c10/cuda/CUDAGuard.h>

 #include <cmath>

-#include "../cuda_compat.h"
+#include "cuda_compat.h"
 #include "cuda_vec_utils.cuh"
 #include "dispatch_utils.h"
-#include "torch_utils.h"

 namespace vllm {

@@ -210,68 +210,64 @@ packed_gelu_tanh_kernel(const packed_t& val) {
    return;                                                                    \
  }                                                                            \
  dim3 grid(num_tokens);                                                       \
-  int cc_major = get_device_prop()->major;                                     \
+  int cc_major = at::cuda::getCurrentDeviceProperties()->major;                \
  int support_vec =                                                            \
      (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128)            \
          ? vllm::VecTraits<true>::ARCH_MAX_VEC_SIZE                           \
          : vllm::VecTraits<false>::ARCH_MAX_VEC_SIZE;                         \
-  int vec_size = support_vec / input.element_size();                           \
+  int vec_size = support_vec / at::elementSize(dtype);                         \
  const bool use_vec = (d % vec_size == 0);                                    \
-  const torch::stable::accelerator::DeviceGuard device_guard(                  \
-      input.get_device_index());                                               \
-  const cudaStream_t stream = get_current_cuda_stream();                       \
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));            \
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();                \
  if (use_vec) {                                                               \
    dim3 block(std::min(d / vec_size, 1024));                                  \
    if (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) {         \
-      VLLM_STABLE_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] {   \
+      VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] {          \
        vllm::act_and_mul_kernel<                                              \
            scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,      \
            KERNEL<scalar_t>,                                                  \
            PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
            ACT_FIRST, true, HAS_CLAMP, true><<<grid, block, 0, stream>>>(     \
-            out.mutable_data_ptr<scalar_t>(),                                  \
-            input.const_data_ptr<scalar_t>(), d, LIMIT);                       \
+            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT);   \
      });                                                                      \
    } else {                                                                   \
-      VLLM_STABLE_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] {   \
+      VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] {          \
        vllm::act_and_mul_kernel<                                              \
            scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,      \
            KERNEL<scalar_t>,                                                  \
            PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
            ACT_FIRST, true, HAS_CLAMP, false><<<grid, block, 0, stream>>>(    \
-            out.mutable_data_ptr<scalar_t>(),                                  \
-            input.const_data_ptr<scalar_t>(), d, LIMIT);                       \
+            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT);   \
      });                                                                      \
    }                                                                          \
  } else {                                                                     \
    dim3 block(std::min(d, 1024));                                             \
-    VLLM_STABLE_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] {     \
+    VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] {            \
      vllm::act_and_mul_kernel<                                                \
          scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,        \
          KERNEL<scalar_t>,                                                    \
          PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>,   \
          ACT_FIRST, false, HAS_CLAMP><<<grid, block, 0, stream>>>(            \
-          out.mutable_data_ptr<scalar_t>(), input.const_data_ptr<scalar_t>(),  \
-          d, LIMIT);                                                           \
+          out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT);     \
    });                                                                        \
  }

-void silu_and_mul(torch::stable::Tensor& out,    // [..., d]
-                  torch::stable::Tensor& input)  // [..., 2 * d]
+void silu_and_mul(torch::Tensor& out,    // [..., d]
+                  torch::Tensor& input)  // [..., 2 * d]
 {
  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
                                true, false, 0.0f);
 }

-void silu_and_mul_clamp(torch::stable::Tensor& out,    // [..., d]
-                        torch::stable::Tensor& input,  // [..., 2 * d]
+void silu_and_mul_clamp(torch::Tensor& out,    // [..., d]
+                        torch::Tensor& input,  // [..., 2 * d]
                        double limit) {
  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
                                true, true, (float)limit);
 }

-void mul_and_silu(torch::stable::Tensor& out,    // [..., d]
-                  torch::stable::Tensor& input)  // [..., 2 * d]
+void mul_and_silu(torch::Tensor& out,    // [..., d]
+                  torch::Tensor& input)  // [..., 2 * d]
 {
  // The difference between mul_and_silu and silu_and_mul is that mul_and_silu
  // applies the silu to the latter half of the input.
@@ -279,15 +275,15 @@ void mul_and_silu(torch::stable::Tensor& out,    // [..., d]
                                false, false, 0.0f);
 }

-void gelu_and_mul(torch::stable::Tensor& out,    // [..., d]
-                  torch::stable::Tensor& input)  // [..., 2 * d]
+void gelu_and_mul(torch::Tensor& out,    // [..., d]
+                  torch::Tensor& input)  // [..., 2 * d]
 {
  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_kernel, vllm::packed_gelu_kernel,
                                true, false, 0.0f);
 }

-void gelu_tanh_and_mul(torch::stable::Tensor& out,    // [..., d]
-                       torch::stable::Tensor& input)  // [..., 2 * d]
+void gelu_tanh_and_mul(torch::Tensor& out,    // [..., d]
+                       torch::Tensor& input)  // [..., 2 * d]
 {
  LAUNCH_ACTIVATION_GATE_KERNEL(
      vllm::gelu_tanh_kernel, vllm::packed_gelu_tanh_kernel, true, false, 0.0f);
@@ -438,20 +434,19 @@ __global__ void swigluoai_and_mul_kernel(
    return;                                                                    \
  }                                                                            \
  dim3 grid(num_tokens);                                                       \
-  int cc_major = get_device_prop()->major;                                     \
+  int cc_major = at::cuda::getCurrentDeviceProperties()->major;                \
  int support_vec =                                                            \
      (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128)            \
          ? vllm::VecTraits<true>::ARCH_MAX_VEC_SIZE                           \
          : vllm::VecTraits<false>::ARCH_MAX_VEC_SIZE;                         \
-  int vec_size = support_vec / input.element_size();                           \
+  int vec_size = support_vec / at::elementSize(dtype);                         \
  const bool use_vec = (d % vec_size == 0);                                    \
-  const torch::stable::accelerator::DeviceGuard device_guard(                  \
-      input.get_device_index());                                               \
-  const cudaStream_t stream = get_current_cuda_stream();                       \
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));            \
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();                \
  if (use_vec) {                                                               \
    dim3 block(std::min(d / vec_size, 1024));                                  \
    if (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) {         \
-      VLLM_STABLE_DISPATCH_FLOATING_TYPES(                                     \
+      VLLM_DISPATCH_FLOATING_TYPES(                                            \
          dtype, "act_and_mul_kernel_with_param", [&] {                        \
            vllm::act_and_mul_kernel_with_param<                               \
                scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,  \
@@ -459,11 +454,11 @@ __global__ void swigluoai_and_mul_kernel(
                PACKED_KERNEL<                                                 \
                    typename vllm::PackedTypeConverter<scalar_t>::Type>,       \
                true, true><<<grid, block, 0, stream>>>(                       \
-                out.mutable_data_ptr<scalar_t>(),                              \
-                input.const_data_ptr<scalar_t>(), d, PARAM);                   \
+                out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d,       \
+                PARAM);                                                        \
          });                                                                  \
    } else {                                                                   \
-      VLLM_STABLE_DISPATCH_FLOATING_TYPES(                                     \
+      VLLM_DISPATCH_FLOATING_TYPES(                                            \
          dtype, "act_and_mul_kernel_with_param", [&] {                        \
            vllm::act_and_mul_kernel_with_param<                               \
                scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,  \
@@ -471,49 +466,45 @@ __global__ void swigluoai_and_mul_kernel(
                PACKED_KERNEL<                                                 \
                    typename vllm::PackedTypeConverter<scalar_t>::Type>,       \
                true, false><<<grid, block, 0, stream>>>(                      \
-                out.mutable_data_ptr<scalar_t>(),                              \
-                input.const_data_ptr<scalar_t>(), d, PARAM);                   \
+                out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d,       \
+                PARAM);                                                        \
          });                                                                  \
    }                                                                          \
  } else {                                                                     \
    dim3 block(std::min(d, 1024));                                             \
-    VLLM_STABLE_DISPATCH_FLOATING_TYPES(                                       \
-        dtype, "act_and_mul_kernel_with_param", [&] {                          \
-          vllm::act_and_mul_kernel_with_param<                                 \
-              scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,    \
-              KERNEL<scalar_t>,                                                \
-              PACKED_KERNEL<                                                   \
-                  typename vllm::PackedTypeConverter<scalar_t>::Type>,         \
-              false><<<grid, block, 0, stream>>>(                              \
-              out.mutable_data_ptr<scalar_t>(),                                \
-              input.const_data_ptr<scalar_t>(), d, PARAM);                     \
-        });                                                                    \
+    VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel_with_param", [&] { \
+      vllm::act_and_mul_kernel_with_param<                                     \
+          scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,        \
+          KERNEL<scalar_t>,                                                    \
+          PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>,   \
+          false><<<grid, block, 0, stream>>>(                                  \
+          out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, PARAM);     \
+    });                                                                        \
  }

-#define LAUNCH_SIGLUOAI_AND_MUL(KERNEL, ALPHA, LIMIT)                         \
-  int d = input.size(-1) / 2;                                                 \
-  int64_t num_tokens = input.numel() / input.size(-1);                        \
-  dim3 grid(num_tokens);                                                      \
-  dim3 block(std::min(d, 1024));                                              \
-  const torch::stable::accelerator::DeviceGuard device_guard(                 \
-      input.get_device_index());                                              \
-  const cudaStream_t stream = get_current_cuda_stream();                      \
-  VLLM_STABLE_DISPATCH_FLOATING_TYPES(                                        \
-      input.scalar_type(), "clamp_swiglu_kernel_with_params", [&] {           \
-        vllm::swigluoai_and_mul_kernel<scalar_t, KERNEL<scalar_t>>            \
-            <<<grid, block, 0, stream>>>(out.mutable_data_ptr<scalar_t>(),    \
-                                         input.const_data_ptr<scalar_t>(), d, \
-                                         ALPHA, LIMIT);                       \
+#define LAUNCH_SIGLUOAI_AND_MUL(KERNEL, ALPHA, LIMIT)                          \
+  int d = input.size(-1) / 2;                                                  \
+  int64_t num_tokens = input.numel() / input.size(-1);                         \
+  dim3 grid(num_tokens);                                                       \
+  dim3 block(std::min(d, 1024));                                               \
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));            \
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();                \
+  VLLM_DISPATCH_FLOATING_TYPES(                                                \
+      input.scalar_type(), "clamp_swiglu_kernel_with_params", [&] {            \
+        vllm::swigluoai_and_mul_kernel<scalar_t, KERNEL<scalar_t>>             \
+            <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(),             \
+                                         input.data_ptr<scalar_t>(), d, ALPHA, \
+                                         LIMIT);                               \
      });

-void fatrelu_and_mul(torch::stable::Tensor& out,    // [..., d],
-                     torch::stable::Tensor& input,  // [..., 2 * d]
+void fatrelu_and_mul(torch::Tensor& out,    // [..., d],
+                     torch::Tensor& input,  // [..., 2 * d]
                     double threshold) {
  LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM(
      vllm::fatrelu_kernel, vllm::packed_fatrelu_kernel, threshold);
 }
-void swigluoai_and_mul(torch::stable::Tensor& out,    // [..., d]
-                       torch::stable::Tensor& input,  // [..., 2 * d]
+void swigluoai_and_mul(torch::Tensor& out,    // [..., d]
+                       torch::Tensor& input,  // [..., 2 * d]
                       double alpha, double limit) {
  LAUNCH_SIGLUOAI_AND_MUL(vllm::swigluoai_and_mul, alpha, limit);
 }
@@ -568,46 +559,45 @@ __global__ void activation_kernel(
 }  // namespace vllm

 // Launch element-wise activation kernel.
-#define LAUNCH_ACTIVATION_KERNEL(KERNEL)                                       \
-  auto dtype = input.scalar_type();                                            \
-  int d = input.size(-1);                                                      \
-  int64_t num_tokens = input.numel() / input.size(-1);                         \
-  if (num_tokens == 0) {                                                       \
-    return;                                                                    \
-  }                                                                            \
-  dim3 grid(num_tokens);                                                       \
-  int cc_major = get_device_prop()->major;                                     \
-  int support_vec =                                                            \
-      (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128)            \
-          ? vllm::VecTraits<true>::ARCH_MAX_VEC_SIZE                           \
-          : vllm::VecTraits<false>::ARCH_MAX_VEC_SIZE;                         \
-  int vec_size = support_vec / input.element_size();                           \
-  const bool use_vec = (d % vec_size == 0);                                    \
-  const torch::stable::accelerator::DeviceGuard device_guard(                  \
-      input.get_device_index());                                               \
-  const cudaStream_t stream = get_current_cuda_stream();                       \
-  if (use_vec) {                                                               \
-    dim3 block(std::min(d / vec_size, 1024));                                  \
-    if (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) {         \
-      VLLM_STABLE_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] {    \
-        vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, true, true>        \
-            <<<grid, block, 0, stream>>>(out.mutable_data_ptr<scalar_t>(),     \
-                                         input.const_data_ptr<scalar_t>(), d); \
-      });                                                                      \
-    } else {                                                                   \
-      VLLM_STABLE_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] {    \
-        vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, true, false>       \
-            <<<grid, block, 0, stream>>>(out.mutable_data_ptr<scalar_t>(),     \
-                                         input.const_data_ptr<scalar_t>(), d); \
-      });                                                                      \
-    }                                                                          \
-  } else {                                                                     \
-    dim3 block(std::min(d, 1024));                                             \
-    VLLM_STABLE_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] {      \
-      vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, false>               \
-          <<<grid, block, 0, stream>>>(out.mutable_data_ptr<scalar_t>(),       \
-                                       input.const_data_ptr<scalar_t>(), d);   \
-    });                                                                        \
+#define LAUNCH_ACTIVATION_KERNEL(KERNEL)                                 \
+  auto dtype = input.scalar_type();                                      \
+  int d = input.size(-1);                                                \
+  int64_t num_tokens = input.numel() / input.size(-1);                   \
+  if (num_tokens == 0) {                                                 \
+    return;                                                              \
+  }                                                                      \
+  dim3 grid(num_tokens);                                                 \
+  int cc_major = at::cuda::getCurrentDeviceProperties()->major;          \
+  int support_vec =                                                      \
+      (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128)      \
+          ? vllm::VecTraits<true>::ARCH_MAX_VEC_SIZE                     \
+          : vllm::VecTraits<false>::ARCH_MAX_VEC_SIZE;                   \
+  int vec_size = support_vec / at::elementSize(dtype);                   \
+  const bool use_vec = (d % vec_size == 0);                              \
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));      \
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();          \
+  if (use_vec) {                                                         \
+    dim3 block(std::min(d / vec_size, 1024));                            \
+    if (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) {   \
+      VLLM_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] {     \
+        vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, true, true>  \
+            <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(),       \
+                                         input.data_ptr<scalar_t>(), d); \
+      });                                                                \
+    } else {                                                             \
+      VLLM_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] {     \
+        vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, true, false> \
+            <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(),       \
+                                         input.data_ptr<scalar_t>(), d); \
+      });                                                                \
+    }                                                                    \
+  } else {                                                               \
+    dim3 block(std::min(d, 1024));                                       \
+    VLLM_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] {       \
+      vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, false>         \
+          <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(),         \
+                                       input.data_ptr<scalar_t>(), d);   \
+    });                                                                  \
  }

 namespace vllm {
@@ -635,20 +625,20 @@ __device__ __forceinline__ T gelu_quick_kernel(const T& x) {

 }  // namespace vllm

-void gelu_new(torch::stable::Tensor& out,    // [..., d]
-              torch::stable::Tensor& input)  // [..., d]
+void gelu_new(torch::Tensor& out,    // [..., d]
+              torch::Tensor& input)  // [..., d]
 {
  LAUNCH_ACTIVATION_KERNEL(vllm::gelu_new_kernel);
 }

-void gelu_fast(torch::stable::Tensor& out,    // [..., d]
-               torch::stable::Tensor& input)  // [..., d]
+void gelu_fast(torch::Tensor& out,    // [..., d]
+               torch::Tensor& input)  // [..., d]
 {
  LAUNCH_ACTIVATION_KERNEL(vllm::gelu_fast_kernel);
 }

-void gelu_quick(torch::stable::Tensor& out,    // [..., d]
-                torch::stable::Tensor& input)  // [..., d]
+void gelu_quick(torch::Tensor& out,    // [..., d]
+                torch::Tensor& input)  // [..., d]
 {
  LAUNCH_ACTIVATION_KERNEL(vllm::gelu_quick_kernel);
 }
@@ -17,18 +17,21 @@
 * limitations under the License.
 */

+#include <torch/all.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
 #include <algorithm>

-#include "../../attention/attention_dtypes.h"
+#include "attention_dtypes.h"
 #include "attention_utils.cuh"
-#include "../../cuda_compat.h"
+#include "../cuda_compat.h"

 #ifdef USE_ROCM
  #include <hip/hip_bf16.h>
-  #include "../../quantization/w8a8/fp8/amd/quant_utils.cuh"
+  #include "../quantization/w8a8/fp8/amd/quant_utils.cuh"
 typedef __hip_bfloat16 __nv_bfloat16;
 #else
-  #include "../../quantization/w8a8/fp8/nvidia/quant_utils.cuh"
+  #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh"
 #endif

 #define MAX(a, b) ((a) > (b) ? (a) : (b))
@@ -18,8 +18,8 @@
 */
 #pragma once

-#include "../../cuda_compat.h"
-#include "../../attention/attention_dtypes.h"
+#include "../cuda_compat.h"
+#include "attention_dtypes.h"

 #include <float.h>
 #include <type_traits>
@@ -1,7 +1,6 @@
 #pragma once

 #include "attention_generic.cuh"
-#include "torch_utils.h"

 #include <stdint.h>
 #ifdef ENABLE_FP8
@@ -31,7 +30,7 @@ inline Fp8KVCacheDataType get_fp8_kv_cache_data_type(
  } else if (dtype_str == "fp8_e5m2") {
    return Fp8KVCacheDataType::kFp8E5M2;
  }
-  TORCH_UTILS_CHECK(false, "Unsupported fp8 kv cache data type: ", dtype_str);
+  TORCH_CHECK(false, "Unsupported fp8 kv cache data type: ", dtype_str);
 }

 // fp8 vector types for quantization of kv cache
@@ -1,14 +1,14 @@
 #include <optional>
+#include <torch/all.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
 #include <algorithm>
 #include <limits>

-#include "../torch_utils.h"
-#include "../dispatch_utils.h"
-#include <torch/headeronly/core/ScalarType.h>
-
-#include "../../attention/attention_dtypes.h"
+#include "attention_dtypes.h"
 #include "attention_utils.cuh"
-#include "../../quantization/w8a8/fp8/common.cuh"
+#include "../quantization/w8a8/fp8/common.cuh"
+#include "../dispatch_utils.h"

 namespace vllm {

@@ -196,17 +196,17 @@ __global__ void merge_attn_states_kernel(
 // The following macro is used to dispatch the conversion function based on
 // the output data type. The FN is a macro that calls a function with
 // template<typename scalar_t>.
-#define DISPATCH_BY_SCALAR_DTYPE(scalar_dtype, fn)                          \
-  {                                                                         \
-    if (scalar_dtype == torch::headeronly::ScalarType::Float) {             \
-      fn(float);                                                            \
-    } else if (scalar_dtype == torch::headeronly::ScalarType::Half) {       \
-      fn(uint16_t);                                                         \
-    } else if (scalar_dtype == torch::headeronly::ScalarType::BFloat16) {   \
-      fn(__nv_bfloat16);                                                    \
-    } else {                                                                \
-      STD_TORCH_CHECK(false, "Unsupported data type of O: ", scalar_dtype); \
-    }                                                                       \
+#define DISPATCH_BY_SCALAR_DTYPE(scalar_dtype, fn)                      \
+  {                                                                     \
+    if (scalar_dtype == at::ScalarType::Float) {                        \
+      fn(float);                                                        \
+    } else if (scalar_dtype == at::ScalarType::Half) {                  \
+      fn(uint16_t);                                                     \
+    } else if (scalar_dtype == at::ScalarType::BFloat16) {              \
+      fn(__nv_bfloat16);                                                \
+    } else {                                                            \
+      TORCH_CHECK(false, "Unsupported data type of O: ", scalar_dtype); \
+    }                                                                   \
  }

 #define LAUNCH_MERGE_ATTN_STATES(scalar_t, output_t, NUM_THREADS,           \
@@ -245,14 +245,11 @@ __global__ void merge_attn_states_kernel(
 */
 template <typename scalar_t>
 void merge_attn_states_launcher(
-    torch::stable::Tensor& output,
-    std::optional<torch::stable::Tensor> output_lse,
-    const torch::stable::Tensor& prefix_output,
-    const torch::stable::Tensor& prefix_lse,
-    const torch::stable::Tensor& suffix_output,
-    const torch::stable::Tensor& suffix_lse,
+    torch::Tensor& output, std::optional<torch::Tensor> output_lse,
+    const torch::Tensor& prefix_output, const torch::Tensor& prefix_lse,
+    const torch::Tensor& suffix_output, const torch::Tensor& suffix_lse,
    const std::optional<int64_t> prefill_tokens_with_context,
-    const std::optional<torch::stable::Tensor>& output_scale) {
+    const std::optional<torch::Tensor>& output_scale) {
  constexpr uint NUM_THREADS = 128;
  const uint num_tokens = output.size(0);
  const uint num_heads = output.size(1);
@@ -261,23 +258,23 @@ void merge_attn_states_launcher(
  const uint output_head_stride = output.stride(1);
  // Thread mapping is based on input BF16 pack_size
  const uint pack_size = 16 / sizeof(scalar_t);
-  STD_TORCH_CHECK(head_size % pack_size == 0,
-                  "headsize must be multiple of pack_size:", pack_size);
+  TORCH_CHECK(head_size % pack_size == 0,
+              "headsize must be multiple of pack_size:", pack_size);

  const uint prefix_num_tokens =
      prefill_tokens_with_context.has_value()
          ? static_cast<uint>(prefill_tokens_with_context.value())
          : num_tokens;
-  STD_TORCH_CHECK(prefix_num_tokens <= num_tokens,
-                  "prefix_num_tokens must be <= num_tokens");
+  TORCH_CHECK(prefix_num_tokens <= num_tokens,
+              "prefix_num_tokens must be <= num_tokens");

  float* output_lse_ptr = nullptr;
  if (output_lse.has_value()) {
-    output_lse_ptr = output_lse.value().mutable_data_ptr<float>();
+    output_lse_ptr = output_lse.value().data_ptr<float>();
  }
  float* output_scale_ptr = nullptr;
  if (output_scale.has_value()) {
-    output_scale_ptr = output_scale.value().mutable_data_ptr<float>();
+    output_scale_ptr = output_scale.value().data_ptr<float>();
  }
  // Process one pack elements per thread. for float, the
  // pack_size is 4 for half/bf16, the pack_size is 8.
@@ -287,15 +284,14 @@ void merge_attn_states_launcher(
  dim3 block(NUM_THREADS);
  dim3 grid((total_threads + NUM_THREADS - 1) / NUM_THREADS);

-  const torch::stable::accelerator::DeviceGuard device_guard(
-      prefix_output.get_device_index());
-  auto stream = get_current_cuda_stream();
+  const c10::cuda::OptionalCUDAGuard device_guard(prefix_output.device());
+  auto stream = at::cuda::getCurrentCUDAStream();

  if (output_scale.has_value()) {
    // FP8 output path - dispatch on output FP8 type
-    VLLM_STABLE_DISPATCH_FP8_TYPES(
-        output.scalar_type(), "merge_attn_states_fp8",
-        [&] { LAUNCH_MERGE_ATTN_STATES(scalar_t, fp8_t, NUM_THREADS, true); });
+    VLLM_DISPATCH_FP8_TYPES(output.scalar_type(), "merge_attn_states_fp8", [&] {
+      LAUNCH_MERGE_ATTN_STATES(scalar_t, fp8_t, NUM_THREADS, true);
+    });
  } else {
    // Original BF16/FP16/FP32 output path
    LAUNCH_MERGE_ATTN_STATES(scalar_t, scalar_t, NUM_THREADS, false);
@@ -309,29 +305,26 @@ void merge_attn_states_launcher(
        suffix_lse, prefill_tokens_with_context, output_scale);       \
  }

-void merge_attn_states(
-    torch::stable::Tensor& output,
-    std::optional<torch::stable::Tensor> output_lse,
-    const torch::stable::Tensor& prefix_output,
-    const torch::stable::Tensor& prefix_lse,
-    const torch::stable::Tensor& suffix_output,
-    const torch::stable::Tensor& suffix_lse,
-    const std::optional<int64_t> prefill_tokens_with_context,
-    const std::optional<torch::stable::Tensor>& output_scale) {
+void merge_attn_states(torch::Tensor& output,
+                       std::optional<torch::Tensor> output_lse,
+                       const torch::Tensor& prefix_output,
+                       const torch::Tensor& prefix_lse,
+                       const torch::Tensor& suffix_output,
+                       const torch::Tensor& suffix_lse,
+                       std::optional<int64_t> prefill_tokens_with_context,
+                       const std::optional<torch::Tensor>& output_scale) {
  if (output_scale.has_value()) {
-    STD_TORCH_CHECK(
-        output.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn ||
-            output.scalar_type() ==
-                torch::headeronly::ScalarType::Float8_e4m3fnuz,
-        "output must be FP8 when output_scale is provided, got: ",
-        output.scalar_type());
+    TORCH_CHECK(output.scalar_type() == at::ScalarType::Float8_e4m3fn ||
+                    output.scalar_type() == at::ScalarType::Float8_e4m3fnuz,
+                "output must be FP8 when output_scale is provided, got: ",
+                output.scalar_type());
  } else {
-    STD_TORCH_CHECK(
-        output.scalar_type() == prefix_output.scalar_type(), "output dtype (",
-        output.scalar_type(), ") must match prefix_output dtype (",
-        prefix_output.scalar_type(), ") when output_scale is not set");
+    TORCH_CHECK(output.scalar_type() == prefix_output.scalar_type(),
+                "output dtype (", output.scalar_type(),
+                ") must match prefix_output dtype (",
+                prefix_output.scalar_type(), ") when output_scale is not set");
  }
  // Always dispatch on prefix_output (input) dtype
-  DISPATCH_BY_SCALAR_DTYPE(prefix_output.scalar_type(),
+  DISPATCH_BY_SCALAR_DTYPE(prefix_output.dtype(),
                           CALL_MERGE_ATTN_STATES_LAUNCHER);
 }
@@ -18,12 +18,13 @@ limitations under the License.
 * Taken from SGLANG PR https://github.com/sgl-project/sglang/pull/6929
 * by Alcanderian JieXin Liang
 */
-#include "libtorch_stable/torch_utils.h"
-
-#include <torch/csrc/stable/library.h>
+#include "core/registration.h"

+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
 #include <cutlass/cutlass.h>
 #include <cutlass/kernel_hardware_info.h>
+#include <torch/all.h>

 #include <cute/tensor.hpp>
 #include <iostream>
@@ -34,27 +35,27 @@ limitations under the License.
 // clang-format off
 #if !defined(CUDA_VERSION) || CUDA_VERSION < 12040
 void sm100_cutlass_mla_decode(
-    torch::stable::Tensor const& out,
-    torch::stable::Tensor const& lse,
-    torch::stable::Tensor const& q_nope,
-    torch::stable::Tensor const& q_pe,
-    torch::stable::Tensor const& kv_c_and_k_pe_cache,
-    torch::stable::Tensor const& seq_lens,
-    torch::stable::Tensor const& page_table,
-    torch::stable::Tensor const& workspace,
+    torch::Tensor const& out,
+    torch::Tensor const& lse,
+    torch::Tensor const& q_nope,
+    torch::Tensor const& q_pe,
+    torch::Tensor const& kv_c_and_k_pe_cache,
+    torch::Tensor const& seq_lens,
+    torch::Tensor const& page_table,
+    torch::Tensor const& workspace,
    double sm_scale,
    int64_t num_kv_splits) {
-  STD_TORCH_CHECK(false, "CUDA version must be >= 12.4 for cutlass_mla_decode");
+  TORCH_CHECK(false, "CUDA version must be >= 12.4 for cutlass_mla_decode");
 }
 int64_t sm100_cutlass_mla_get_workspace_size(int64_t max_seq_len, int64_t num_batches, int64_t sm_count, int64_t num_kv_splits) {
-  STD_TORCH_CHECK(false, "CUDA version must be >= 12.4 for cutlass_mla_get_workspace_size");
+  TORCH_CHECK(false, "CUDA version must be >= 12.4 for cutlass_mla_get_workspace_size");
 }
 #else

 #define CUTLASS_CHECK(status)                                                       \
  {                                                                                 \
    cutlass::Status error = status;                                                 \
-    STD_TORCH_CHECK(error == cutlass::Status::kSuccess, cutlassGetStatusString(error)); \
+    TORCH_CHECK(error == cutlass::Status::kSuccess, cutlassGetStatusString(error)); \
  }

 using namespace cute;
@@ -99,23 +100,23 @@ struct MlaSm100 {

 template <typename T>
 typename T::Fmha::Arguments args_from_options(
-    torch::stable::Tensor const& out,
-    torch::stable::Tensor const& lse,
-    torch::stable::Tensor const& q_nope,
-    torch::stable::Tensor const& q_pe,
-    torch::stable::Tensor const& kv_c_and_k_pe_cache,
-    torch::stable::Tensor const& seq_lens,
-    torch::stable::Tensor const& page_table,
+    at::Tensor const& out,
+    at::Tensor const& lse,
+    at::Tensor const& q_nope,
+    at::Tensor const& q_pe,
+    at::Tensor const& kv_c_and_k_pe_cache,
+    at::Tensor const& seq_lens,
+    at::Tensor const& page_table,
    double sm_scale,
    int64_t num_kv_splits) {
  cutlass::KernelHardwareInfo hw_info;
-  hw_info.device_id = q_nope.get_device_index();
+  hw_info.device_id = q_nope.device().index();
  hw_info.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id);

-  int batches = q_nope.size(0);
-  int page_count_per_seq = page_table.size(1);
-  int page_count_total = kv_c_and_k_pe_cache.size(0);
-  int page_size = kv_c_and_k_pe_cache.size(1);
+  int batches = q_nope.sizes()[0];
+  int page_count_per_seq = page_table.sizes()[1];
+  int page_count_total = kv_c_and_k_pe_cache.sizes()[0];
+  int page_size = kv_c_and_k_pe_cache.sizes()[1];
  int max_seq_len = page_size * page_count_per_seq;
  using TileShapeH = typename T::TileShapeH;
  using TileShapeD = typename T::TileShapeD;
@@ -185,14 +186,14 @@ typename T::Fmha::Arguments args_from_options(

 template <typename Element, typename ElementOut, bool IsPaged128, typename PersistenceOption>
 void runMla(
-    torch::stable::Tensor const& out,
-    torch::stable::Tensor const& lse,
-    torch::stable::Tensor const& q_nope,
-    torch::stable::Tensor const& q_pe,
-    torch::stable::Tensor const& kv_c_and_k_pe_cache,
-    torch::stable::Tensor const& seq_lens,
-    torch::stable::Tensor const& page_table,
-    torch::stable::Tensor const& workspace,
+    at::Tensor const& out,
+    at::Tensor const& lse,
+    at::Tensor const& q_nope,
+    at::Tensor const& q_pe,
+    at::Tensor const& kv_c_and_k_pe_cache,
+    at::Tensor const& seq_lens,
+    at::Tensor const& page_table,
+    at::Tensor const& workspace,
    double sm_scale,
    int64_t num_kv_splits,
    cudaStream_t stream) {
@@ -219,37 +220,37 @@ void runMla(
  }()

 void sm100_cutlass_mla_decode(
-    torch::stable::Tensor const& out,
-    torch::stable::Tensor const& lse,
-    torch::stable::Tensor const& q_nope,
-    torch::stable::Tensor const& q_pe,
-    torch::stable::Tensor const& kv_c_and_k_pe_cache,
-    torch::stable::Tensor const& seq_lens,
-    torch::stable::Tensor const& page_table,
-    torch::stable::Tensor const& workspace,
+    torch::Tensor const& out,
+    torch::Tensor const& lse,
+    torch::Tensor const& q_nope,
+    torch::Tensor const& q_pe,
+    torch::Tensor const& kv_c_and_k_pe_cache,
+    torch::Tensor const& seq_lens,
+    torch::Tensor const& page_table,
+    torch::Tensor const& workspace,
    double sm_scale,
    int64_t num_kv_splits) {
-  auto in_dtype = q_nope.scalar_type();
-  torch::stable::accelerator::DeviceGuard device_guard(q_nope.get_device_index());
-  const cudaStream_t stream = get_current_cuda_stream(q_nope.get_device_index());
-  const int page_size = kv_c_and_k_pe_cache.size(1);
-
+  auto in_dtype = q_nope.dtype();
+  at::cuda::CUDAGuard device_guard{(char)q_nope.get_device()};
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream(q_nope.get_device());
+  const int page_size = kv_c_and_k_pe_cache.sizes()[1];
+  
  // NOTE(alcanderian): IsPersistent has bug with manual split_kv.
  // Kernel will hang if batch is too large with large num_kv_splits. (for example bs=8, num_kv_splits=8)
  // Maybe per batch split kv will fix this.
  DISPATCH_BOOL(page_size == 128, IsPaged128, [&] {
    DISPATCH_BOOL(num_kv_splits <= 1, NotManualSplitKV, [&] {
-      if (in_dtype == torch::headeronly::ScalarType::Half) {
+      if (in_dtype == at::ScalarType::Half) {
        runMla<cutlass::half_t, cutlass::half_t, IsPaged128, IsPersistent<NotManualSplitKV>>(
          out, lse, q_nope, q_pe, kv_c_and_k_pe_cache, seq_lens, page_table, workspace, sm_scale, num_kv_splits, stream);
-      } else if (in_dtype == torch::headeronly::ScalarType::BFloat16) {
+      } else if (in_dtype == at::ScalarType::BFloat16) {
        runMla<cutlass::bfloat16_t, cutlass::bfloat16_t, IsPaged128, IsPersistent<NotManualSplitKV>>(
          out, lse, q_nope, q_pe, kv_c_and_k_pe_cache, seq_lens, page_table, workspace, sm_scale, num_kv_splits, stream);
-      } else if (in_dtype == torch::headeronly::ScalarType::Float8_e4m3fn) {
+      } else if (in_dtype == at::ScalarType::Float8_e4m3fn) {
        runMla<cutlass::float_e4m3_t, cutlass::bfloat16_t, IsPaged128, IsPersistent<NotManualSplitKV>>(
          out, lse, q_nope, q_pe, kv_c_and_k_pe_cache, seq_lens, page_table, workspace, sm_scale, num_kv_splits, stream);
      } else {
-        STD_TORCH_CHECK(false, "Unsupported input data type of MLA");
+        TORCH_CHECK(false, "Unsupported input data type of MLA");
      }
      return true;
    });
@@ -279,12 +280,12 @@ int64_t sm100_cutlass_mla_get_workspace_size(int64_t max_seq_len, int64_t num_ba

 #endif

-STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
-  m.impl("sm100_cutlass_mla_decode", TORCH_BOX(&sm100_cutlass_mla_decode));
+TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
+  m.impl("sm100_cutlass_mla_decode", &sm100_cutlass_mla_decode);
 }

-STABLE_TORCH_LIBRARY_IMPL(_C, CompositeExplicitAutograd, m) {
-  m.impl("sm100_cutlass_mla_get_workspace_size", TORCH_BOX(&sm100_cutlass_mla_get_workspace_size));
+TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CatchAll, m) {
+  m.impl("sm100_cutlass_mla_get_workspace_size", &sm100_cutlass_mla_get_workspace_size);
 }

 // clang-format on
@@ -16,9 +16,8 @@
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
-#include "../torch_utils.h"
 #include "attention_kernels.cuh"
-#include "../../cuda_compat.h"
+#include "../cuda_compat.h"

 #define MAX(a, b) ((a) > (b) ? (a) : (b))
 #define MIN(a, b) ((a) < (b) ? (a) : (b))
@@ -45,15 +44,13 @@ template <typename T, typename CACHE_T, int BLOCK_SIZE,
          vllm::Fp8KVCacheDataType KV_DTYPE, bool IS_BLOCK_SPARSE,
          int NUM_THREADS = 128>
 void paged_attention_v1_launcher(
-    torch::stable::Tensor& out, torch::stable::Tensor& query,
-    torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache,
-    int num_kv_heads, float scale, torch::stable::Tensor& block_tables,
-    torch::stable::Tensor& seq_lens, int max_seq_len,
-    const std::optional<torch::stable::Tensor>& alibi_slopes,
-    torch::stable::Tensor& k_scale, torch::stable::Tensor& v_scale,
-    const int tp_rank, const int blocksparse_local_blocks,
-    const int blocksparse_vert_stride, const int blocksparse_block_size,
-    const int blocksparse_head_sliding_step) {
+    torch::Tensor& out, torch::Tensor& query, torch::Tensor& key_cache,
+    torch::Tensor& value_cache, int num_kv_heads, float scale,
+    torch::Tensor& block_tables, torch::Tensor& seq_lens, int max_seq_len,
+    const std::optional<torch::Tensor>& alibi_slopes, torch::Tensor& k_scale,
+    torch::Tensor& v_scale, const int tp_rank,
+    const int blocksparse_local_blocks, const int blocksparse_vert_stride,
+    const int blocksparse_block_size, const int blocksparse_head_sliding_step) {
  int num_seqs = query.size(0);
  int num_heads = query.size(1);
  int head_size = query.size(2);
@@ -72,8 +69,8 @@ void paged_attention_v1_launcher(
  T* query_ptr = reinterpret_cast<T*>(query.data_ptr());
  CACHE_T* key_cache_ptr = reinterpret_cast<CACHE_T*>(key_cache.data_ptr());
  CACHE_T* value_cache_ptr = reinterpret_cast<CACHE_T*>(value_cache.data_ptr());
-  int* block_tables_ptr = block_tables.mutable_data_ptr<int>();
-  int* seq_lens_ptr = seq_lens.mutable_data_ptr<int>();
+  int* block_tables_ptr = block_tables.data_ptr<int>();
+  int* seq_lens_ptr = seq_lens.data_ptr<int>();
  const float* k_scale_ptr = reinterpret_cast<const float*>(k_scale.data_ptr());
  const float* v_scale_ptr = reinterpret_cast<const float*>(v_scale.data_ptr());

@@ -88,9 +85,8 @@ void paged_attention_v1_launcher(

  dim3 grid(num_heads, num_seqs, 1);
  dim3 block(NUM_THREADS);
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      query.get_device_index());
-  const cudaStream_t stream = get_current_cuda_stream();
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  switch (head_size) {
    // NOTE(woosuk): To reduce the compilation time, we only compile for the
    // head sizes that we use in the model. However, we can easily extend this
@@ -123,7 +119,7 @@ void paged_attention_v1_launcher(
      LAUNCH_PAGED_ATTENTION_V1(256);
      break;
    default:
-      STD_TORCH_CHECK(false, "Unsupported head size: ", head_size);
+      TORCH_CHECK(false, "Unsupported head size: ", head_size);
      break;
  }
 }
@@ -145,43 +141,43 @@ void paged_attention_v1_launcher(

 // NOTE(woosuk): To reduce the compilation time, we omitted block sizes
 // 1, 2, 4, 64, 128, 256.
-#define CALL_V1_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE)             \
-  switch (block_size) {                                               \
-    case 8:                                                           \
-      CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE);             \
-      break;                                                          \
-    case 16:                                                          \
-      CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE);            \
-      break;                                                          \
-    case 32:                                                          \
-      CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE);            \
-      break;                                                          \
-    default:                                                          \
-      STD_TORCH_CHECK(false, "Unsupported block size: ", block_size); \
-      break;                                                          \
+#define CALL_V1_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE)         \
+  switch (block_size) {                                           \
+    case 8:                                                       \
+      CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE);         \
+      break;                                                      \
+    case 16:                                                      \
+      CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE);        \
+      break;                                                      \
+    case 32:                                                      \
+      CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE);        \
+      break;                                                      \
+    default:                                                      \
+      TORCH_CHECK(false, "Unsupported block size: ", block_size); \
+      break;                                                      \
  }

 void paged_attention_v1(
-    torch::stable::Tensor& out,    // [num_seqs, num_heads, head_size]
-    torch::stable::Tensor& query,  // [num_seqs, num_heads, head_size]
-    torch::stable::Tensor&
+    torch::Tensor& out,    // [num_seqs, num_heads, head_size]
+    torch::Tensor& query,  // [num_seqs, num_heads, head_size]
+    torch::Tensor&
        key_cache,  // [num_blocks, num_heads, head_size/x, block_size, x]
-    torch::stable::Tensor&
+    torch::Tensor&
        value_cache,       // [num_blocks, num_heads, head_size, block_size]
    int64_t num_kv_heads,  // [num_heads]
    double scale,
-    torch::stable::Tensor& block_tables,  // [num_seqs, max_num_blocks_per_seq]
-    torch::stable::Tensor& seq_lens,      // [num_seqs]
+    torch::Tensor& block_tables,  // [num_seqs, max_num_blocks_per_seq]
+    torch::Tensor& seq_lens,      // [num_seqs]
    int64_t block_size, int64_t max_seq_len,
-    const std::optional<torch::stable::Tensor>& alibi_slopes,
-    const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale,
-    torch::stable::Tensor& v_scale, const int64_t tp_rank,
+    const std::optional<torch::Tensor>& alibi_slopes,
+    const std::string& kv_cache_dtype, torch::Tensor& k_scale,
+    torch::Tensor& v_scale, const int64_t tp_rank,
    const int64_t blocksparse_local_blocks,
    const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
    const int64_t blocksparse_head_sliding_step) {
  const bool is_block_sparse = (blocksparse_vert_stride > 1);

-  DISPATCH_BY_KV_CACHE_DTYPE(query.scalar_type(), kv_cache_dtype,
+  DISPATCH_BY_KV_CACHE_DTYPE(query.dtype(), kv_cache_dtype,
                             CALL_V1_LAUNCHER_BLOCK_SIZE)
 }

@@ -16,9 +16,8 @@
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
-#include "../torch_utils.h"
 #include "attention_kernels.cuh"
-#include "../../cuda_compat.h"
+#include "../cuda_compat.h"

 #define MAX(a, b) ((a) > (b) ? (a) : (b))
 #define MIN(a, b) ((a) < (b) ? (a) : (b))
@@ -45,16 +44,14 @@ template <typename T, typename CACHE_T, int BLOCK_SIZE,
          vllm::Fp8KVCacheDataType KV_DTYPE, bool IS_BLOCK_SPARSE,
          int NUM_THREADS = 128, int PARTITION_SIZE = 512>
 void paged_attention_v2_launcher(
-    torch::stable::Tensor& out, torch::stable::Tensor& exp_sums,
-    torch::stable::Tensor& max_logits, torch::stable::Tensor& tmp_out,
-    torch::stable::Tensor& query, torch::stable::Tensor& key_cache,
-    torch::stable::Tensor& value_cache, int num_kv_heads, float scale,
-    torch::stable::Tensor& block_tables, torch::stable::Tensor& seq_lens,
-    int max_seq_len, const std::optional<torch::stable::Tensor>& alibi_slopes,
-    torch::stable::Tensor& k_scale, torch::stable::Tensor& v_scale,
-    const int tp_rank, const int blocksparse_local_blocks,
-    const int blocksparse_vert_stride, const int blocksparse_block_size,
-    const int blocksparse_head_sliding_step) {
+    torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits,
+    torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,
+    torch::Tensor& value_cache, int num_kv_heads, float scale,
+    torch::Tensor& block_tables, torch::Tensor& seq_lens, int max_seq_len,
+    const std::optional<torch::Tensor>& alibi_slopes, torch::Tensor& k_scale,
+    torch::Tensor& v_scale, const int tp_rank,
+    const int blocksparse_local_blocks, const int blocksparse_vert_stride,
+    const int blocksparse_block_size, const int blocksparse_head_sliding_step) {
  int num_seqs = query.size(0);
  int num_heads = query.size(1);
  int head_size = query.size(2);
@@ -76,8 +73,8 @@ void paged_attention_v2_launcher(
  T* query_ptr = reinterpret_cast<T*>(query.data_ptr());
  CACHE_T* key_cache_ptr = reinterpret_cast<CACHE_T*>(key_cache.data_ptr());
  CACHE_T* value_cache_ptr = reinterpret_cast<CACHE_T*>(value_cache.data_ptr());
-  int* block_tables_ptr = block_tables.mutable_data_ptr<int>();
-  int* seq_lens_ptr = seq_lens.mutable_data_ptr<int>();
+  int* block_tables_ptr = block_tables.data_ptr<int>();
+  int* seq_lens_ptr = seq_lens.data_ptr<int>();
  const float* k_scale_ptr = reinterpret_cast<const float*>(k_scale.data_ptr());
  const float* v_scale_ptr = reinterpret_cast<const float*>(v_scale.data_ptr());

@@ -94,9 +91,8 @@ void paged_attention_v2_launcher(
  int reduce_shared_mem_size = 2 * max_num_partitions * sizeof(float);

  dim3 block(NUM_THREADS);
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      query.get_device_index());
-  const cudaStream_t stream = get_current_cuda_stream();
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  switch (head_size) {
    // NOTE(woosuk): To reduce the compilation time, we only compile for the
    // head sizes that we use in the model. However, we can easily extend this
@@ -129,7 +125,7 @@ void paged_attention_v2_launcher(
      LAUNCH_PAGED_ATTENTION_V2(256);
      break;
    default:
-      STD_TORCH_CHECK(false, "Unsupported head size: ", head_size);
+      TORCH_CHECK(false, "Unsupported head size: ", head_size);
      break;
  }
 }
@@ -152,48 +148,46 @@ void paged_attention_v2_launcher(

 // NOTE(woosuk): To reduce the compilation time, we omitted block sizes
 // 1, 2, 4, 64, 128, 256.
-#define CALL_V2_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE)             \
-  switch (block_size) {                                               \
-    case 8:                                                           \
-      CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE);             \
-      break;                                                          \
-    case 16:                                                          \
-      CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE);            \
-      break;                                                          \
-    case 32:                                                          \
-      CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE);            \
-      break;                                                          \
-    default:                                                          \
-      STD_TORCH_CHECK(false, "Unsupported block size: ", block_size); \
-      break;                                                          \
+#define CALL_V2_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE)         \
+  switch (block_size) {                                           \
+    case 8:                                                       \
+      CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE);         \
+      break;                                                      \
+    case 16:                                                      \
+      CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE);        \
+      break;                                                      \
+    case 32:                                                      \
+      CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE);        \
+      break;                                                      \
+    default:                                                      \
+      TORCH_CHECK(false, "Unsupported block size: ", block_size); \
+      break;                                                      \
  }

 void paged_attention_v2(
-    torch::stable::Tensor& out,  // [num_seqs, num_heads, head_size]
-    torch::stable::Tensor&
-        exp_sums,  // [num_seqs, num_heads, max_num_partitions]
-    torch::stable::Tensor&
-        max_logits,  // [num_seqs, num_heads, max_num_partitions]
-    torch::stable::Tensor&
+    torch::Tensor& out,         // [num_seqs, num_heads, head_size]
+    torch::Tensor& exp_sums,    // [num_seqs, num_heads, max_num_partitions]
+    torch::Tensor& max_logits,  // [num_seqs, num_heads, max_num_partitions]
+    torch::Tensor&
        tmp_out,  // [num_seqs, num_heads, max_num_partitions, head_size]
-    torch::stable::Tensor& query,  // [num_seqs, num_heads, head_size]
-    torch::stable::Tensor&
+    torch::Tensor& query,  // [num_seqs, num_heads, head_size]
+    torch::Tensor&
        key_cache,  // [num_blocks, num_heads, head_size/x, block_size, x]
-    torch::stable::Tensor&
+    torch::Tensor&
        value_cache,       // [num_blocks, num_heads, head_size, block_size]
    int64_t num_kv_heads,  // [num_heads]
    double scale,
-    torch::stable::Tensor& block_tables,  // [num_seqs, max_num_blocks_per_seq]
-    torch::stable::Tensor& seq_lens,      // [num_seqs]
+    torch::Tensor& block_tables,  // [num_seqs, max_num_blocks_per_seq]
+    torch::Tensor& seq_lens,      // [num_seqs]
    int64_t block_size, int64_t max_seq_len,
-    const std::optional<torch::stable::Tensor>& alibi_slopes,
-    const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale,
-    torch::stable::Tensor& v_scale, const int64_t tp_rank,
+    const std::optional<torch::Tensor>& alibi_slopes,
+    const std::string& kv_cache_dtype, torch::Tensor& k_scale,
+    torch::Tensor& v_scale, const int64_t tp_rank,
    const int64_t blocksparse_local_blocks,
    const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
    const int64_t blocksparse_head_sliding_step) {
  const bool is_block_sparse = (blocksparse_vert_stride > 1);
-  DISPATCH_BY_KV_CACHE_DTYPE(query.scalar_type(), kv_cache_dtype,
+  DISPATCH_BY_KV_CACHE_DTYPE(query.dtype(), kv_cache_dtype,
                             CALL_V2_LAUNCHER_BLOCK_SIZE)
 }

@@ -0,0 +1,401 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+#include <assert.h>
+
+#include <cuda.h>
+
+#include <torch/all.h>
+
+__device__ int64_t save_blocks(int* block_offset, int64_t range_start,
+                               int64_t range_end, int64_t block_size,
+                               int64_t input_block_count, int64_t kv_seqlen) {
+  if (range_start >= kv_seqlen) {
+    return input_block_count;
+  }
+  if (range_end > kv_seqlen) {
+    range_end = kv_seqlen;
+  }
+  int64_t current_block_count = input_block_count;
+  for (int idx = range_start; idx < range_end; idx += block_size) {
+    block_offset[current_block_count++] = idx;
+  }
+  return current_block_count;
+}
+
+__global__ void convert_vertical_slash_indexes_kernel(
+    const int* q_seqlens,         // [BATCH, ]
+    const int* kv_seqlens,        // [BATCH, ]
+    const int* vertical_indexes,  // [BATCH, N_HEADS, NNZ_V]
+    const int* slash_indexes,     // [BATCH, N_HEADS, NNZ_S]
+    int* block_count,             // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
+    int* block_offset,  // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_S]
+    int* column_count,  // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
+    int* column_index,  // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_V]
+    int64_t N_HEADS, int64_t N_ROWS, int64_t BLOCK_SIZE_M, int64_t BLOCK_SIZE_N,
+    int64_t NNZ_V, int64_t NNZ_S,
+    bool causal  // True for intra, False for succ
+) {
+  const int batch_idx = blockIdx.y;
+  const int head_idx = blockIdx.x;
+  const int group_idx = blockIdx.z;
+
+  int64_t q_seqlen = q_seqlens[batch_idx];
+  int64_t kv_seqlen = kv_seqlens[batch_idx];
+  int64_t block_idx_m = group_idx * blockDim.x + threadIdx.x;
+  int64_t start_m = block_idx_m * BLOCK_SIZE_M;
+  if (start_m >= q_seqlen) {
+    return;
+  }
+  int64_t end_m = start_m + BLOCK_SIZE_M;
+  vertical_indexes += (batch_idx * N_HEADS + head_idx) * NNZ_V;
+  slash_indexes += (batch_idx * N_HEADS + head_idx) * NNZ_S;
+  int64_t row_offset = (batch_idx * N_HEADS + head_idx) * N_ROWS + block_idx_m;
+  block_count += row_offset;
+  block_offset += row_offset * NNZ_S;
+  column_count += row_offset;
+  column_index += row_offset * NNZ_V;
+
+  bool has_slash = true;
+  int64_t tmp_col_cnt = 0, tmp_blk_cnt = 0;
+  int64_t s = 0, v = 0;
+  int64_t v_idx = vertical_indexes[v++];
+  int64_t s_idx = slash_indexes[s++];
+  if (causal) {
+    while (s_idx >= end_m + (kv_seqlen - q_seqlen) && s < NNZ_S) {
+      s_idx = slash_indexes[s++];
+    }
+    if (s_idx > end_m + (kv_seqlen - q_seqlen)) has_slash = false;
+    s_idx = max((kv_seqlen - q_seqlen) + end_m - s_idx, BLOCK_SIZE_M);
+  } else {
+    while (s_idx >= end_m + kv_seqlen && s < NNZ_S) {
+      s_idx = slash_indexes[s++];
+    }
+    if (s_idx > end_m + kv_seqlen) has_slash = false;
+    s_idx = max(kv_seqlen + end_m - s_idx, BLOCK_SIZE_M);
+  }
+
+  int64_t range_start = s_idx - BLOCK_SIZE_M, range_end = s_idx;
+  if (!has_slash) {
+    if (causal) {
+      range_start = (kv_seqlen - q_seqlen) + end_m;
+      range_end = (kv_seqlen - q_seqlen) + end_m + BLOCK_SIZE_N;
+    } else {
+      range_start = kv_seqlen;
+      range_end = kv_seqlen + BLOCK_SIZE_N;
+    }
+  }
+
+  bool slash_finished = false;
+  while (1) {
+    if (v_idx < range_end) {
+      if (v_idx < range_start) {
+        column_index[tmp_col_cnt++] = v_idx;
+      }
+      if (v < NNZ_V) {
+        v_idx = vertical_indexes[v++];
+      } else {
+        if (causal)
+          v_idx = end_m + BLOCK_SIZE_N + (kv_seqlen - q_seqlen);
+        else
+          v_idx = end_m + BLOCK_SIZE_N + kv_seqlen;
+      }
+    } else {
+      if ((s < NNZ_S && causal) ||
+          (s < NNZ_S && !causal && slash_indexes[s] >= start_m)) {
+        if (causal)
+          s_idx = max((kv_seqlen - q_seqlen) + end_m - slash_indexes[s++],
+                      BLOCK_SIZE_M);
+        else
+          s_idx = max(kv_seqlen + end_m - slash_indexes[s++], BLOCK_SIZE_M);
+      } else {
+        if (v == NNZ_V || (v_idx > range_start && causal)) {
+          // add the last vertical if no more slash
+          if (v == NNZ_V && !causal && v_idx < kv_seqlen) {
+            column_index[tmp_col_cnt++] = v_idx;
+          }
+          tmp_blk_cnt = save_blocks(block_offset, range_start, range_end,
+                                    BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen);
+          break;
+        } else {
+          if (causal) {
+            range_start = (kv_seqlen - q_seqlen) + end_m;
+            range_end = (kv_seqlen - q_seqlen) + end_m + BLOCK_SIZE_N;
+          } else {
+            // if slash_finished but there are vertical left, save current
+            // blocks
+            tmp_blk_cnt = save_blocks(block_offset, range_start, range_end,
+                                      BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen);
+            range_start = kv_seqlen;
+            range_end = kv_seqlen + BLOCK_SIZE_N;
+          }
+          slash_finished = true;
+        }
+      }
+      if (!slash_finished) {
+        if (s_idx > range_end + BLOCK_SIZE_M) {
+          tmp_blk_cnt = save_blocks(block_offset, range_start, range_end,
+                                    BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen);
+          range_start = s_idx - BLOCK_SIZE_M;
+          range_end = s_idx;
+        } else if (s_idx > range_end) {
+          range_end += BLOCK_SIZE_M;
+        }
+      }
+    }
+  }
+
+  block_count[0] = tmp_blk_cnt;
+  column_count[0] = tmp_col_cnt;
+}
+
+void convert_vertical_slash_indexes_64x64(
+    const int* q_seqlens,         // [BATCH, ]
+    const int* kv_seqlens,        // [BATCH, ]
+    const int* vertical_indexes,  // [BATCH, N_HEADS, NNZ_V]
+    const int* slash_indexes,     // [BATCH, N_HEADS, NNZ_S]
+    int* block_count,             // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
+    int* block_offset,  // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_S]
+    int* column_count,  // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
+    int* column_index,  // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_V]
+    int64_t BATCH_SIZE, int64_t N_HEADS, int64_t N_ROWS, int64_t BLOCK_SIZE_M,
+    int64_t BLOCK_SIZE_N, int64_t NNZ_V, int64_t NNZ_S, bool causal) {
+  const int N_THREADS = 64;
+  const dim3 dimBlock(N_THREADS);
+  const dim3 dimGrid(N_HEADS, BATCH_SIZE, (N_ROWS + N_THREADS - 1) / N_THREADS);
+  convert_vertical_slash_indexes_kernel<<<dimGrid, dimBlock>>>(
+      q_seqlens, kv_seqlens, vertical_indexes, slash_indexes, block_count,
+      block_offset, column_count, column_index, N_HEADS, N_ROWS, BLOCK_SIZE_M,
+      BLOCK_SIZE_N, NNZ_V, NNZ_S, causal);
+}
+
+/**
+ * Implements the Algorithm 4 in paper https://arxiv.org/abs/2407.02490.
+ *
+ * This function builds the index of each row of blocks from vertical indices
+ * and slash indices. The vertical indices are treated as points, while the
+ * slash indices are converted as ranges. The output consists of the merged
+ * ranges and separate column indices, where the ranges are represented by
+ * block indices.
+ *
+ * The implementation is referenced from the original MInference repo:
+ * https://github.com/microsoft/MInference/blob/main/csrc/vertical_slash_index.cu.
+ */
+void convert_vertical_slash_indexes(
+    torch::Tensor& block_count,      // [BATCH, N_HEADS, NUM_ROWS]
+    torch::Tensor& block_offset,     // [BATCH, N_HEADS, NUM_ROWS, NNZ_S]
+    torch::Tensor& column_count,     // [BATCH, N_HEADS, NUM_ROWS]
+    torch::Tensor& column_index,     // [BATCH, N_HEADS, NUM_ROWS, NNZ_V]
+    torch::Tensor q_seqlens,         // [BATCH, ]
+    torch::Tensor kv_seqlens,        // [BATCH, ]
+    torch::Tensor vertical_indexes,  // [BATCH, N_HEADS, NNZ_V]
+    torch::Tensor slash_indexes,     // [BATCH, N_HEADS, NNZ_S]
+    int64_t context_size, int64_t block_size_M, int64_t block_size_N,
+    bool causal) {
+  cudaSetDevice(q_seqlens.get_device());
+
+  int batch_size = slash_indexes.size(0);
+  int num_heads = slash_indexes.size(1);
+  int nnz_slash = slash_indexes.size(2);
+  int nnz_vertical = vertical_indexes.size(2);
+  int num_rows = (context_size + block_size_M - 1) / block_size_M;
+
+  convert_vertical_slash_indexes_64x64(
+      q_seqlens.data_ptr<int>(), kv_seqlens.data_ptr<int>(),
+      vertical_indexes.data_ptr<int>(), slash_indexes.data_ptr<int>(),
+      block_count.data_ptr<int>(), block_offset.data_ptr<int>(),
+      column_count.data_ptr<int>(), column_index.data_ptr<int>(), batch_size,
+      num_heads, num_rows, block_size_M, block_size_N, nnz_vertical, nnz_slash,
+      causal);
+}
+
+__global__ void convert_vertical_slash_indexes_kernel_mergehead(
+    const int* q_seqlens,         // [BATCH, ]
+    const int* kv_seqlens,        // [BATCH, ]
+    const int* vertical_indexes,  // [BATCH, N_HEADS, NNZ_V]
+    const int* slash_indexes,     // [BATCH, N_HEADS, NNZ_S]
+    const int* per_head_vertical_topkv, const int* per_head_slash_topkv,
+    int* block_count,   // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
+    int* block_offset,  // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_S]
+    int* column_count,  // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
+    int* column_index,  // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_V]
+    int64_t N_HEADS, int64_t N_ROWS, int64_t BLOCK_SIZE_M, int64_t BLOCK_SIZE_N,
+    int64_t NNZ_V, int64_t NNZ_S,
+    bool causal  // True for intra, False for succ
+) {
+  const int batch_idx = blockIdx.y;
+  const int head_idx = blockIdx.x;
+  const int group_idx = blockIdx.z;
+
+  int64_t q_seqlen = q_seqlens[batch_idx];
+  int64_t kv_seqlen = kv_seqlens[batch_idx];
+  int64_t block_idx_m = group_idx * blockDim.x + threadIdx.x;
+  int64_t start_m = block_idx_m * BLOCK_SIZE_M;
+  if (start_m >= q_seqlen) {
+    return;
+  }
+  int64_t end_m = start_m + BLOCK_SIZE_M;
+  vertical_indexes += (batch_idx * N_HEADS + head_idx) * NNZ_V;
+  slash_indexes += (batch_idx * N_HEADS + head_idx) * NNZ_S;
+  int64_t row_offset = (batch_idx * N_HEADS + head_idx) * N_ROWS + block_idx_m;
+  block_count += row_offset;
+  block_offset += row_offset * NNZ_S;
+  column_count += row_offset;
+  column_index += row_offset * NNZ_V;
+
+  // MergeHead: each head has it's unique max topk NNZ_V，NNZ_S. (NNZ_V，NNZ_S
+  // above is buffer size, use to compute offset)
+  NNZ_S = per_head_slash_topkv[head_idx];
+  NNZ_V = per_head_vertical_topkv[head_idx];
+
+  bool has_slash = true;
+  int64_t tmp_col_cnt = 0, tmp_blk_cnt = 0;
+  int64_t s = 0, v = 0;
+  int64_t v_idx = vertical_indexes[v++];
+  int64_t s_idx = slash_indexes[s++];
+  if (causal) {
+    while (s_idx >= end_m + (kv_seqlen - q_seqlen) && s < NNZ_S) {
+      s_idx = slash_indexes[s++];
+    }
+    if (s_idx > end_m + (kv_seqlen - q_seqlen)) has_slash = false;
+    s_idx = max((kv_seqlen - q_seqlen) + end_m - s_idx, BLOCK_SIZE_M);
+  } else {
+    while (s_idx >= end_m + kv_seqlen && s < NNZ_S) {
+      s_idx = slash_indexes[s++];
+    }
+    if (s_idx > end_m + kv_seqlen) has_slash = false;
+    s_idx = max(kv_seqlen + end_m - s_idx, BLOCK_SIZE_M);
+  }
+
+  int64_t range_start = s_idx - BLOCK_SIZE_M, range_end = s_idx;
+  if (!has_slash) {
+    if (causal) {
+      range_start = (kv_seqlen - q_seqlen) + end_m;
+      range_end = (kv_seqlen - q_seqlen) + end_m + BLOCK_SIZE_N;
+    } else {
+      range_start = kv_seqlen;
+      range_end = kv_seqlen + BLOCK_SIZE_N;
+    }
+  }
+
+  bool slash_finished = false;
+  while (1) {
+    if (v_idx < range_end) {
+      if (v_idx < range_start) {
+        column_index[tmp_col_cnt++] = v_idx;
+      }
+      if (v < NNZ_V) {
+        v_idx = vertical_indexes[v++];
+      } else {
+        if (causal)
+          v_idx = end_m + BLOCK_SIZE_N + (kv_seqlen - q_seqlen);
+        else
+          v_idx = end_m + BLOCK_SIZE_N + kv_seqlen;
+      }
+    } else {
+      if ((s < NNZ_S && causal) ||
+          (s < NNZ_S && !causal && slash_indexes[s] >= start_m)) {
+        if (causal)
+          s_idx = max((kv_seqlen - q_seqlen) + end_m - slash_indexes[s++],
+                      BLOCK_SIZE_M);
+        else
+          s_idx = max(kv_seqlen + end_m - slash_indexes[s++], BLOCK_SIZE_M);
+      } else {
+        if (v == NNZ_V || (v_idx > range_start && causal)) {
+          // add the last vertical if no more slash
+          if (v == NNZ_V && !causal && v_idx < kv_seqlen) {
+            column_index[tmp_col_cnt++] = v_idx;
+          }
+          tmp_blk_cnt = save_blocks(block_offset, range_start, range_end,
+                                    BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen);
+          break;
+        } else {
+          if (causal) {
+            range_start = (kv_seqlen - q_seqlen) + end_m;
+            range_end = (kv_seqlen - q_seqlen) + end_m + BLOCK_SIZE_N;
+          } else {
+            // if slash_finished but there are vertical left, save current
+            // blocks
+            tmp_blk_cnt = save_blocks(block_offset, range_start, range_end,
+                                      BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen);
+            range_start = kv_seqlen;
+            range_end = kv_seqlen + BLOCK_SIZE_N;
+          }
+          slash_finished = true;
+        }
+      }
+      if (!slash_finished) {
+        if (s_idx > range_end + BLOCK_SIZE_M) {
+          tmp_blk_cnt = save_blocks(block_offset, range_start, range_end,
+                                    BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen);
+          range_start = s_idx - BLOCK_SIZE_M;
+          range_end = s_idx;
+        } else if (s_idx > range_end) {
+          range_end += BLOCK_SIZE_M;
+        }
+      }
+    }
+  }
+
+  block_count[0] = tmp_blk_cnt;
+  column_count[0] = tmp_col_cnt;
+}
+
+void convert_vertical_slash_indexes_64x64_mergehead(
+    const int* q_seqlens,         // [BATCH, ]
+    const int* kv_seqlens,        // [BATCH, ]
+    const int* vertical_indexes,  // [BATCH, N_HEADS, NNZ_V]
+    const int* slash_indexes,     // [BATCH, N_HEADS, NNZ_S]
+    int* per_head_vertical_topkv, int* per_head_slash_topkv,
+    int* block_count,   // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
+    int* block_offset,  // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_S]
+    int* column_count,  // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
+    int* column_index,  // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_V]
+    int64_t BATCH_SIZE, int64_t N_HEADS, int64_t N_ROWS, int64_t BLOCK_SIZE_M,
+    int64_t BLOCK_SIZE_N, int64_t NNZ_V, int64_t NNZ_S, bool causal) {
+  const int N_THREADS = 64;
+  const dim3 dimBlock(N_THREADS);
+  const dim3 dimGrid(N_HEADS, BATCH_SIZE, (N_ROWS + N_THREADS - 1) / N_THREADS);
+  convert_vertical_slash_indexes_kernel_mergehead<<<dimGrid, dimBlock>>>(
+      q_seqlens, kv_seqlens, vertical_indexes, slash_indexes,
+      per_head_vertical_topkv, per_head_slash_topkv, block_count, block_offset,
+      column_count, column_index, N_HEADS, N_ROWS, BLOCK_SIZE_M, BLOCK_SIZE_N,
+      NNZ_V, NNZ_S, causal);
+}
+
+/**
+ * Implements the Algorithm 4 in paper https://arxiv.org/abs/2407.02490.
+ *
+ * Like the above convert_vertical_slash_indexes, but with
+ * pre-computed vertical and slash counts.
+ */
+void convert_vertical_slash_indexes_mergehead(
+    torch::Tensor& block_count,            // [BATCH, N_HEADS, NUM_ROWS]
+    torch::Tensor& block_offset,           // [BATCH, N_HEADS, NUM_ROWS, NNZ_S]
+    torch::Tensor& column_count,           // [BATCH, N_HEADS, NUM_ROWS]
+    torch::Tensor& column_index,           // [BATCH, N_HEADS, NUM_ROWS, NNZ_V]
+    torch::Tensor q_seqlens,               // [BATCH, ]
+    torch::Tensor kv_seqlens,              // [BATCH, ]
+    torch::Tensor vertical_indexes,        // [BATCH, N_HEADS, NNZ_V]
+    torch::Tensor slash_indexes,           // [BATCH, N_HEADS, NNZ_S]
+    torch::Tensor vertical_indices_count,  // [N_HEADS, ]
+    torch::Tensor slash_indices_count,     // [N_HEADS, ]
+    int64_t context_size, int64_t block_size_M, int64_t block_size_N,
+    bool causal) {
+  cudaSetDevice(q_seqlens.get_device());
+
+  int batch_size = slash_indexes.size(0);
+  int num_heads = slash_indexes.size(1);
+  int nnz_slash = slash_indexes.size(2);
+  int nnz_vertical = vertical_indexes.size(2);
+  int num_rows = (context_size + block_size_M - 1) / block_size_M;
+
+  convert_vertical_slash_indexes_64x64_mergehead(
+      q_seqlens.data_ptr<int>(), kv_seqlens.data_ptr<int>(),
+      vertical_indexes.data_ptr<int>(), slash_indexes.data_ptr<int>(),
+      vertical_indices_count.data_ptr<int>(),
+      slash_indices_count.data_ptr<int>(), block_count.data_ptr<int>(),
+      block_offset.data_ptr<int>(), column_count.data_ptr<int>(),
+      column_index.data_ptr<int>(), batch_size, num_heads, num_rows,
+      block_size_M, block_size_N, nnz_vertical, nnz_slash, causal);
+}
@@ -12,8 +12,7 @@ void swap_blocks(torch::Tensor& src, torch::Tensor& dst,

 void swap_blocks_batch(const torch::Tensor& src_ptrs,
                       const torch::Tensor& dst_ptrs,
-                       const torch::Tensor& sizes,
-                       bool is_src_access_order_any);
+                       const torch::Tensor& sizes);

 void reshape_and_cache(torch::Tensor& key, torch::Tensor& value,
                       torch::Tensor& key_cache, torch::Tensor& value_cache,
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
Yongye ZhuandClaude Opus 4.7	58c8a5eaa5	[Attention][TokenSpeed MLA] Also warm up prefill kernel from decode impl The prefill backend may be paired with flash_attn / trtllm in production — in that case the prefill backend's __init__ never runs and the prefill kernel's first call pays a 1.5–2 minute JIT cost. Add the same idempotent `warmup_compile_prefill` invocation to TokenspeedMLAImpl.__init__ (the decode-side backend, always present when tokenspeed is selected). The function dedupes by config key, so the double call is a no-op when both backends are tokenspeed. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>	2026-05-06 07:45:16 +00:00
Yongye Zhu	c4547482ca	update version Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>	2026-05-06 07:28:20 +00:00
Yongye Zhu	91ef0afcb2	adding to cuda.txt dependency Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>	2026-05-06 07:20:48 +00:00
Yongye Zhu	97cd2c41ad	fix precommit Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>	2026-05-06 03:14:49 +00:00
Yongye ZhuandClaude Opus 4.7	c001535038	[Attention][TokenSpeed MLA] Force prefill V tensor contiguous before kernel call `v` arrives at both `run_prefill_new_tokens` and `run_prefill_context_chunk` as the second half of `kv_nope.split([qk_nope_head_dim, v_head_dim], dim=-1)` in mla_attention.py — a non-contiguous view along the last dim. The kernel internally does `v.reshape(1, total_kv, h_k, 1, d_v)` which silently copies when the input is non-contiguous; pull that copy out so the layout is predictable at the kernel boundary. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>	2026-05-06 02:53:46 +00:00
Yongye ZhuandClaude Opus 4.7	de6bc297df	[Attention][TokenSpeed MLA] Surface install hint when package missing Previously a user explicitly selecting TOKENSPEED_MLA without `tokenspeed_mla` installed got either a generic "required dependencies not available" message (prefill backend) or a raw ModuleNotFoundError deep inside forward_mqa at the first request (decode backend). Now both backends fail at startup with the exact install command: `uv pip install tokenspeed-mla`. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>	2026-05-06 02:49:09 +00:00
Yongye ZhuandClaude Opus 4.7	0012818287	[Attention][TokenSpeed MLA] Fix trtllm LSE parity test: log2 → natural log trtllm_ragged_attention_deepseek returns LSE in log2; tokenspeed and merge_attn_states use natural log. Multiply the trtllm reference by ln 2 before comparison. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>	2026-05-06 02:49:09 +00:00
Yongye ZhuandClaude Opus 4.7	73cd7e25ae	[Attention][TokenSpeed MLA] Warm up BF16 prefill compile, drop seq_lens computation Pre-JIT both BF16 and FP8 prefill kernels at backend init since the dtype isn't visible from `__init__` — depends on `use_prefill_query_quantization`. Move the per-forward `seq_lens` computation into `prepare_metadata` and document the cuda-graph padding interaction with `query_start_loc`. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>	2026-05-06 02:49:09 +00:00
Yongye ZhuandClaude Opus 4.7	964c6eb485	[Attention][TokenSpeed MLA] Fix decode FP8 numerics: pass output_scale and assert FP8 Q The decode kernel needs both bmm scales to recover correct outputs from an FP8 KV cache: bmm1 (softmax_scale = scale * q_scale * k_scale) and bmm2 (output_scale = k_scale, since V is stored as V_real / k_scale). We were only passing bmm1, which left bmm2 = 1.0 and produced silently wrong output. Also assert query dtype is float8_e4m3fn on entry to forward_mqa. supports_quant_query_input=True (inherited from MLACommonImpl) tells the upstream pipeline to FP8-quantize Q via _decode_concat_quant_fp8_op; the kernel is shape-specialized for FP8 Q + FP8 KV, so any other dtype here means the upstream quant path didn't run and the kernel will produce garbage. Failing loud beats failing silent. Verified: gsm8k matches reference with TOKENSPEED_MLA decode + FLASH_ATTN prefill on Kimi-K2.5-NVFP4 / TP=4 / B200. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>	2026-05-06 02:49:09 +00:00
Yongye ZhuandClaude Opus 4.7	d0e6514bf8	[Attention] Add TOKENSPEED_MLA backend for DeepSeek R1 prefill + decode on Blackwell Wires the tokenspeed_mla CuTe DSL kernels into vLLM as a new MLA backend, covering both prefill (tokenspeed_mla_prefill) and decode (tokenspeed_mla_decode). Targets Blackwell (SM100) with FP8 KV cache and DeepSeek R1 MLA dimensions; users opt in via -ac '{"backend":"TOKENSPEED_MLA","mla_prefill_backend":"TOKENSPEED_MLA"}'. Includes numeric parity tests against the trtllm reference kernels. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>	2026-05-06 02:49:09 +00:00