Init

Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
2026-07-13 07:16:33 +00:00
1585 changed files with 19604 additions and 141601 deletions
@@ -18,8 +18,6 @@ steps:
      TERM: "xterm-256color"
    retry:
      automatic:
-        - exit_status: 1  # Transient Docker/BuildKit failure
-          limit: 1
        - exit_status: -1  # Agent was lost
          limit: 1
        - exit_status: -10  # Agent was lost
@@ -48,8 +46,6 @@ steps:
      VLLM_BRANCH: "$BUILDKITE_COMMIT"
    retry:
      automatic:
-        - exit_status: 1  # Transient Docker/BuildKit failure
-          limit: 1
        - exit_status: -1  # Agent was lost
          limit: 1
        - exit_status: -10  # Agent was lost
@@ -76,8 +72,6 @@ steps:
      VLLM_BRANCH: "$BUILDKITE_COMMIT"
    retry:
      automatic:
-        - exit_status: 1  # Transient Docker/BuildKit failure
-          limit: 1
        - exit_status: -1  # Agent was lost
          limit: 1
        - exit_status: -10  # Agent was lost
@@ -18,8 +18,6 @@ steps:
  - tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
  - tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
  - tests/kernels/mamba/test_cpu_short_conv.py
-  - tests/kernels/mamba/test_causal_conv1d.py
-  - tests/kernels/mamba/test_mamba_ssm.py
  commands:
    - |
      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
@@ -30,9 +28,7 @@ steps:
      pytest -x -v -s tests/kernels/test_onednn.py
      pytest -x -v -s tests/kernels/test_awq_int4_to_int8.py
      pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
-      pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
-      pytest -x -v -s tests/kernels/mamba/test_causal_conv1d.py
-      pytest -x -v -s tests/kernels/mamba/test_mamba_ssm.py"
+      pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py"

 # Note: SDE can't be downloaded from CI host because of AWS WAF
 # - label: CPU-Compatibility Tests
@@ -145,22 +141,9 @@ steps:
  commands:
    - |
      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 45m "
-      pytest -x -v -s tests/models/multimodal/generation --ignore=tests/models/multimodal/generation/test_pixtral.py --ignore=tests/models/multimodal/generation/test_qwen2_5_vl.py -m cpu_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB"
+      pytest -x -v -s tests/models/multimodal/generation --ignore=tests/models/multimodal/generation/test_pixtral.py -m cpu_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB"
  parallelism: 4

- label: CPU-Qwen2.5-VL Multimodal Tests
-  depends_on: []
-  device: intel_cpu
-  no_plugin: true
-  source_file_dependencies:
-  # - vllm/
-  - vllm/model_executor/layers/rotary_embedding
-  - tests/models/multimodal/generation/
-  commands:
-    - |
-      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 40m "
-      VLLM_CI_ENV=0 pytest -x -v -s tests/models/multimodal/generation/test_qwen2_5_vl.py"
-
 - label: "Arm CPU Test"
  depends_on: []
  soft_fail: false
@@ -18,7 +18,7 @@ steps:
  - label: "XPU example Test"
    depends_on:
      - image-build-xpu
-    timeout_in_minutes: 50
+    timeout_in_minutes: 30
    optional: true
    device: intel_gpu
    agent_tags:
@@ -39,7 +39,7 @@ steps:
  - label: "XPU V1 test"
    depends_on:
      - image-build-xpu
-    timeout_in_minutes: 70
+    timeout_in_minutes: 30
    optional: true
    device: intel_gpu
    agent_tags:
@@ -60,7 +60,7 @@ steps:
  - label: "XPU server test"
    depends_on:
      - image-build-xpu
-    timeout_in_minutes: 45
+    timeout_in_minutes: 30
    optional: true
    device: intel_gpu
    agent_tags:
@@ -177,18 +177,6 @@ BRANCH=$4
 IMAGE_TAG=$5
 IMAGE_TAG_LATEST=${6:-} # only used for main branch, optional

-# When TORCH_NIGHTLY=1, build the base CI image against PyTorch nightly so the
-# entire existing pipeline runs on nightly torch (CUDA/GPU lane only). Delegate
-# to the dedicated nightly build (PYTORCH_NIGHTLY=1, CUDA 13.0) and tag it at the
-# normal IMAGE_TAG that every test step already pulls -- no separate image tag,
-# no duplicate "vLLM Against PyTorch Nightly" pipeline section.
-if [[ "${TORCH_NIGHTLY:-0}" == "1" ]]; then
-    echo "--- :warning: TORCH_NIGHTLY=1 -- building base image on PyTorch nightly"
-    SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
-    exec "${SCRIPT_DIR}/image_build_torch_nightly.sh" \
-        "${REGISTRY}" "${REPO}" "${BUILDKITE_COMMIT}" "${BRANCH}" "${IMAGE_TAG}"
-fi
-
 # build config
 TARGET="test-ci"
 VLLM_BAKE_FILE_PATH="${VLLM_BAKE_FILE_PATH:-docker/docker-bake.hcl}"
@@ -19,14 +19,13 @@ if docker manifest inspect "$IMAGE" >/dev/null 2>&1; then
  echo "Image found"
 else
  echo "Image not found, proceeding with build..."
-  # build for arm64 GPU targets: Grace/GH200 (sm_90),
-  # Blackwell/Thor (sm_100/sm_103/sm_110), and DGX Spark/GB10
+  # build for arm64 GPU targets: Grace/GH200 (sm_90) and DGX Spark/GB10
  # (sm_121, family-covered by 12.0 under CUDA 13)
  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 10.0 11.0 12.0" \
+    --build-arg torch_cuda_arch_list="9.0 12.0" \
    --build-arg USE_SCCACHE=1 \
    --build-arg buildkite_commit="$BUILDKITE_COMMIT" \
    --tag "$IMAGE" \
@@ -3,7 +3,7 @@ depends_on:
  - image-build-xpu
 steps:
 - label: XPU Sleep Mode
-  timeout_in_minutes: 45
+  timeout_in_minutes: 30
  device: intel_gpu
  agent_tags:
    label: production
@@ -86,7 +86,7 @@ steps:
      pytest -v -s lora/test_punica_ops.py::test_add_lora_fused_moe_early_exit'

 - label: LoRA Punica FP8/XPU Ops
-  timeout_in_minutes: 60
+  timeout_in_minutes: 45
  device: intel_gpu
  agent_tags:
    label: production
@@ -3,7 +3,7 @@ depends_on:
  - image-build-xpu
 steps:
 - label: V1 Core + KV + Metrics
-  timeout_in_minutes: 45
+  timeout_in_minutes: 30
  device: intel_gpu
  agent_tags:
    label: production
@@ -33,7 +33,7 @@ steps:
      pytest -v -s v1/executor'

 - label: V1 Sample + Logits
-  timeout_in_minutes: 90
+  timeout_in_minutes: 30
  device: intel_gpu
  agent_tags:
    label: production
@@ -72,7 +72,9 @@ steps:
      pytest -v -s v1/test_oracle.py &&
      pytest -v -s v1/test_request.py &&
      pytest -v -s v1/test_outputs.py &&
-      pytest -v -s v1/sample'
+      pytest -v -s v1/sample/test_topk_topp_sampler.py &&
+      pytest -v -s v1/sample/test_logprobs.py &&
+      pytest -v -s v1/sample/test_logprobs_e2e.py'

 - label: Basic Models Tests (Initialization)
  timeout_in_minutes: 60
@@ -152,7 +154,7 @@ steps:

 - label: Regression
  key: regression
-  timeout_in_minutes: 50
+  timeout_in_minutes: 30
  device: intel_gpu
  agent_tags:
    label: production
@@ -186,7 +188,7 @@ steps:

 - label: Metrics, Tracing (2 GPUs)
  key: metrics-tracing-2-gpus
-  timeout_in_minutes: 45
+  timeout_in_minutes: 30
  num_devices: 2
  device: intel_gpu
  agent_tags:
@@ -222,7 +224,7 @@ steps:

 - label: Async Engine, Inputs, Utils, Worker
  key: async-engine-inputs-utils-worker
-  timeout_in_minutes: 55
+  timeout_in_minutes: 30
  device: intel_gpu
  agent_tags:
    label: production
@@ -4,7 +4,7 @@ depends_on:
 steps:
 - label: Distributed Model Tests (2 GPUs)
  key: distributed-model-tests-2-gpus
-  timeout_in_minutes: 65
+  timeout_in_minutes: 50
  device: intel_gpu
  agent_tags:
    label: production
@@ -4,7 +4,7 @@ depends_on:
 steps:
 - label: "Multi-Modal Models (Standard) 1: qwen2"
  key: multi-modal-models-standard-1-qwen2
-  timeout_in_minutes: 70
+  timeout_in_minutes: 45
  device: intel_gpu
  agent_tags:
    label: production
@@ -29,7 +29,7 @@ steps:

 - label: "Multi-Modal Models (Standard) 2: qwen3 + gemma"
  key: multi-modal-models-standard-2-qwen3-gemma
-  timeout_in_minutes: 70
+  timeout_in_minutes: 45
  device: intel_gpu
  agent_tags:
    label: production
@@ -52,7 +52,7 @@ steps:

 - label: "Multi-Modal Models (Standard) 3: llava + qwen2_vl"
  key: multi-modal-models-standard-3-llava-qwen2-vl
-  timeout_in_minutes: 65
+  timeout_in_minutes: 45
  device: intel_gpu
  agent_tags:
    label: production
@@ -100,7 +100,7 @@ steps:

 - label: Multi-Modal Processor # 44min
  key: multi-modal-processor
-  timeout_in_minutes: 60
+  timeout_in_minutes: 45
  device: intel_gpu
  agent_tags:
    label: production
@@ -17,7 +17,7 @@ steps:
  - label: "XPU example Test"
    depends_on:
      - image-build-xpu
-    timeout_in_minutes: 50
+    timeout_in_minutes: 30
    device: intel_gpu
    agent_tags:
      label: production
@@ -76,7 +76,7 @@ steps:
  - label: "XPU V1 test"
    depends_on:
      - image-build-xpu
-    timeout_in_minutes: 70
+    timeout_in_minutes: 30
    device: intel_gpu
    agent_tags:
      label: production
@@ -99,13 +99,12 @@ steps:
        pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py --ignore=v1/worker/test_worker_memory_snapshot.py &&
        pytest -v -s v1/structured_output &&
        pytest -v -s v1/test_serial_utils.py &&
-        pytest -v -s v1/e2e/general/test_correctness_sliding_window.py --deselect="tests/v1/e2e/general/test_correctness_sliding_window.py::test_sliding_window_retrieval[True-1-5-google/gemma-3-1b-it]" &&
        pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py --ignore=v1/spec_decode/test_speculators_correctness.py &&
        pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py --ignore=v1/kv_connector/unit/test_offloading_connector.py'
  - label: "XPU server test"
    depends_on:
      - image-build-xpu
-    timeout_in_minutes: 45
+    timeout_in_minutes: 30
    device: intel_gpu
    agent_tags:
      label: production
@@ -1,7 +1,6 @@
 # For hf script, without -t option (tensor parallel size).
 # bash .buildkite/lm-eval-harness/run-lm-eval-mmlupro-vllm-baseline.sh -m meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8 -l 250 -t 8 -f 5
 model_name: "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8"
-rocm_safetensors_load_strategy: lazy
 required_gpu_arch:
  - gfx942
  - gfx950
@@ -72,11 +72,6 @@ def launch_lm_eval(eval_config, tp_size):
    if moe_backend is not None:
        model_args += f"moe_backend={moe_backend},"

-    if current_platform.is_rocm():
-        rocm_load_strategy = eval_config.get("rocm_safetensors_load_strategy")
-        if rocm_load_strategy is not None:
-            model_args += f"safetensors_load_strategy={rocm_load_strategy},"
-
    env_vars = eval_config.get("env_vars", None)
    with scoped_env_vars(env_vars):
        results = lm_eval.simple_evaluate(
@@ -3,8 +3,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 #
 # Append a build artifact line to the Buildkite annotation.
-# Usage: annotate-build-artifact.sh <label> <value> <context>
+# Usage: annotate-build-artifact.sh <label> <value>
 set -e
 echo "- **${1}**: \`${2}\`" | \
-  buildkite-agent annotate --append --style 'info' \
-    --context "${3:?context is required}"
+  buildkite-agent annotate --append --style 'info' --context 'release-artifacts'
@@ -1,32 +0,0 @@
-#!/usr/bin/env bash
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-#
-# Build the macOS arm64 CPU wheel natively on a macOS agent (the `macmini`
-# queue) into artifacts/dist/ for upload-nightly-wheels.sh.
-
-set -euo pipefail
-
-# The Rust frontend build needs protoc.
-if ! command -v protoc >/dev/null 2>&1; then
-  brew install protobuf
-fi
-
-# upload-nightly-wheels.sh expects exactly one wheel.
-rm -rf artifacts/dist
-mkdir -p artifacts/dist
-
-export VLLM_TARGET_DEVICE=cpu
-export VLLM_REQUIRE_RUST_FRONTEND=1
-export MACOSX_DEPLOYMENT_TARGET=11.0
-# uv's CPython is universal2; force an arm64-only build and tag so the wheel
-# isn't mislabelled universal2 and installed on Intel Macs where import fails.
-export ARCHFLAGS="-arch arm64"
-export _PYTHON_HOST_PLATFORM="macosx-11.0-arm64"
-export CMAKE_BUILD_PARALLEL_LEVEL="${CMAKE_BUILD_PARALLEL_LEVEL:-4}"
-
-uv venv --python 3.12
-uv pip install -r requirements/build/cpu.txt --index-strategy unsafe-best-match
-uv build --wheel --no-build-isolation -o artifacts/dist
-
-ls -l artifacts/dist/*.whl
@@ -29,11 +29,7 @@ if python3 -c "import torch; assert torch.version.hip" 2>/dev/null; then
        TORCH_INDEX_URL=""
    fi
 else
-    if [ "${TORCH_NIGHTLY:-0}" = "1" ]; then
-        TORCH_INDEX_URL="https://download.pytorch.org/whl/nightly/cu130"
-    else
-        TORCH_INDEX_URL="https://download.pytorch.org/whl/cu130"
-    fi
+    TORCH_INDEX_URL="https://download.pytorch.org/whl/cu130"
 fi
 echo ">>> Using PyTorch index: ${TORCH_INDEX_URL:-PyPI default}"

@@ -15,9 +15,9 @@ set -euo pipefail

 DEFAULT_REPO_SLUG="vllm-project/vllm"
 DEFAULT_CI_HCL_SOURCE="docker/ci-rocm.hcl"
-DEFAULT_CI_BASE_CONTENT_FILES="requirements/common.txt requirements/rocm.txt requirements/test/rocm.txt docker/Dockerfile.rocm_base docker/ci-rocm.hcl docker/docker-bake-rocm.hcl tools/install_torchcodec_rocm.sh tools/install_protoc.sh rust-toolchain.toml tests/vllm_test_utils .buildkite/scripts/ci-bake-rocm.sh .buildkite/scripts/rocm/build-ci-base.sh"
+DEFAULT_CI_BASE_CONTENT_FILES="requirements/common.txt requirements/rocm.txt requirements/test/rocm.txt docker/Dockerfile.rocm_base docker/ci-rocm.hcl docker/docker-bake-rocm.hcl tools/install_torchcodec_rocm.sh tests/vllm_test_utils .buildkite/scripts/ci-bake-rocm.sh .buildkite/scripts/rocm/build-ci-base.sh"
 DEFAULT_CI_BASE_DOCKERFILE="docker/Dockerfile.rocm"
-DEFAULT_CI_BASE_DOCKERFILE_STAGES="base rust_toolchain_input_0 rust_toolchain_input_1 rust-toolchain-input rust-toolchain build_nixl build_rocshmem build_deepep mori_base ci_base"
+DEFAULT_CI_BASE_DOCKERFILE_STAGES="base build_rixl build_rocshmem build_deepep mori_base ci_base"
 DEFAULT_CI_BASE_METADATA_VERSION="1"
 IMAGE_EXISTED_BEFORE_BUILD=0

@@ -285,7 +285,7 @@ get_content_arg_names() {
    fi | awk 'NF && !seen[$0]++'
 }

-compute_ci_base_content_hash_once() {
+compute_ci_base_content_hash() {
    local -a content_paths=()
    local -a content_args=()
    local dockerfile="${CI_BASE_DOCKERFILE:-}"
@@ -301,8 +301,7 @@ compute_ci_base_content_hash_once() {
        if [[ -n "${dockerfile}" ]]; then
            printf 'dockerfile:%s\n' "${dockerfile}"
            printf 'resolved-build-args:\n'
-            hash_dockerfile_arg_values "${dockerfile}" "${content_args[@]}" \
-                || return 1
+            hash_dockerfile_arg_values "${dockerfile}" "${content_args[@]}"
            if [[ -n "${stages}" ]]; then
                printf 'dockerfile-stages:%s\n' "${stages}"
                if [[ -f "${dockerfile}" ]]; then
@@ -315,53 +314,6 @@ compute_ci_base_content_hash_once() {
    } | sha256sum | cut -d' ' -f1
 }

-compute_ci_base_content_hash() {
-    local attempts="${CI_BASE_HASH_ATTEMPTS:-3}"
-    local delay_secs="${CI_BASE_HASH_RETRY_DELAY:-5}"
-    local attempt=0
-    local hash=""
-    local failed=0
-    local -a hashes=()
-
-    if [[ ! "${attempts}" =~ ^[1-9][0-9]*$ ]]; then
-        echo "Invalid CI_BASE_HASH_ATTEMPTS: ${attempts}" >&2
-        return 1
-    fi
-    if [[ ! "${delay_secs}" =~ ^[0-9]+$ ]]; then
-        echo "Invalid CI_BASE_HASH_RETRY_DELAY: ${delay_secs}" >&2
-        return 1
-    fi
-
-    for ((attempt = 1; attempt <= attempts; attempt++)); do
-        if ! hash=$(compute_ci_base_content_hash_once); then
-            echo "ci_base content hash calculation ${attempt}/${attempts} failed" >&2
-            failed=1
-        else
-            hashes+=("${hash}")
-            echo "ci_base content hash calculation ${attempt}/${attempts}: ${hash}" >&2
-        fi
-
-        if ((attempt < attempts)); then
-            sleep "${delay_secs}"
-        fi
-    done
-
-    if ((failed)) || ((${#hashes[@]} != attempts)); then
-        echo "Could not calculate a reliable ci_base content hash" >&2
-        return 1
-    fi
-
-    for hash in "${hashes[@]:1}"; do
-        if [[ "${hash}" != "${hashes[0]}" ]]; then
-            echo "ci_base content hash changed between calculations" >&2
-            printf '  observed: %s\n' "${hashes[@]}" >&2
-            return 1
-        fi
-    done
-
-    printf '%s\n' "${hashes[0]}"
-}
-
 extract_dockerfile_arg_default() {
    local dockerfile="$1"
    local arg_name="$2"
@@ -414,11 +366,7 @@ hash_dockerfile_arg_values() {
        printf 'arg:%s=%s\n' "${arg_name}" "${arg_value:-<empty>}"
        if [[ "${arg_name}" == "BASE_IMAGE" && -n "${arg_value}" ]]; then
            digest=$(resolve_image_digest "${arg_value}")
-            if [[ -z "${digest}" ]]; then
-                echo "Failed to resolve digest for BASE_IMAGE=${arg_value}" >&2
-                return 1
-            fi
-            printf 'arg:%s.digest=%s\n' "${arg_name}" "${digest}"
+            printf 'arg:%s.digest=%s\n' "${arg_name}" "${digest:-unknown}"
        fi
    done
 }
@@ -816,7 +764,7 @@ configure_ci_base_image_refs() {
        fi
        set_buildkite_metadata "rocm-ci-base-image" "${CI_BASE_IMAGE_TAG}"
        set_buildkite_metadata "rocm-ci-base-image-content" "${content_tag}"
-        set_buildkite_metadata "rocm-ci-base-image-commit" "${CI_BASE_IMAGE_TAG_COMMIT_REF:-}"
+        set_buildkite_metadata "rocm-ci-base-image-commit" "${CI_BASE_IMAGE_TAG_COMMIT:-}"
        set_buildkite_metadata "rocm-ci-base-image-stable" "${CI_BASE_IMAGE_TAG_STABLE:-}"
        return 0
    fi
@@ -1159,8 +1107,8 @@ ci_base_metadata_pairs() {
    metadata_pair "vllm.rocm.nic_backend" "$(resolve_dockerfile_arg_value "${dockerfile}" "NIC_BACKEND")"
    metadata_pair "vllm.rocm.ainic_version" "$(resolve_dockerfile_arg_value "${dockerfile}" "AINIC_VERSION")"
    metadata_pair "vllm.rocm.ubuntu_codename" "$(resolve_dockerfile_arg_value "${dockerfile}" "UBUNTU_CODENAME")"
-    metadata_pair "vllm.rocm.nixl_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "NIXL_REPO")"
-    metadata_pair "vllm.rocm.nixl_commit" "${NIXL_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "NIXL_BRANCH")}"
+    metadata_pair "vllm.rocm.rixl_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "RIXL_REPO")"
+    metadata_pair "vllm.rocm.rixl_commit" "${RIXL_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "RIXL_BRANCH")}"
    metadata_pair "vllm.rocm.ucx_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "UCX_REPO")"
    metadata_pair "vllm.rocm.ucx_commit" "${UCX_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "UCX_BRANCH")}"
    metadata_pair "vllm.rocm.rocshmem_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "ROCSHMEM_REPO")"
@@ -1169,7 +1117,7 @@ ci_base_metadata_pairs() {
    metadata_pair "vllm.rocm.deepep_commit" "${DEEPEP_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_BRANCH")}"
    metadata_pair "vllm.rocm.deepep_nic" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_NIC")"
    metadata_pair "vllm.rocm.deepep_rocm_arch" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_ROCM_ARCH")"
-    metadata_pair "vllm.rocm.nixl_cache_key" "${NIXL_CACHE_KEY:-}"
+    metadata_pair "vllm.rocm.rixl_cache_key" "${RIXL_CACHE_KEY:-}"
    metadata_pair "vllm.rocm.rocshmem_cache_key" "${ROCSHMEM_CACHE_KEY:-}"
    metadata_pair "vllm.rocm.deepep_cache_key" "${DEEPEP_CACHE_KEY:-}"

@@ -1263,24 +1211,12 @@ uses_rocm_csrc_cache() {
    esac
 }

-uses_rocm_rust_cache() {
-    case "${TARGET}" in
-        rust-rocm-ci|test-rocm-ci|test-rocm-ci-with-wheel|test-rocm-ci-with-artifacts|export-wheel-rocm)
-            return 0
-            ;;
-        *)
-            return 1
-            ;;
-    esac
-}
-
 compute_rocm_csrc_content_hash() {
    local bake_dir=""
    local dockerfile_rocm=""
    local -a content_paths=(
        "requirements/common.txt"
        "requirements/rocm.txt"
-        "pyproject.toml"
        "setup.py"
        "CMakeLists.txt"
        "cmake"
@@ -1324,56 +1260,6 @@ compute_rocm_csrc_content_hash_if_needed() {
    echo "ROCm csrc content cache ref: ${ROCM_CSRC_CONTENT_CACHE_REF}"
 }

-compute_rocm_rust_content_hash() {
-    local bake_dir=""
-    local dockerfile_rocm=""
-    local -a content_paths=(
-        "requirements/build/rust.txt"
-        "rust/Cargo.lock"
-        "rust/Cargo.toml"
-        "rust/proto"
-        "rust/src"
-        "rust-toolchain.toml"
-        "tools/build_rust.py"
-        "tools/install_protoc.sh"
-        "build_rust.sh"
-    )
-    local -a content_args=()
-
-    bake_dir=$(dirname "${VLLM_BAKE_FILE}")
-    dockerfile_rocm="${bake_dir}/Dockerfile.rocm"
-    mapfile -t content_args < <(
-        get_content_arg_names "${dockerfile_rocm}" "base rust_toolchain_input_0 rust_toolchain_input_1 rust-toolchain-input rust_input_0 rust_input_1 rust-input rust-toolchain rust-build" "${ROCM_RUST_CONTENT_ARGS:-}"
-    )
-
-    {
-        printf 'rust-input-files-hash:%s\n' "$(compute_content_hash "${content_paths[@]}")"
-        printf 'dockerfile:%s\n' "${dockerfile_rocm}"
-        printf 'resolved-build-args:\n'
-        hash_dockerfile_arg_values "${dockerfile_rocm}" "${content_args[@]}"
-        printf 'dockerfile-stages:base rust_toolchain_input_0 rust_toolchain_input_1 rust-toolchain-input rust_input_0 rust_input_1 rust-input rust-toolchain rust-build\n'
-        if [[ -f "${dockerfile_rocm}" ]]; then
-            hash_dockerfile_stages "${dockerfile_rocm}" "base rust_toolchain_input_0 rust_toolchain_input_1 rust-toolchain-input rust_input_0 rust_input_1 rust-input rust-toolchain rust-build"
-        else
-            printf 'missing:%s\n' "${dockerfile_rocm}"
-        fi
-    } | sha256sum | cut -d' ' -f1
-}
-
-compute_rocm_rust_content_hash_if_needed() {
-    local cache_repo="${DOCKERHUB_CACHE_REPO:-rocm/vllm-ci-cache}"
-
-    if [[ "${ROCM_RUST_CONTENT_CACHE:-1}" == "0" ]] || ! uses_rocm_rust_cache; then
-        return 0
-    fi
-
-    ROCM_RUST_CONTENT_HASH=$(compute_rocm_rust_content_hash)
-    ROCM_RUST_CONTENT_CACHE_REF="${cache_repo}:rust-rocm-input-${ROCM_RUST_CONTENT_HASH}"
-    export ROCM_RUST_CONTENT_HASH
-    export ROCM_RUST_CONTENT_CACHE_REF
-    echo "ROCm Rust content cache ref: ${ROCM_RUST_CONTENT_CACHE_REF}"
-}
-
 write_hcl_string_list_entries() {
    local indent="$1"
    local value=""
@@ -1431,7 +1317,6 @@ write_rocm_build_arg_override() {
                "${CI_BASE_DOCKERFILE_STAGES:-${DEFAULT_CI_BASE_DOCKERFILE_STAGES}}" \
                "${CI_BASE_CONTENT_ARGS:-}"
            get_content_arg_names "${dockerfile_rocm}" "base csrc-build" "${ROCM_CSRC_CONTENT_ARGS:-}"
-            get_content_arg_names "${dockerfile_rocm}" "base rust_toolchain_input_0 rust_toolchain_input_1 rust-toolchain-input rust_input_0 rust_input_1 rust-input rust-toolchain rust-build" "${ROCM_RUST_CONTENT_ARGS:-}"
        } | awk 'NF && !seen[$0]++'
    )

@@ -1480,133 +1365,46 @@ validate_cache_export_mode() {
    esac
 }

-validate_content_cache_export_mode() {
-    local mode="$1"
-    local env_name="$2"
-
-    case "${mode}" in
-        missing|always|never)
-            ;;
-        *)
-            echo "Error: ${env_name} must be one of: missing, always, never"
-            exit 1
-            ;;
-    esac
-}
-
-should_export_content_cache_ref() {
-    local cache_ref="$1"
-    local cache_name="$2"
-    local mode="${ROCM_CONTENT_CACHE_EXPORT_MODE:-missing}"
-
-    case "${mode}" in
-        always)
-            echo "${cache_name} content cache export mode is always; exporting ${cache_ref}"
-            return 0
-            ;;
-        never)
-            echo "${cache_name} content cache export mode is never; not exporting ${cache_ref}"
-            return 1
-            ;;
-        missing|"")
-            if docker buildx imagetools inspect "${cache_ref}" >/dev/null 2>&1; then
-                echo "${cache_name} content cache exists; not re-exporting ${cache_ref}"
-                return 1
-            fi
-            echo "${cache_name} content cache missing; will export ${cache_ref}"
-            return 0
-            ;;
-        *)
-            echo "Error: ROCM_CONTENT_CACHE_EXPORT_MODE must be one of: missing, always, never"
-            exit 1
-            ;;
-    esac
-}
-
 write_rocm_cache_override() {
    local cache_repo="${DOCKERHUB_CACHE_REPO:-rocm/vllm-ci-cache}"
-    local content_cache_export_mode="${ROCM_CONTENT_CACHE_EXPORT_MODE:-missing}"
    local csrc_cache_to_mode="${ROCM_CSRC_CACHE_TO_MODE:-max}"
-    local rust_cache_to_mode="${ROCM_RUST_CACHE_TO_MODE:-max}"
    local rocm_cache_to_mode="${ROCM_FINAL_CACHE_TO_MODE:-min}"
-    local -a csrc_content_cache_from=()
-    local -a rust_content_cache_from=()
-    local -a combined_content_cache_from=()
+    local -a content_cache_from=()
    local -a csrc_cache_to=()
-    local -a rust_cache_to=()
    local -a rocm_cache_to=()
    local -a export_wheel_cache_to=()
-    local export_csrc_cache=1
-    local export_rust_cache=1

-    if ! uses_rocm_csrc_cache && ! uses_rocm_rust_cache; then
+    if ! uses_rocm_csrc_cache; then
        return 0
    fi

-    validate_content_cache_export_mode \
-        "${content_cache_export_mode}" \
-        "ROCM_CONTENT_CACHE_EXPORT_MODE"
    validate_cache_export_mode "${csrc_cache_to_mode}" "ROCM_CSRC_CACHE_TO_MODE"
-    validate_cache_export_mode "${rust_cache_to_mode}" "ROCM_RUST_CACHE_TO_MODE"
    validate_cache_export_mode "${rocm_cache_to_mode}" "ROCM_FINAL_CACHE_TO_MODE"
-    echo "ROCm content cache export mode: ${content_cache_export_mode}"
    echo "ROCm csrc cache export mode: ${csrc_cache_to_mode}"
-    echo "ROCm Rust cache export mode: ${rust_cache_to_mode}"
    echo "ROCm final image cache export mode: ${rocm_cache_to_mode}"

    if [[ -n "${ROCM_CSRC_CONTENT_CACHE_REF:-}" ]]; then
-        csrc_content_cache_from+=("type=registry,ref=${ROCM_CSRC_CONTENT_CACHE_REF}")
-        if should_export_content_cache_ref "${ROCM_CSRC_CONTENT_CACHE_REF}" "ROCm csrc"; then
-            csrc_cache_to+=(
-                "type=registry,ref=${ROCM_CSRC_CONTENT_CACHE_REF},mode=${csrc_cache_to_mode},ignore-error=true"
-            )
-        else
-            export_csrc_cache=0
-        fi
+        content_cache_from+=("type=registry,ref=${ROCM_CSRC_CONTENT_CACHE_REF}")
+        csrc_cache_to+=(
+            "type=registry,ref=${ROCM_CSRC_CONTENT_CACHE_REF},mode=${csrc_cache_to_mode},ignore-error=true"
+        )
    fi

-    if [[ -n "${ROCM_RUST_CONTENT_CACHE_REF:-}" ]]; then
-        rust_content_cache_from+=("type=registry,ref=${ROCM_RUST_CONTENT_CACHE_REF}")
-        if should_export_content_cache_ref "${ROCM_RUST_CONTENT_CACHE_REF}" "ROCm Rust"; then
-            rust_cache_to+=(
-                "type=registry,ref=${ROCM_RUST_CONTENT_CACHE_REF},mode=${rust_cache_to_mode},ignore-error=true"
-            )
-        else
-            export_rust_cache=0
-        fi
-    fi
-
-    combined_content_cache_from=("${csrc_content_cache_from[@]}" "${rust_content_cache_from[@]}")
-
    # Docker Hub cache exports are best-effort. A cache-only target failure can
    # otherwise cancel the sibling image target before its manifest is pushed.
    if [[ -n "${BUILDKITE_COMMIT:-}" ]]; then
-        if [[ ${export_csrc_cache} -eq 1 ]]; then
-            csrc_cache_to+=(
-                "type=registry,ref=${cache_repo}:csrc-rocm-${BUILDKITE_COMMIT},mode=${csrc_cache_to_mode},ignore-error=true"
-            )
-        fi
-        if [[ ${export_rust_cache} -eq 1 ]]; then
-            rust_cache_to+=(
-                "type=registry,ref=${cache_repo}:rust-rocm-${BUILDKITE_COMMIT},mode=${rust_cache_to_mode},ignore-error=true"
-            )
-        fi
+        csrc_cache_to+=(
+            "type=registry,ref=${cache_repo}:csrc-rocm-${BUILDKITE_COMMIT},mode=${csrc_cache_to_mode},ignore-error=true"
+        )
        rocm_cache_to+=(
            "type=registry,ref=${cache_repo}:rocm-${BUILDKITE_COMMIT},mode=${rocm_cache_to_mode},ignore-error=true"
        )
    fi

    if [[ -n "${ROCM_CACHE_BRANCH_TAG:-}" ]]; then
-        if [[ ${export_csrc_cache} -eq 1 ]]; then
-            csrc_cache_to+=(
-                "type=registry,ref=${cache_repo}:csrc-rocm-branch-${ROCM_CACHE_BRANCH_TAG},mode=${csrc_cache_to_mode},ignore-error=true"
-            )
-        fi
-        if [[ ${export_rust_cache} -eq 1 ]]; then
-            rust_cache_to+=(
-                "type=registry,ref=${cache_repo}:rust-rocm-branch-${ROCM_CACHE_BRANCH_TAG},mode=${rust_cache_to_mode},ignore-error=true"
-            )
-        fi
+        csrc_cache_to+=(
+            "type=registry,ref=${cache_repo}:csrc-rocm-branch-${ROCM_CACHE_BRANCH_TAG},mode=${csrc_cache_to_mode},ignore-error=true"
+        )
        rocm_cache_to+=(
            "type=registry,ref=${cache_repo}:rocm-branch-${ROCM_CACHE_BRANCH_TAG},mode=${rocm_cache_to_mode},ignore-error=true"
        )
@@ -1624,7 +1422,7 @@ target "csrc-rocm-ci" {
  cache-from = concat(
    get_cache_from_rocm_csrc(),
 EOF
-        write_hcl_string_list "    " "${csrc_content_cache_from[@]}"
+        write_hcl_string_list "    " "${content_cache_from[@]}"
        cat <<EOF
  )
 EOF
@@ -1632,23 +1430,11 @@ EOF
        cat <<EOF
 }

-target "rust-rocm-ci" {
-  cache-from = concat(
-    get_cache_from_rocm_rust(),
-EOF
-        write_hcl_string_list "    " "${rust_content_cache_from[@]}"
-        cat <<EOF
-  )
-EOF
-        write_hcl_string_list_attr "  " "cache-to" "${rust_cache_to[@]}"
-        cat <<EOF
-}
-
 target "test-rocm-ci" {
  cache-from = concat(
    get_cache_from_rocm(),
 EOF
-        write_hcl_string_list "    " "${combined_content_cache_from[@]}"
+        write_hcl_string_list "    " "${content_cache_from[@]}"
        cat <<EOF
  )
 EOF
@@ -1660,7 +1446,7 @@ target "export-wheel-rocm" {
  cache-from = concat(
    get_cache_from_rocm(),
 EOF
-        write_hcl_string_list "    " "${combined_content_cache_from[@]}"
+        write_hcl_string_list "    " "${content_cache_from[@]}"
        cat <<EOF
  )
 EOF
@@ -1686,7 +1472,7 @@ extract_dependency_pins() {
        return 0
    fi

-    for var in NIXL_BRANCH UCX_BRANCH ROCSHMEM_BRANCH DEEPEP_BRANCH; do
+    for var in RIXL_BRANCH UCX_BRANCH ROCSHMEM_BRANCH DEEPEP_BRANCH; do
        if [[ -n "${!var:-}" ]]; then
            echo "Using provided ${var}: ${!var}"
            continue
@@ -1706,30 +1492,30 @@ extract_dependency_pins() {
 compute_dependency_cache_keys() {
    local bake_dir=""
    local dockerfile_rocm=""
-    local nixl_branch=""
+    local rixl_branch=""
    local ucx_branch=""
    local rocshmem_branch=""
    local deepep_branch=""
-    local nixl_material=""
+    local rixl_material=""
    local rocshmem_material=""
    local deepep_material=""

    bake_dir=$(dirname "${VLLM_BAKE_FILE}")
    dockerfile_rocm="${bake_dir}/Dockerfile.rocm"
-    nixl_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "NIXL_BRANCH")
+    rixl_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "RIXL_BRANCH")
    ucx_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "UCX_BRANCH")
    rocshmem_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "ROCSHMEM_BRANCH")
    deepep_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "DEEPEP_BRANCH")

-    if [[ -n "${nixl_branch}" && -n "${ucx_branch}" ]]; then
-        nixl_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_nixl")
-        NIXL_CACHE_KEY=$(
+    if [[ -n "${rixl_branch}" && -n "${ucx_branch}" ]]; then
+        rixl_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_rixl")
+        RIXL_CACHE_KEY=$(
            compose_dependency_cache_key \
-                "${nixl_branch}-ucx-${ucx_branch}" \
-                "${nixl_material}"
+                "${rixl_branch}-ucx-${ucx_branch}" \
+                "${rixl_material}"
        )
-        export NIXL_CACHE_KEY
-        echo "NIXL dependency cache key: ${NIXL_CACHE_KEY}"
+        export RIXL_CACHE_KEY
+        echo "RIXL dependency cache key: ${RIXL_CACHE_KEY}"
    fi

    if [[ -n "${rocshmem_branch}" ]]; then
@@ -1780,11 +1566,11 @@ dependency_cache_ref_for_target() {
    local cache_repo="${DOCKERHUB_CACHE_REPO:-rocm/vllm-ci-cache}"

    case "${target}" in
-        nixl-rocm-ci)
-            if [[ -n "${NIXL_CACHE_KEY:-}" ]]; then
-                printf '%s\n' "${cache_repo}:nixl-rocm-${NIXL_CACHE_KEY}"
-            elif [[ -n "${NIXL_BRANCH:-}" ]]; then
-                printf '%s\n' "${cache_repo}:nixl-rocm-${NIXL_BRANCH}-ucx-${UCX_BRANCH:-}"
+        rixl-rocm-ci)
+            if [[ -n "${RIXL_CACHE_KEY:-}" ]]; then
+                printf '%s\n' "${cache_repo}:rixl-rocm-${RIXL_CACHE_KEY}"
+            elif [[ -n "${RIXL_BRANCH:-}" ]]; then
+                printf '%s\n' "${cache_repo}:rixl-rocm-${RIXL_BRANCH}-ucx-${UCX_BRANCH:-}"
            fi
            ;;
        rocshmem-rocm-ci)
@@ -1815,7 +1601,7 @@ add_dependency_cache_target() {

 resolve_ci_base_dependency_targets() {
    local mode="${ROCM_DEP_CACHE_EXPORT_MODE:-missing}"
-    local nixl_ref=""
+    local rixl_ref=""
    local rocshmem_ref=""
    local deepep_ref=""

@@ -1824,7 +1610,7 @@ resolve_ci_base_dependency_targets() {
    case "${mode}" in
        always)
            echo "ROCM_DEP_CACHE_EXPORT_MODE=always; exporting all dependency caches serially"
-            for target in nixl-rocm-ci rocshmem-rocm-ci deepep-rocm-ci; do
+            for target in rixl-rocm-ci rocshmem-rocm-ci deepep-rocm-ci; do
                if [[ -n "$(dependency_cache_ref_for_target "${target}")" ]]; then
                    add_dependency_cache_target "${target}"
                fi
@@ -1844,13 +1630,13 @@ resolve_ci_base_dependency_targets() {
            ;;
    esac

-    if [[ "${mode}" != "always" && -n "${NIXL_CACHE_KEY:-}" ]]; then
-        nixl_ref=$(dependency_cache_ref_for_target "nixl-rocm-ci")
-        if dependency_cache_ref_exists "${nixl_ref}"; then
-            echo "NIXL dependency cache exists: ${nixl_ref}"
+    if [[ "${mode}" != "always" && -n "${RIXL_CACHE_KEY:-}" ]]; then
+        rixl_ref=$(dependency_cache_ref_for_target "rixl-rocm-ci")
+        if dependency_cache_ref_exists "${rixl_ref}"; then
+            echo "RIXL dependency cache exists: ${rixl_ref}"
        else
-            echo "NIXL dependency cache missing; will seed: ${nixl_ref}"
-            add_dependency_cache_target "nixl-rocm-ci"
+            echo "RIXL dependency cache missing; will seed: ${rixl_ref}"
+            add_dependency_cache_target "rixl-rocm-ci"
        fi
    fi

@@ -1950,8 +1736,8 @@ confirm_remote_image_push() {
        fi

        if [[ -z "${remote_revision}" \
-              && ${IMAGE_EXISTED_BEFORE_BUILD} -eq 0 ]] \
-            && image_tag_is_commit_scoped; then
+              && ${IMAGE_EXISTED_BEFORE_BUILD} -eq 0 \
+              && image_tag_is_commit_scoped ]]; then
            echo "Remote image exists under a commit-scoped tag; accepting push despite missing revision label."
            return 0
        fi
@@ -2081,57 +1867,36 @@ upload_wheel_artifacts_if_present() {
    local wheel_dir="./wheel-export"
    local artifact_dir="artifacts/vllm-rocm-install"
    local archive_name="vllm-rocm-install.tar.gz"
-    local metadata_dir="${wheel_dir}/.vllm-ci-artifact"
-    local native_base_image=""
    local whl=""
    local whl_name=""
-    local -a wheels=()

    if ! should_upload_wheel_artifacts; then
        return 0
    fi

-    if [[ -d "${wheel_dir}" ]]; then
-        mapfile -t wheels < <(find "${wheel_dir}" -maxdepth 1 -type f -name '*.whl' -print)
-    fi
-    if [[ ${#wheels[@]} -ne 1 ]]; then
-        echo "Expected exactly one ROCm wheel in ${wheel_dir}; found ${#wheels[@]}" >&2
-        return 1
-    fi
-    whl="${wheels[0]}"
-    whl_name=$(basename "${whl}")
-    native_base_image="${CI_BASE_IMAGE_TAG_COMMIT_REF:-${CI_BASE_IMAGE:-}}"
-    if [[ -z "${native_base_image}" ]]; then
-        echo "Native ROCm artifact requires a ci_base image reference" >&2
-        return 1
+    if [[ ! -d "${wheel_dir}" ]] || ! ls "${wheel_dir}"/*.whl >/dev/null 2>&1; then
+        echo "No ROCm wheel artifacts found in ${wheel_dir}"
+        return 0
    fi

    echo "--- :package: Uploading ROCm vLLM install artifact"
-    rm -rf "${artifact_dir}" "${metadata_dir}"
-    mkdir -p "${artifact_dir}" "${metadata_dir}"
-
-    printf '%s\n' "${BUILDKITE_COMMIT:-local}" > "${metadata_dir}/commit.txt"
-    printf '%s\n' "${native_base_image}" > "${metadata_dir}/native-base-image.txt"
-    printf '%s\n' "${CI_BASE_IMAGE:-}" > "${metadata_dir}/ci-base-image.txt"
-    printf '%s\n' "${IMAGE_TAG:-}" > "${metadata_dir}/fallback-image.txt"
-    printf '%s\n' "${whl_name}" > "${metadata_dir}/wheel-filename.txt"
+    mkdir -p "${artifact_dir}"

    tar -C "${wheel_dir}" -czf "${artifact_dir}/${archive_name}" .
-    (
-        cd "${artifact_dir}"
-        sha256sum "${archive_name}" > "${archive_name}.sha256"
-    )
    echo "Created ${archive_name}: $(du -sh "${artifact_dir}/${archive_name}" | cut -f1)"
-    cp "${metadata_dir}"/*.txt "${artifact_dir}/"
-    cp "${whl}" "${artifact_dir}/${whl_name}"
-    echo "Copied ${whl_name}: $(du -sh "${artifact_dir}/${whl_name}" | cut -f1)"
+    printf '%s\n' "${CI_BASE_IMAGE:-}" > "${artifact_dir}/ci-base-image.txt"
+    printf '%s\n' "${IMAGE_TAG:-}" > "${artifact_dir}/fallback-image.txt"
+
+    for whl in "${wheel_dir}"/*.whl; do
+        [[ -f "${whl}" ]] || continue
+        whl_name=$(basename "${whl}")
+        cp "${whl}" "${artifact_dir}/${whl_name}"
+        echo "Copied ${whl_name}: $(du -sh "${artifact_dir}/${whl_name}" | cut -f1)"
+    done

    if command -v buildkite-agent >/dev/null 2>&1; then
-        buildkite-agent artifact upload "${artifact_dir}/*" || return 1
+        buildkite-agent artifact upload "${artifact_dir}/*"
        echo "ROCm vLLM install artifacts uploaded to ${artifact_dir}/"
-    elif [[ "${BUILDKITE:-false}" == "true" ]]; then
-        echo "buildkite-agent not found; cannot upload required ROCm artifacts" >&2
-        return 1
    else
        echo "Not in Buildkite, skipping artifact upload"
    fi
@@ -2155,7 +1920,6 @@ main() {
    compute_dependency_cache_keys
    write_ci_base_label_override
    compute_rocm_csrc_content_hash_if_needed
-    compute_rocm_rust_content_hash_if_needed
    write_rocm_cache_override
    resolve_ci_base_dependency_targets
    print_bake_config
@@ -2163,11 +1927,6 @@ main() {
        echo "BAKE_PRINT_ONLY=1 set; skipping build"
        return 0
    fi
-    if should_upload_wheel_artifacts; then
-        # wheel-export is an output directory, not a BuildKit cache. Starting
-        # clean prevents a failed/retried export from packaging a stale wheel.
-        rm -rf ./wheel-export
-    fi
    seed_dependency_caches_if_needed
    run_bake
    upload_wheel_artifacts_if_present
@@ -45,10 +45,8 @@ $PYTHON .buildkite/scripts/generate-nightly-index.py --version "$SUBPATH" --curr
 echo "Uploading indices to $S3_COMMIT_PREFIX"
 aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "$S3_COMMIT_PREFIX"

-# copy to /nightly/ only when enabled for a main branch build that is not a PR
-if [[ "${UPDATE_NIGHTLY_INDEX:-1}" == "1" && \
-      "$BUILDKITE_BRANCH" == "main" && \
-      "$BUILDKITE_PULL_REQUEST" == "false" ]]; then
+# copy to /nightly/ only if it is on the main branch and not a PR
+if [[ "$BUILDKITE_BRANCH" == "main" && "$BUILDKITE_PULL_REQUEST" == "false" ]]; then
    echo "Uploading indices to overwrite /nightly/"
    aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/nightly/"
 fi
@@ -69,7 +67,7 @@ pure_version="${version%%+*}"
 echo "Pure version (without variant): $pure_version"

 # re-generate and copy to /<pure_version>/ only if it does not have "dev" in the version
-if [[ "${UPDATE_VERSION_INDEX:-1}" == "1" && "$version" != *"dev"* ]]; then
+if [[ "$version" != *"dev"* ]]; then
    echo "Re-generating indices for /$pure_version/"
    rm -rf "${INDICES_OUTPUT_DIR:?}"
    mkdir -p "$INDICES_OUTPUT_DIR"
@@ -1,7 +1,7 @@
 #!/bin/bash

-# This script runs ROCm tests either directly in a native CI pod or inside the
-# corresponding Docker container. Multi-node tests continue to use Docker.
+# This script runs tests inside the corresponding ROCm docker container.
+# It handles both single-node and multi-node test configurations.
 #
 # Multi-node detection: Instead of matching on fragile group names, we detect
 # multi-node jobs structurally by looking for the bracket command syntax
@@ -34,27 +34,10 @@ set -o pipefail
 : "${CLICOLOR_FORCE:=1}"
 : "${PY_COLORS:=1}"
 : "${ROCM_DOCKER_TTY:=1}"
-: "${PYTHONFAULTHANDLER:=1}"
-: "${PYTEST_TIMEOUT:=2400}"
 if [[ " ${PYTEST_ADDOPTS:-} " != *" --color"* ]]; then
  PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }--color=yes"
 fi
-if [[ " ${PYTEST_ADDOPTS:-} " != *" --durations="* ]]; then
-  PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }--durations=25"
-fi
-if [[ " ${PYTEST_ADDOPTS:-} " != *" --durations-min="* ]]; then
-  PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }--durations-min=1.0"
-fi
-# Dump stacks after 25 minutes, then stop an individual test after 40 minutes.
-if [[ " ${PYTEST_ADDOPTS:-} " != *" faulthandler_timeout="* ]]; then
-  PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }-o faulthandler_timeout=1500"
-fi
-if [[ " ${PYTEST_ADDOPTS:-} " != *" --timeout-method="* &&
-  " ${PYTEST_ADDOPTS:-} " != *" --timeout-method "* ]]; then
-  PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }--timeout-method=thread"
-fi
-export BUILDKIT_PROGRESS TERM FORCE_COLOR CLICOLOR_FORCE PY_COLORS PYTEST_ADDOPTS PYTEST_TIMEOUT ROCM_DOCKER_TTY
-export PYTHONFAULTHANDLER
+export BUILDKIT_PROGRESS TERM FORCE_COLOR CLICOLOR_FORCE PY_COLORS PYTEST_ADDOPTS ROCM_DOCKER_TTY

 # Export Python path for commands that run directly on the host. Containerized
 # tests set this to /vllm-workspace below so spawned Python processes do not
@@ -70,28 +53,6 @@ report_docker_usage() {
  docker system df || true
 }

-clear_ci_orchestration_env() {
-  unset -v \
-    VLLM_TEST_GROUP_NAME \
-    VLLM_CI_REQUIRE_PERSISTENT_HF_CACHE \
-    VLLM_CI_ARTIFACT_STEP \
-    VLLM_TEST_CACHE \
-    VLLM_CI_EXECUTION_MODE \
-    VLLM_CI_WORKSPACE \
-    VLLM_CI_REQUIRE_WORKSPACE_MOUNT \
-    VLLM_TEST_COMMANDS \
-    VLLM_CI_BRANCH \
-    VLLM_CI_BASE_IMAGE \
-    VLLM_CI_FALLBACK_IMAGE \
-    VLLM_CI_DOCKER_DISABLED \
-    VLLM_CI_ARTIFACT_GLOB \
-    VLLM_CI_ARTIFACT_CHECKSUM_GLOB \
-    VLLM_CI_EXPECTED_GPU_COUNT \
-    VLLM_CI_USE_ARTIFACTS \
-    VLLM_CI_RESULTS_ROOT \
-    VLLM_ALLOW_DEPRECATED_BEAM_SEARCH
-}
-
 cleanup_network() {
  local max_nodes=${NUM_NODES:-2}
  for node in $(seq 0 $((max_nodes - 1))); do
@@ -184,11 +145,7 @@ prepare_artifact_image() {
  fi

  cp "${wheel_dir}"/*.whl "${context_dir}/wheels/" || return 1
-  tar -C "${wheel_dir}" \
-    --exclude='*.whl' \
-    --exclude='.vllm-ci-artifact' \
-    --exclude='./.vllm-ci-artifact' \
-    -cf - . \
+  tar -C "${wheel_dir}" --exclude='*.whl' -cf - . \
    | tar -C "${workspace_dir}" -xf - || return 1
  cat > "${context_dir}/Dockerfile" <<'EOF'
 ARG BASE_IMAGE
@@ -211,259 +168,6 @@ EOF
  return 0
 }

-is_native_runtime() {
-  [[ "${AMD_CI_RUNTIME:-}" == "native" || "${NATIVE_CI:-}" == "true" ]]
-}
-
-validate_native_workspace() {
-  local workspace_dir="${VLLM_CI_WORKSPACE:-/vllm-workspace}"
-  local workspace_real=""
-  local checkout_real=""
-  local workspace_mount=""
-
-  mkdir -p "${workspace_dir}" || return 1
-  workspace_real=$(readlink -m "${workspace_dir}") || return 1
-  if [[ -n "${BUILDKITE_BUILD_CHECKOUT_PATH:-}" ]]; then
-    checkout_real=$(readlink -m "${BUILDKITE_BUILD_CHECKOUT_PATH}") || return 1
-    if [[ "${checkout_real}" == "${workspace_real}" \
-      || "${checkout_real}" == "${workspace_real}/"* \
-      || "${workspace_real}" == "${checkout_real}/"* ]]; then
-      echo "Refusing to replace ${workspace_real}; it overlaps the Buildkite checkout ${checkout_real}" >&2
-      return 1
-    fi
-  fi
-  if [[ "${VLLM_CI_REQUIRE_WORKSPACE_MOUNT:-1}" == "1" ]]; then
-    if ! command -v findmnt >/dev/null 2>&1; then
-      echo "findmnt is required to verify the native workspace mount" >&2
-      return 1
-    fi
-    workspace_mount=$(findmnt -n -T "${workspace_real}" -o TARGET 2>/dev/null || true)
-    if [[ "$(readlink -m "${workspace_mount:-/}")" != "${workspace_real}" ]]; then
-      echo "Native CI requires a dedicated volume mounted at ${workspace_real}" >&2
-      return 1
-    fi
-  fi
-}
-
-prepare_native_workspace() {
-  if [[ "${VLLM_CI_USE_ARTIFACTS:-0}" != "1" ]]; then
-    echo "Native CI requires VLLM_CI_USE_ARTIFACTS=1"
-    return 1
-  fi
-  if ! command -v buildkite-agent >/dev/null 2>&1; then
-    echo "buildkite-agent not found; cannot download ROCm wheel artifact"
-    return 1
-  fi
-  validate_native_workspace || return 1
-
-  local artifact_glob="${VLLM_CI_ARTIFACT_GLOB:-artifacts/vllm-rocm-install/vllm-rocm-install.tar.gz}"
-  local artifact_checksum_glob="${VLLM_CI_ARTIFACT_CHECKSUM_GLOB:-${artifact_glob}.sha256}"
-  local artifact_step="${VLLM_CI_ARTIFACT_STEP:-image-build-amd}"
-  local archive=""
-  local checksum=""
-  local download_dir=""
-  local metadata_dir=""
-  local recorded_base=""
-  local recorded_commit=""
-  local recorded_wheel=""
-  local workspace_dir="${VLLM_CI_WORKSPACE:-/vllm-workspace}"
-  local wheel_dir=""
-  local attempt=0
-  local attempt_dir=""
-  local -a archives=()
-  local -a checksums=()
-  local -a wheels=()
-
-  artifact_work_dir=$(mktemp -d -t vllm-rocm-artifact.XXXXXX) || return 1
-  wheel_dir="${artifact_work_dir}/wheels"
-  mkdir -p "${wheel_dir}" || return 1
-
-  echo "--- Downloading ROCm wheel artifact from ${artifact_step} (native in-pod)"
-  for attempt in 1 2 3; do
-    attempt_dir="${artifact_work_dir}/download-${attempt}"
-    rm -rf "${attempt_dir}" || return 1
-    mkdir -p "${attempt_dir}" || return 1
-    if buildkite-agent artifact download \
-      "${artifact_glob}" "${attempt_dir}" --step "${artifact_step}" \
-      && buildkite-agent artifact download \
-        "${artifact_checksum_glob}" "${attempt_dir}" --step "${artifact_step}"; then
-      download_dir="${attempt_dir}"
-      break
-    fi
-    echo "Artifact download attempt ${attempt}/3 failed"
-    if [[ "${attempt}" -lt 3 ]]; then
-      sleep $((attempt * 2))
-    fi
-  done
-  if [[ -z "${download_dir}" ]]; then
-    echo "Failed to download ${artifact_glob} and ${artifact_checksum_glob} from ${artifact_step}"
-    return 1
-  fi
-
-  mapfile -t archives < <(
-    find "${download_dir}" -name "vllm-rocm-install.tar.gz" -type f -print
-  )
-  mapfile -t checksums < <(
-    find "${download_dir}" -name "vllm-rocm-install.tar.gz.sha256" -type f -print
-  )
-  if [[ ${#archives[@]} -ne 1 || ${#checksums[@]} -ne 1 ]]; then
-    echo "Expected exactly one ROCm archive and checksum; found ${#archives[@]} archive(s) and ${#checksums[@]} checksum(s)" >&2
-    return 1
-  fi
-  archive="${archives[0]}"
-  checksum="${checksums[0]}"
-  if [[ "$(dirname "${archive}")" != "$(dirname "${checksum}")" ]]; then
-    echo "ROCm archive and checksum were downloaded to different directories" >&2
-    return 1
-  fi
-  (
-    cd "$(dirname "${archive}")"
-    sha256sum -c "$(basename "${checksum}")"
-  ) || return 1
-
-  tar --no-same-owner -xzf "${archive}" -C "${wheel_dir}" || return 1
-  mapfile -t wheels < <(
-    find "${wheel_dir}" -maxdepth 1 -type f -name '*.whl' -print
-  )
-  if [[ ${#wheels[@]} -ne 1 ]]; then
-    echo "ROCm artifact must contain exactly one top-level wheel; found ${#wheels[@]}" >&2
-    return 1
-  fi
-  metadata_dir="${wheel_dir}/.vllm-ci-artifact"
-  for metadata_file in commit.txt native-base-image.txt wheel-filename.txt; do
-    if [[ ! -s "${metadata_dir}/${metadata_file}" ]]; then
-      echo "ROCm artifact metadata is missing ${metadata_file}" >&2
-      return 1
-    fi
-  done
-  for metadata_file in ci-base-image.txt fallback-image.txt; do
-    if [[ ! -f "${metadata_dir}/${metadata_file}" ]]; then
-      echo "ROCm artifact metadata is missing ${metadata_file}" >&2
-      return 1
-    fi
-  done
-
-  recorded_commit=$(tr -d '\r\n' < "${metadata_dir}/commit.txt")
-  recorded_base=$(tr -d '\r\n' < "${metadata_dir}/native-base-image.txt")
-  recorded_wheel=$(tr -d '\r\n' < "${metadata_dir}/wheel-filename.txt")
-  if [[ -z "${BUILDKITE_COMMIT:-}" || "${recorded_commit}" != "${BUILDKITE_COMMIT}" ]]; then
-    echo "ROCm artifact commit ${recorded_commit} does not match ${BUILDKITE_COMMIT:-unset}" >&2
-    return 1
-  fi
-  if [[ -z "${VLLM_CI_BASE_IMAGE:-}" || "${recorded_base}" != "${VLLM_CI_BASE_IMAGE}" ]]; then
-    echo "ROCm artifact base ${recorded_base} does not match ${VLLM_CI_BASE_IMAGE:-unset}" >&2
-    return 1
-  fi
-  if [[ "${recorded_wheel}" != "$(basename "${wheels[0]}")" ]]; then
-    echo "ROCm artifact wheel manifest ${recorded_wheel} does not match $(basename "${wheels[0]}")" >&2
-    return 1
-  fi
-  for required_dir in tests .buildkite requirements; do
-    if [[ ! -d "${wheel_dir}/${required_dir}" ]]; then
-      echo "ROCm wheel artifact did not contain ${required_dir}/" >&2
-      return 1
-    fi
-  done
-
-  echo "--- Installing ROCm wheel into pod environment"
-  python3 -m pip install --no-deps --force-reinstall "${wheels[0]}" || return 1
-
-  echo "--- Preparing ${workspace_dir} from artifact"
-  find "${workspace_dir}" -mindepth 1 -maxdepth 1 -exec rm -rf -- {} + || return 1
-  tar -C "${wheel_dir}" \
-    --exclude='*.whl' \
-    --exclude='.vllm-ci-artifact' \
-    --exclude='./.vllm-ci-artifact' \
-    -cf - . | tar --no-same-owner -C "${workspace_dir}" -xf - || return 1
-  if [[ ! -d "${workspace_dir}/tests" ]]; then
-    echo "Failed to stage the native test workspace" >&2
-    return 1
-  fi
-
-  return 0
-}
-
-initialize_native_environment() {
-  local job_id="${BUILDKITE_JOB_ID:-${BUILDKITE_PARALLEL_JOB:-local}}"
-  local job_id_suffix=""
-  local native_root=""
-  local hf_mount=""
-
-  if [[ "$(id -u)" -ne 0 ]]; then
-    echo "Native ROCm CI currently requires the ci_base container to run as root" >&2
-    return 1
-  fi
-
-  job_id="${job_id//[^A-Za-z0-9_.-]/_}"
-  job_id_suffix="${job_id##*-}"
-  job_id_suffix="${job_id_suffix:0:12}"
-  native_root="/tmp/vllm-native-${job_id}"
-  TMPDIR="/tmp/vllm-${job_id_suffix}/tmp"
-  VLLM_RPC_BASE_PATH="/tmp"
-  TORCHINDUCTOR_CACHE_DIR="${native_root}/cache/torchinductor"
-  TRITON_CACHE_DIR="${native_root}/cache/triton"
-  VLLM_CACHE_ROOT="${native_root}/cache/vllm"
-  XDG_CACHE_HOME="${native_root}/cache/xdg"
-  : "${HF_HOME:=/home/buildkite-agent/huggingface}"
-  : "${HF_HUB_DOWNLOAD_TIMEOUT:=300}"
-  : "${HF_HUB_ETAG_TIMEOUT:=60}"
-  export TMPDIR VLLM_RPC_BASE_PATH
-  export TORCHINDUCTOR_CACHE_DIR TRITON_CACHE_DIR VLLM_CACHE_ROOT XDG_CACHE_HOME
-  export HF_HOME HF_HUB_DOWNLOAD_TIMEOUT HF_HUB_ETAG_TIMEOUT
-  export PYTORCH_ROCM_ARCH=""
-
-  mkdir -p "${TMPDIR}" \
-    "${TORCHINDUCTOR_CACHE_DIR}" \
-    "${TRITON_CACHE_DIR}" \
-    "${VLLM_CACHE_ROOT}" \
-    "${XDG_CACHE_HOME}" \
-    "${HF_HOME}" || return 1
-
-  echo "Native compile caches: VLLM_CACHE_ROOT=${VLLM_CACHE_ROOT} TORCHINDUCTOR_CACHE_DIR=${TORCHINDUCTOR_CACHE_DIR}"
-
-  if [[ "${VLLM_CI_REQUIRE_PERSISTENT_HF_CACHE:-0}" == "1" ]]; then
-    if ! command -v findmnt >/dev/null 2>&1; then
-      echo "findmnt is required to verify the native Hugging Face cache mount" >&2
-      return 1
-    fi
-    hf_mount=$(findmnt -n -T "${HF_HOME}" -o TARGET 2>/dev/null || true)
-    if [[ -z "${hf_mount}" || "${hf_mount}" == "/" ]]; then
-      echo "Native CI requires a persistent volume mounted at or above ${HF_HOME}" >&2
-      return 1
-    fi
-  fi
-}
-
-run_native_preflight() {
-  local expected_gpus="${VLLM_CI_EXPECTED_GPU_COUNT:-1}"
-
-  if [[ ! "${expected_gpus}" =~ ^[0-9]+$ ]]; then
-    echo "Invalid VLLM_CI_EXPECTED_GPU_COUNT=${expected_gpus}" >&2
-    return 1
-  fi
-
-  python3 -c "import encodings, importlib.metadata as im, importlib.util as iu; [im.version(d) for d in ('transformers', 'torch', 'ray', 'sympy', 'markupsafe', 'vllm')]; missing=[m for m in ('torch.utils.model_zoo', 'transformers.models.nomic_bert', 'ray.dag', 'sympy.physics', 'markupsafe._speedups') if iu.find_spec(m) is None]; assert not missing, missing" || return 1
-
-  if [[ "${expected_gpus}" == "0" ]]; then
-    echo "Native CPU-only AMD job: skipping ROCm device validation"
-    return 0
-  fi
-
-  echo "--- ROCm info"
-  rocminfo || return 1
-  VLLM_CI_EXPECTED_GPU_COUNT="${expected_gpus}" python3 - <<'PY'
-import os
-
-import torch
-
-expected = int(os.environ["VLLM_CI_EXPECTED_GPU_COUNT"])
-assert torch.version.hip, "PyTorch is not a ROCm build"
-assert torch.cuda.is_available(), "ROCm GPU is not available to PyTorch"
-actual = torch.cuda.device_count()
-assert actual == expected, f"Expected {expected} ROCm GPU(s), found {actual}"
-PY
-}
-
 is_multi_node() {
  local cmds="$1"
  # Primary signal: NUM_NODES environment variable set by the pipeline
@@ -646,58 +350,7 @@ re_quote_pytest_markers() {
 # Main
 ###############################################################################

-if is_native_runtime; then
-  echo "--- Native in-pod ROCm CI (AMD_CI_RUNTIME=${AMD_CI_RUNTIME:-unset}, NATIVE_CI=${NATIVE_CI:-unset})"
-  artifact_work_dir=""
-
-  cleanup_native_workspace() {
-    if [[ -n "${artifact_work_dir}" ]]; then
-      rm -rf "${artifact_work_dir}"
-    fi
-  }
-  trap cleanup_native_workspace EXIT
-
-  if [[ -n "${VLLM_TEST_COMMANDS:-}" ]]; then
-    commands="${VLLM_TEST_COMMANDS}"
-    commands_source="env"
-  else
-    commands="$*"
-    commands_source="argv"
-    if [[ -z "$commands" ]]; then
-      echo "Error: No test commands provided for native CI." >&2
-      exit 1
-    fi
-  fi
-
-  if [[ "$commands_source" == "argv" ]]; then
-    commands=$(re_quote_pytest_markers "$commands")
-  fi
-
-  if is_multi_node "$commands"; then
-    echo "Native CI does not support multi-node jobs yet."
-    exit 1
-  fi
-
-  if ! initialize_native_environment; then
-    echo "Failed to initialize the native test environment"
-    exit 1
-  fi
-  if ! prepare_native_workspace; then
-    echo "Failed to prepare native test workspace"
-    exit 1
-  fi
-
-  export PYTHONPATH="${VLLM_CI_WORKSPACE:-/vllm-workspace}"
-
-  echo "Native test commands: $commands"
-  run_native_preflight || exit 1
-  # Keep AMD CI orchestration variables out of vLLM's runtime environment.
-  clear_ci_orchestration_env
-  /bin/bash -o pipefail -c "${commands}"
-  handle_pytest_exit "$?"
-fi
-
-# --- GPU initialization for legacy Docker execution ---
+# --- GPU initialization ---
 echo "--- ROCm info"
 rocminfo

@@ -799,27 +452,25 @@ fi

 echo "Final commands: $commands"

-standalone_merge_base_env=()
-if [[ "$commands" == *python_only_compile.sh* ]]; then
-  # The ROCm test image often ships /vllm-workspace without .git. Resolve the
-  # wheels.vllm.ai commit from the agent checkout for this test only.
-  vllm_standalone_merge_base=""
-  checkout="${BUILDKITE_BUILD_CHECKOUT_PATH:-}"
-  if [[ -z "${checkout}" || ! -d "${checkout}" ]]; then
-    checkout="."
-  fi
-  # Pass safe.directory per-command because Buildkite uses mixed user IDs.
-  if git -c "safe.directory=${checkout}" -C "${checkout}" rev-parse --is-inside-work-tree >/dev/null 2>&1; then
-    vllm_standalone_merge_base="$(
-      git -c "safe.directory=${checkout}" -C "${checkout}" merge-base HEAD origin/main 2>/dev/null || true
-    )"
-  fi
-  if [[ -z "${vllm_standalone_merge_base}" ]]; then
-    vllm_standalone_merge_base="${BUILDKITE_COMMIT:-}"
-  fi
-  echo "INFO: passing CI_STANDALONE_MERGE_BASE into container: ${vllm_standalone_merge_base}"
-  standalone_merge_base_env=(-e "CI_STANDALONE_MERGE_BASE=${vllm_standalone_merge_base}")
+# The ROCm test image often ships /vllm-workspace without .git (artifact tarball unpack).
+# tests/standalone_tests/python_only_compile.sh uses merge-base(HEAD, origin/main) for
+# wheels.vllm.ai; compute on the agent (full git checkout) and pass into the container.
+vllm_standalone_merge_base=""
+checkout="${BUILDKITE_BUILD_CHECKOUT_PATH:-}"
+if [[ -z "${checkout}" || ! -d "${checkout}" ]]; then
+  checkout="."
 fi
+# Pass safe.directory per-command (-c) because buildkite runs will always fail
+# the next check on git 2.35.2+ due to mixed uses of root and buildkite-agent/uids.
+if git -c "safe.directory=${checkout}" -C "${checkout}" rev-parse --is-inside-work-tree >/dev/null 2>&1; then
+  vllm_standalone_merge_base="$(
+    git -c "safe.directory=${checkout}" -C "${checkout}" merge-base HEAD origin/main 2>/dev/null || true
+  )"
+fi
+if [[ -z "${vllm_standalone_merge_base}" ]]; then
+  vllm_standalone_merge_base="${BUILDKITE_COMMIT:-}"
+fi
+echo "INFO: passing VLLM_STANDALONE_MERGE_BASE into container: ${vllm_standalone_merge_base}"

 MYPYTHONPATH="/vllm-workspace"

@@ -850,7 +501,6 @@ else
 fi

 # --- Route: multi-node vs single-node ---
-clear_ci_orchestration_env
 if is_multi_node "$commands"; then
  echo "--- Multi-node job detected"
  export DCKR_VER=$(docker --version | sed 's/Docker version \(.*\), build .*/\1/')
@@ -939,9 +589,7 @@ else
    -e FORCE_COLOR \
    -e CLICOLOR_FORCE \
    -e PY_COLORS \
-    -e PYTHONFAULTHANDLER \
    -e PYTEST_ADDOPTS \
-    -e PYTEST_TIMEOUT \
    -v "${HF_CACHE}:${HF_MOUNT}" \
    -e "HF_HOME=${HF_MOUNT}" \
    -e "PYTHONPATH=${MYPYTHONPATH}" \
@@ -951,7 +599,7 @@ else
    -e "VLLM_CACHE_ROOT=${CONTAINER_CACHE_ROOT}/vllm" \
    -e "XDG_CACHE_HOME=${CONTAINER_CACHE_ROOT}/xdg" \
    -e "PYTORCH_ROCM_ARCH=" \
-    "${standalone_merge_base_env[@]}" \
+    -e "VLLM_STANDALONE_MERGE_BASE=${vllm_standalone_merge_base}" \
    --name "${container_name}" \
    "${image_name}" \
    /bin/bash -c "${CONTAINER_PREFLIGHT} && ${commands}"
@@ -1,11 +1,10 @@
 #!/bin/bash
 set -euox pipefail

-export VLLM_CPU_KVCACHE_SPACE=1
+export VLLM_CPU_KVCACHE_SPACE=1 
 export VLLM_CPU_CI_ENV=1
-# Skip torch.compile via vLLM's --enforce-eager flag (passed below) instead of
-# TORCH_COMPILE_DISABLE=1, which torch 2.12 no longer treats as a silent no-op
-# when callers specify fullgraph=True.
+# Reduce sub-processes for acceleration
+export TORCH_COMPILE_DISABLE=1 
 export VLLM_ENABLE_V1_MULTIPROCESSING=0

 SDE_ARCHIVE="sde-external-10.7.0-2026-02-18-lin.tar.xz"
@@ -50,15 +49,15 @@ wait_for_pid_and_check_log() {
 }

 # Test Sky Lake (AVX512F)
-./sde/sde64 -skl -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 --enforce-eager > test_0.log 2>&1 &
+./sde/sde64 -skl -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 > test_0.log 2>&1 &
 PID_TEST_0=$!

 # Test Cascade Lake (AVX512F + VNNI)
-./sde/sde64 -clx -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 --enforce-eager > test_1.log 2>&1 &
+./sde/sde64 -clx -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 > test_1.log 2>&1 &
 PID_TEST_1=$!

 # Test Cooper Lake (AVX512F + VNNI + BF16)
-./sde/sde64 -cpx -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 --enforce-eager > test_2.log 2>&1 &
+./sde/sde64 -cpx -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 > test_2.log 2>&1 &
 PID_TEST_2=$!

 wait_for_pid_and_check_log $PID_TEST_0 test_0.log
@@ -40,9 +40,7 @@ function cpu_tests() {
    pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py
    pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
    pytest -x -v -s tests/kernels/moe/test_cpu_int4_moe.py
-    pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py
-    pytest -x -v -s tests/kernels/mamba/test_causal_conv1d.py
-    pytest -x -v -s tests/kernels/mamba/test_mamba_ssm.py"
+    pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py"

  # skip tests requiring model downloads if HF_TOKEN is not set
  # due to rate-limits
@@ -99,4 +97,3 @@ function cpu_tests() {
 # All of CPU tests are expected to be finished less than 40 mins.
 export -f cpu_tests
 timeout 2h bash -c cpu_tests
-
@@ -35,7 +35,6 @@ case "${test_suite}" in
    pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py --ignore=v1/worker/test_worker_memory_snapshot.py
    pytest -v -s v1/structured_output
    pytest -v -s v1/test_serial_utils.py
-    pytest -v -s v1/e2e/general/test_correctness_sliding_window.py --deselect="tests/v1/e2e/general/test_correctness_sliding_window.py::test_sliding_window_retrieval[True-1-5-google/gemma-3-1b-it]"
    pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py --ignore=v1/spec_decode/test_speculators_correctness.py
    pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py --ignore=v1/kv_connector/unit/test_offloading_connector.py
    ;;
@@ -130,22 +130,6 @@ docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm
 docker push vllm/vllm-openai-rocm:latest-base
 docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base

-# ---- XPU ----
-
-docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-xpu
-
-docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-xpu vllm/vllm-openai-xpu:latest-x86_64
-docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-xpu vllm/vllm-openai-xpu:v${RELEASE_VERSION}-x86_64
-docker push vllm/vllm-openai-xpu:latest-x86_64
-docker push vllm/vllm-openai-xpu:v${RELEASE_VERSION}-x86_64
-
-docker manifest rm vllm/vllm-openai-xpu:latest || true
-docker manifest rm vllm/vllm-openai-xpu:v${RELEASE_VERSION} || true
-docker manifest create vllm/vllm-openai-xpu:latest vllm/vllm-openai-xpu:latest-x86_64 --amend
-docker manifest create vllm/vllm-openai-xpu:v${RELEASE_VERSION} vllm/vllm-openai-xpu:v${RELEASE_VERSION}-x86_64 --amend
-docker manifest push vllm/vllm-openai-xpu:latest
-docker manifest push vllm/vllm-openai-xpu:v${RELEASE_VERSION}
-
 # ---- CPU ----
 # CPU images are behind separate block steps and may not have been built.
 # All-or-nothing: inspect both arches first, then either publish everything
@@ -13,18 +13,6 @@ metadata_get() {
    fi
 }

-use_ci_base_if_present() {
-    local ci_base_image=""
-
-    ci_base_image="$(metadata_get rocm-ci-base-image)"
-    if [[ -z "${ci_base_image}" ]]; then
-        return 1
-    fi
-
-    export CI_BASE_IMAGE="${ci_base_image}"
-    echo "Using ROCm ci_base image selected by the preceding build step: ${CI_BASE_IMAGE}"
-}
-
 use_refreshed_base_if_present() {
    local base_refreshed=""

@@ -34,12 +22,15 @@ use_refreshed_base_if_present() {
    fi

    export BASE_IMAGE
+    export CI_BASE_IMAGE
    export IMAGE_TAG_LATEST

    BASE_IMAGE="$(metadata_get rocm-base-image)"
+    CI_BASE_IMAGE="$(metadata_get rocm-ci-base-image)"
    IMAGE_TAG_LATEST="$(metadata_get rocm-ci-image-descriptive)"

    echo "Using refreshed ROCm base image for test image: ${BASE_IMAGE}"
+    echo "Using refreshed ROCm ci_base image for test image: ${CI_BASE_IMAGE}"
    if [[ -n "${IMAGE_TAG_LATEST}" ]]; then
        echo "Also tagging full ROCm CI image as: ${IMAGE_TAG_LATEST}"
    fi
@@ -50,8 +41,6 @@ use_refreshed_base_if_present() {
 main() {
    local base_refreshed=0

-    use_ci_base_if_present || true
-
    if use_refreshed_base_if_present; then
        base_refreshed=1
    fi
@@ -21,20 +21,16 @@ export CARGO_HOME="${CARGO_HOME:-$HOME/.cargo}"
 export RUSTUP_HOME="${RUSTUP_HOME:-$HOME/.rustup}"
 export PATH="$CARGO_HOME/bin:$PATH"

-PROTOC_VERSION="${PROTOC_VERSION:-31.1}"
-CARGO_BINSTALL_VERSION="${CARGO_BINSTALL_VERSION:-1.20.1}"
-UV_VERSION="${UV_VERSION:-0.11.28}"
-PYO3_PYTHON_VERSION="${PYO3_PYTHON_VERSION:-3.12}"
-
-CARGO_SORT_VERSION_REQ="${CARGO_SORT_VERSION_REQ:-2}"
-CARGO_DENY_VERSION_REQ="${CARGO_DENY_VERSION_REQ:-0.20}"
-CARGO_NEXTEST_VERSION_REQ="${CARGO_NEXTEST_VERSION_REQ:-0.9}"
-
 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)
@@ -49,17 +45,16 @@ install_protoc() {
      ;;
  esac

-  local url="https://github.com/protocolbuffers/protobuf/releases/download/v${PROTOC_VERSION}/protoc-${PROTOC_VERSION}-linux-${arch}.zip"
+  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 ${PROTOC_VERSION}"
+  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"
-  protoc --version
 }

 rust_toolchain() {
@@ -80,48 +75,66 @@ install_rust_toolchain() {
 }

 install_cargo_binstall() {
-  log_section "Installing cargo-binstall ${CARGO_BINSTALL_VERSION}"
+  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/v${CARGO_BINSTALL_VERSION}/install-from-binstall-release.sh" \
-    | env BINSTALL_VERSION="$CARGO_BINSTALL_VERSION" bash
-  cargo-binstall -V
+    https://raw.githubusercontent.com/cargo-bins/cargo-binstall/main/install-from-binstall-release.sh \
+    | bash
 }

 install_cargo_sort() {
-  log_section "Installing cargo-sort ${CARGO_SORT_VERSION_REQ}"
-  cargo binstall --no-confirm --force "cargo-sort@${CARGO_SORT_VERSION_REQ}"
+  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_deny() {
-  log_section "Installing cargo-deny ${CARGO_DENY_VERSION_REQ}"
-  cargo binstall --no-confirm --force "cargo-deny@${CARGO_DENY_VERSION_REQ}"
+  if command -v cargo-deny >/dev/null 2>&1; then
+    return
+  fi
+
+  log_section "Installing cargo-deny"
+  install_cargo_binstall
+  cargo binstall --no-confirm cargo-deny
 }

 install_cargo_nextest() {
-  log_section "Installing cargo-nextest ${CARGO_NEXTEST_VERSION_REQ}"
-  cargo binstall \
-    --no-confirm \
-    --force \
-    --secure \
-    "cargo-nextest@${CARGO_NEXTEST_VERSION_REQ}"
+  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() {
-  log_section "Installing uv ${UV_VERSION}"
-  curl -L --proto '=https' --tlsv1.2 -sSf \
-    "https://github.com/astral-sh/uv/releases/download/${UV_VERSION}/uv-installer.sh" \
+  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
-  uv --version
 }

 setup_pyo3_python() {
-  log_section "Installing Python ${PYO3_PYTHON_VERSION} for PyO3 tests"
-  uv python install "$PYO3_PYTHON_VERSION"
+  local python_version="${PYO3_PYTHON_VERSION:-3.12}"
+
+  log_section "Installing Python ${python_version} for PyO3 tests"
+  uv python install "$python_version"
  PYO3_PYTHON="$(uv python find \
    --managed-python \
    --no-project \
    --resolve-links \
-    "$PYO3_PYTHON_VERSION")"
+    "$python_version")"
  export PYO3_PYTHON

  local python_libdir
@@ -143,7 +156,6 @@ PY
 }

 run_style_clippy() {
-  install_cargo_binstall
  install_cargo_sort
  install_cargo_deny

@@ -174,7 +186,6 @@ run_style_clippy() {
 run_tests() {
  install_uv
  setup_pyo3_python
-  install_cargo_binstall
  install_cargo_nextest

  log_section "Running cargo nextest"
@@ -23,7 +23,6 @@ NC='\033[0m' # No Color
 # Default configuration
 PIPELINE="ci"
 DRY_RUN=true
-TORCH_NIGHTLY=false

 usage() {
    cat <<EOF
@@ -35,14 +34,12 @@ Sets RUN_ALL=1 and NIGHTLY=1 environment variables.
 SAFETY: Dry-run by default. Use --execute to actually trigger a build.

 Options:
-    --execute        Actually trigger the build (default: dry-run)
-    --pipeline       Buildkite pipeline slug (default: ${PIPELINE})
-    --commit         Override commit SHA (default: current HEAD)
-    --branch         Override branch name (default: current branch)
-    --message        Custom build message (default: auto-generated)
-    --torch-nightly  Also build and run the full suite against torch nightly
-                     (sets TORCH_NIGHTLY=1)
-    --help           Show this help message
+    --execute       Actually trigger the build (default: dry-run)
+    --pipeline      Buildkite pipeline slug (default: ${PIPELINE})
+    --commit        Override commit SHA (default: current HEAD)
+    --branch        Override branch name (default: current branch)
+    --message       Custom build message (default: auto-generated)
+    --help          Show this help message

 Prerequisites:
    - bk CLI installed: brew tap buildkite/buildkite && brew install buildkite/buildkite/bk
@@ -52,7 +49,6 @@ Examples:
    $(basename "$0")                        # Dry-run, show what would happen
    $(basename "$0") --execute              # Actually trigger the build
    $(basename "$0") --pipeline ci-shadow   # Dry-run with different pipeline
-    $(basename "$0") --torch-nightly        # Dry-run a full torch-nightly run
 EOF
    exit 1
 }
@@ -100,10 +96,6 @@ while [[ $# -gt 0 ]]; do
            MESSAGE="$2"
            shift 2
            ;;
-        --torch-nightly)
-            TORCH_NIGHTLY=true
-            shift
-            ;;
        --help|-h)
            usage
            ;;
@@ -179,17 +171,11 @@ if [[ $(echo "$REMOTE_BRANCHES" | wc -l) -gt 5 ]]; then
 fi
 echo ""

-# Environment variables passed to the build.
-BUILD_ENV=("RUN_ALL=1" "NIGHTLY=1")
-if [[ "$TORCH_NIGHTLY" == true ]]; then
-    BUILD_ENV+=("TORCH_NIGHTLY=1")
-fi
-
 log_info "Pipeline: ${PIPELINE}"
 log_info "Branch: ${BRANCH}"
 log_info "Commit: ${COMMIT}"
 log_info "Message: ${MESSAGE}"
-log_info "Environment: ${BUILD_ENV[*]}"
+log_info "Environment: RUN_ALL=1, NIGHTLY=1"
 echo ""

 # Build the command
@@ -201,10 +187,9 @@ CMD=(bk build create
    --commit "${COMMIT}"
    --branch "${BRANCH}"
    --message "${MESSAGE}"
+    --env "RUN_ALL=1"
+    --env "NIGHTLY=1"
 )
-for env_var in "${BUILD_ENV[@]}"; do
-    CMD+=(--env "${env_var}")
-done

 if [[ "$DRY_RUN" == true ]]; then
    echo "=========================================="
@@ -225,14 +210,8 @@ if [[ "$DRY_RUN" == true ]]; then
    echo "    --commit '$(escape_for_shell "${COMMIT}")' \\"
    echo "    --branch '$(escape_for_shell "${BRANCH}")' \\"
    echo "    --message '$(escape_for_shell "${MESSAGE}")' \\"
-    last_idx=$(( ${#BUILD_ENV[@]} - 1 ))
-    for i in "${!BUILD_ENV[@]}"; do
-        if [[ $i -eq $last_idx ]]; then
-            echo "    --env '$(escape_for_shell "${BUILD_ENV[$i]}")'"
-        else
-            echo "    --env '$(escape_for_shell "${BUILD_ENV[$i]}")' \\"
-        fi
-    done
+    echo "    --env 'RUN_ALL=1' \\"
+    echo "    --env 'NIGHTLY=1'"
    echo ""
    echo "=========================================="
    echo -e "${YELLOW}To actually trigger this build, run:${NC}"
@@ -6,14 +6,8 @@ set -ex
 # manylinux platform tag with auditwheel.
 # Index generation is handled separately by generate-and-upload-nightly-index.sh.

-# auditwheel is Linux-only; macOS wheels already carry a valid tag, so skip the
-# manylinux retag for them.
-WHEEL_PLATFORM="${VLLM_WHEEL_PLATFORM:-linux}"
-
-if [[ "$WHEEL_PLATFORM" == "linux" ]]; then
-  # shellcheck source=lib/manylinux.sh
-  source .buildkite/scripts/lib/manylinux.sh
-fi
+# shellcheck source=lib/manylinux.sh
+source .buildkite/scripts/lib/manylinux.sh

 BUCKET="vllm-wheels"
 SUBPATH=$BUILDKITE_COMMIT
@@ -33,10 +27,8 @@ wheel="${wheel_files[0]}"

 # ========= detect manylinux tag and rename ==========

-if [[ "$WHEEL_PLATFORM" == "linux" ]]; then
-  wheel="$(apply_manylinux_tag "$wheel")"
-  echo "Renamed wheel to: $wheel"
-fi
+wheel="$(apply_manylinux_tag "$wheel")"
+echo "Renamed wheel to: $wheel"

 # Extract the version from the wheel
 version=$(unzip -p "$wheel" '**/METADATA' | grep '^Version: ' | cut -d' ' -f2)
@@ -113,8 +113,8 @@ $PYTHON .buildkite/scripts/generate-nightly-index.py \
 echo "Uploading indices to $S3_COMMIT_PREFIX"
 aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "$S3_COMMIT_PREFIX"

-# Only scheduled nightly builds should update the moving nightly index.
-if [[ "${NIGHTLY:-0}" == "1" ]]; then
+# Update rocm/nightly/ if on main branch and not a PR
+if [[ "$BUILDKITE_BRANCH" == "main" && "$BUILDKITE_PULL_REQUEST" == "false" ]] || [[ "$NIGHTLY" == "1" ]]; then
    echo "Updating rocm/nightly/ index..."
    aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/rocm/nightly/"
 fi
@@ -147,7 +147,7 @@ echo ""
 echo "Install command (by commit):"
 echo "  pip install vllm --extra-index-url https://${BUCKET}.s3.amazonaws.com/$ROCM_SUBPATH/"
 echo ""
-if [[ "${NIGHTLY:-0}" == "1" ]]; then
+if [[ "$BUILDKITE_BRANCH" == "main" ]] || [[ "$NIGHTLY" == "1" ]]; then
    echo "Install command (nightly):"
    echo "  pip install vllm --extra-index-url https://${BUCKET}.s3.amazonaws.com/rocm/nightly/"
 fi
@@ -1,21 +0,0 @@
-#!/bin/bash
-
-set -ex
-
-ORIG_TAG_NAME="$BUILDKITE_COMMIT"
-REPO="vllm/vllm-openai-xpu"
-
-echo "Pushing original XPU tag ${ORIG_TAG_NAME}-xpu to nightly tags in ${REPO}"
-
-aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
-docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:"$ORIG_TAG_NAME"-x86_64-xpu
-
-docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:"$ORIG_TAG_NAME"-x86_64-xpu ${REPO}:nightly-x86_64
-docker push ${REPO}:nightly-x86_64
-
-docker manifest rm ${REPO}:nightly || true
-docker manifest rm ${REPO}:nightly-"$BUILDKITE_COMMIT" || true
-docker manifest create ${REPO}:nightly ${REPO}:nightly-x86_64 --amend
-docker manifest create ${REPO}:nightly-"$BUILDKITE_COMMIT" ${REPO}:nightly-x86_64 --amend
-docker manifest push ${REPO}:nightly
-docker manifest push ${REPO}:nightly-"$BUILDKITE_COMMIT"
@@ -4,7 +4,7 @@ depends_on:
 steps:
 - label: V1 attention (H100-MI300)
  key: v1-attention-h100-mi300
-  timeout_in_minutes: 85
+  timeout_in_minutes: 30
  device: h100
  source_file_dependencies:
    - vllm/config/attention.py
@@ -12,13 +12,11 @@ steps:
    - vllm/v1/attention
    - tests/v1/attention
  commands:
-    - pytest -v -s v1/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
-  parallelism: 2
+    - pytest -v -s v1/attention
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 125
+      device: mi325_1
+      timeout_in_minutes: 70
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -32,7 +30,7 @@ steps:

 - label: V1 attention (B200)
  key: v1-attention-b200
-  timeout_in_minutes: 80
+  timeout_in_minutes: 30
  device: b200-k8s
  source_file_dependencies:
    - vllm/config/attention.py
@@ -40,5 +38,4 @@ steps:
    - vllm/v1/attention
    - tests/v1/attention
  commands:
-    - pytest -v -s v1/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
-  parallelism: 2
+    - pytest -v -s v1/attention
@@ -4,7 +4,7 @@ depends_on:
 steps:
 - label: Basic Correctness
  key: basic-correctness
-  timeout_in_minutes: 68
+  timeout_in_minutes: 30
  device: h200_18gb
  source_file_dependencies:
  - vllm/
@@ -18,8 +18,7 @@ steps:
  - pytest -v -s basic_correctness/test_cpu_offload.py
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 60
+      device: mi325_1
+      timeout_in_minutes: 50
      depends_on:
      - image-build-amd
@@ -4,7 +4,7 @@ depends_on:
 steps:
 - label: Benchmarks CLI Test
  key: benchmarks-cli-test
-  timeout_in_minutes: 45
+  timeout_in_minutes: 20
  device: h200_18gb
  source_file_dependencies:
  - vllm/
@@ -13,9 +13,7 @@ steps:
  - pytest -v -s benchmarks/
  mirror:
    amd:
-      dind: false
      device: mi300_1
-      timeout_in_minutes: 40
      depends_on:
      - image-build-amd

@@ -25,7 +23,7 @@ steps:
  num_gpus: 2
  optional: true
  working_dir: "/vllm-workspace/"
-  timeout_in_minutes: 20
+  timeout_in_minutes: 10
  source_file_dependencies:
  - benchmarks/attention_benchmarks/
  - vllm/v1/attention/
@@ -4,7 +4,7 @@ depends_on:
 steps:
 - label: Sequence Parallel Correctness Tests (2 GPUs)
  key: sequence-parallel-correctness-tests-2-gpus
-  timeout_in_minutes: 80
+  timeout_in_minutes: 50
  working_dir: "/vllm-workspace/"
  num_devices: 2
  source_file_dependencies:
@@ -19,7 +19,7 @@ steps:

 - label: Sequence Parallel Correctness Tests (2xH100)
  key: sequence-parallel-correctness-tests-2xh100
-  timeout_in_minutes: 75
+  timeout_in_minutes: 50
  working_dir: "/vllm-workspace/"
  device: h100
  optional: true
@@ -30,7 +30,7 @@ steps:

 - label: AsyncTP Correctness Tests (2xH100)
  key: asynctp-correctness-tests-2xh100
-  timeout_in_minutes: 30
+  timeout_in_minutes: 50
  working_dir: "/vllm-workspace/"
  device: h100
  optional: true
@@ -41,7 +41,7 @@ steps:

 - label: AsyncTP Correctness Tests (B200)
  key: asynctp-correctness-tests-b200
-  timeout_in_minutes: 30
+  timeout_in_minutes: 50
  working_dir: "/vllm-workspace/"
  device: b200-k8s
  optional: true
@@ -52,7 +52,7 @@ steps:

 - label: Distributed Compile Unit Tests (2xH100)
  key: distributed-compile-unit-tests-2xh100
-  timeout_in_minutes: 45
+  timeout_in_minutes: 20
  working_dir: "/vllm-workspace/"
  device: h100
  num_devices: 2
@@ -66,7 +66,7 @@ steps:

 - label: Fusion and Compile Unit Tests (2xB200)
  key: fusion-and-compile-unit-tests-2xb200
-  timeout_in_minutes: 30
+  timeout_in_minutes: 20
  working_dir: "/vllm-workspace/"
  device: b200-k8s
  source_file_dependencies:
@@ -96,7 +96,7 @@ steps:

 - label: Fusion E2E Quick (H100)
  key: fusion-e2e-quick-h100
-  timeout_in_minutes: 25
+  timeout_in_minutes: 15
  working_dir: "/vllm-workspace/"
  device: h100
  num_devices: 1
@@ -115,7 +115,7 @@ steps:

 - label: Fusion E2E Config Sweep (H100)
  key: fusion-e2e-config-sweep-h100
-  timeout_in_minutes: 25
+  timeout_in_minutes: 30
  working_dir: "/vllm-workspace/"
  device: h100
  num_devices: 1
@@ -149,7 +149,7 @@ steps:

 - label: Fusion E2E TP2 Quick (H100)
  key: fusion-e2e-tp2-quick-h100
-  timeout_in_minutes: 35
+  timeout_in_minutes: 20
  working_dir: "/vllm-workspace/"
  device: h100
  num_devices: 2
@@ -167,7 +167,7 @@ steps:

 - label: Fusion E2E TP2 AR-RMS Config Sweep (H100)
  key: fusion-e2e-tp2-ar-rms-config-sweep-h100
-  timeout_in_minutes: 30
+  timeout_in_minutes: 40
  working_dir: "/vllm-workspace/"
  device: h100
  num_devices: 2
@@ -207,7 +207,7 @@ steps:

 - label: Fusion E2E TP2 (B200)
  key: fusion-e2e-tp2-b200
-  timeout_in_minutes: 45
+  timeout_in_minutes: 20
  working_dir: "/vllm-workspace/"
  device: b200-k8s
  num_devices: 2
@@ -4,7 +4,7 @@ depends_on:
 steps:
 - label: Platform Tests
  key: platform-tests
-  timeout_in_minutes: 20
+  timeout_in_minutes: 15
  device: h200_18gb
  source_file_dependencies:
  - vllm/envs.py
@@ -18,18 +18,14 @@ steps:
    - pytest -v -s cuda/test_platform_no_cuda_init.py

 - label: Cudagraph
-  device: h200_35gb
  key: cudagraph
-  timeout_in_minutes: 30
+  timeout_in_minutes: 20
  source_file_dependencies:
  - tests/v1/cudagraph
  - vllm/v1/cudagraph_dispatcher.py
  - vllm/config/compilation.py
  - vllm/compilation
-  - vllm/v1/worker/encoder_cudagraph.py
-  - vllm/v1/worker/encoder_cudagraph_defs.py
  commands:
    - pytest -v -s v1/cudagraph/test_cudagraph_dispatch.py
    - pytest -v -s v1/cudagraph/test_cudagraph_mode.py
-    - pytest -v -s v1/cudagraph/test_breakable_cudagraph.py
-    - pytest -v -s v1/cudagraph/test_encoder_cudagraph.py
+    - pytest -v -s v1/cudagraph/test_breakable_cudagraph.py
@@ -4,7 +4,7 @@ depends_on:
 steps:
 - label: Distributed NixlConnector PD accuracy (4 GPUs)
  key: distributed-nixlconnector-pd-accuracy-4-gpus
-  timeout_in_minutes: 55
+  timeout_in_minutes: 30
  working_dir: "/vllm-workspace/tests"
  num_devices: 4
  source_file_dependencies:
@@ -15,9 +15,8 @@ steps:
    - bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
  mirror:
    amd:
-      dind: false
      device: mi300_4
-      timeout_in_minutes: 60
+      timeout_in_minutes: 110
      depends_on:
        - image-build-amd
      source_file_dependencies:
@@ -30,7 +29,7 @@ steps:

 - label: Distributed FlashInfer NixlConnector PD accuracy (4 GPUs)
  key: distributed-flashinfer-nixlconnector-pd-accuracy-4-gpus
-  timeout_in_minutes: 55
+  timeout_in_minutes: 30
  working_dir: "/vllm-workspace/tests"
  num_devices: 4
  source_file_dependencies:
@@ -66,9 +65,8 @@ steps:
    - DP_EP=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
  mirror:
    amd:
-      dind: false
      device: mi300_4
-      timeout_in_minutes: 40
+      timeout_in_minutes: 50
      depends_on:
        - image-build-amd
      source_file_dependencies:
@@ -81,7 +79,7 @@ steps:

 - label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs)
  key: crosslayer-kv-layout-distributed-nixlconnector-pd-accuracy-tests-4-gpus
-  timeout_in_minutes: 55
+  timeout_in_minutes: 30
  working_dir: "/vllm-workspace/tests"
  num_devices: 4
  source_file_dependencies:
@@ -92,9 +90,8 @@ steps:
    - CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
  mirror:
    amd:
-      dind: false
      device: mi300_4
-      timeout_in_minutes: 60
+      timeout_in_minutes: 110
      depends_on:
        - image-build-amd
      source_file_dependencies:
@@ -107,7 +104,7 @@ steps:

 - label: Hybrid SSM NixlConnector PD accuracy tests (4 GPUs)
  key: hybrid-ssm-nixlconnector-pd-accuracy-tests-4-gpus
-  timeout_in_minutes: 60
+  timeout_in_minutes: 25
  working_dir: "/vllm-workspace/tests"
  num_devices: 4
  source_file_dependencies:
@@ -118,9 +115,8 @@ steps:
    - HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
  mirror:
    amd:
-      dind: false
      device: mi300_4
-      timeout_in_minutes: 55
+      timeout_in_minutes: 60
      depends_on:
        - image-build-amd
      source_file_dependencies:
@@ -147,7 +143,7 @@ steps:

 - label: MultiConnector (Nixl+Offloading) PD accuracy (2 GPUs)
  key: multiconnector-nixl-offloading-pd-accuracy-2-gpus
-  timeout_in_minutes: 40
+  timeout_in_minutes: 30
  working_dir: "/vllm-workspace/tests"
  num_devices: 2
  source_file_dependencies:
@@ -162,7 +158,7 @@ steps:

 - label: NixlConnector PD + Spec Decode acceptance (2 GPUs)
  key: nixlconnector-pd-spec-decode-acceptance-2-gpus
-  timeout_in_minutes: 45
+  timeout_in_minutes: 30
  device: a100
  working_dir: "/vllm-workspace/tests"
  num_devices: 2
@@ -175,9 +171,8 @@ steps:
    - bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh
  mirror:
    amd:
-      dind: false
      device: mi300_2
-      timeout_in_minutes: 45
+      timeout_in_minutes: 60
      depends_on:
        - image-build-amd
      source_file_dependencies:
@@ -187,11 +182,11 @@ steps:
        - vllm/platforms/rocm.py
      commands:
        - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
-        - KV_CACHE_MEMORY_BYTES=8G ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh
+        - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh

 - label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs)
  key: multiconnector-nixl-offloading-pd-edge-cases-2-gpus
-  timeout_in_minutes: 25
+  timeout_in_minutes: 30
  working_dir: "/vllm-workspace/tests"
  num_devices: 2
  source_file_dependencies:
@@ -203,25 +198,3 @@ steps:
  commands:
    - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
    - bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
-
-# P TP 4 - D DPEP 4 test case for DSv4-Flash
- label: DSv4-Flash Disaggregated DP EP
-  key: dsv4-flash-disaggregated
-  timeout_in_minutes: 60
-  device: h200
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  num_devices: 8
-  env:
-    ENABLE_HMA_FLAG: "1"
-    DP_EP: "1"
-    GPU_MEMORY_UTILIZATION: "0.85"
-    PREFILLER_TP_SIZE: "4"
-    DECODER_TP_SIZE: "4"
-    PREFILL_BLOCK_SIZE: "256"
-    DECODE_BLOCK_SIZE: "256"
-    MODEL_NAMES: "deepseek-ai/DeepSeek-V4-Flash"
-    VLLM_SERVE_EXTRA_ARGS: "--trust-remote-code,--kv-cache-dtype,fp8"
-  commands:
-    - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
-    - bash v1/kv_connector/nixl_integration/run_accuracy_test.sh
@@ -4,7 +4,7 @@ depends_on:
 steps:
 - label: Distributed Comm Ops
  key: distributed-comm-ops
-  timeout_in_minutes: 25
+  timeout_in_minutes: 20
  working_dir: "/vllm-workspace/tests"
  num_devices: 2
  source_file_dependencies:
@@ -18,7 +18,7 @@ steps:

 - label: Distributed DP Tests (2 GPUs)
  key: distributed-dp-tests-2-gpus
-  timeout_in_minutes: 35
+  timeout_in_minutes: 20
  working_dir: "/vllm-workspace/tests"
  num_devices: 2
  source_file_dependencies:
@@ -39,9 +39,7 @@ steps:
  - DP_SIZE=2 pytest -v -s entrypoints/openai/test_multi_api_servers.py
  mirror:
    amd:
-      dind: false
      device: mi300_2
-      timeout_in_minutes: 45
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -57,7 +55,7 @@ steps:

 - label: Distributed Compile + RPC Tests (2 GPUs)
  key: distributed-compile-rpc-tests-2-gpus
-  timeout_in_minutes: 65
+  timeout_in_minutes: 20
  working_dir: "/vllm-workspace/tests"
  num_devices: 2
  source_file_dependencies:
@@ -80,7 +78,7 @@ steps:

 - label: Distributed Torchrun + Shutdown Tests (2 GPUs)
  key: distributed-torchrun-shutdown-tests-2-gpus
-  timeout_in_minutes: 30
+  timeout_in_minutes: 20
  working_dir: "/vllm-workspace/tests"
  num_devices: 2
  source_file_dependencies:
@@ -135,7 +133,7 @@ steps:

 - label: Distributed DP Tests (4 GPUs)
  key: distributed-dp-tests-4-gpus
-  timeout_in_minutes: 45
+  timeout_in_minutes: 30
  working_dir: "/vllm-workspace/tests"
  num_devices: 4
  source_file_dependencies:
@@ -156,7 +154,7 @@ steps:

 - label: Distributed Compile + Comm (4 GPUs)
  key: distributed-compile-comm-4-gpus
-  timeout_in_minutes: 70
+  timeout_in_minutes: 30
  working_dir: "/vllm-workspace/tests"
  num_devices: 4
  source_file_dependencies:
@@ -178,7 +176,7 @@ steps:

 - label: Distributed Tests (8xH100)
  key: distributed-tests-8xh100
-  timeout_in_minutes: 20
+  timeout_in_minutes: 10
  device: h100
  num_devices: 8
  working_dir: "/vllm-workspace/tests"
@@ -214,7 +212,7 @@ steps:

 - label: Distributed Tests (2xH100-2xMI300)
  key: distributed-tests-2xh100-2xmi300
-  timeout_in_minutes: 30
+  timeout_in_minutes: 15
  device: h100
  optional: true
  working_dir: "/vllm-workspace/"
@@ -261,7 +259,7 @@ steps:

 - label: Pipeline + Context Parallelism (4 GPUs)
  key: pipeline-context-parallelism-4-gpus
-  timeout_in_minutes: 55
+  timeout_in_minutes: 60
  working_dir: "/vllm-workspace/tests"
  num_devices: 4
  source_file_dependencies:
@@ -276,7 +274,7 @@ steps:

 - label: RayExecutorV2 (4 GPUs)
  key: rayexecutorv2-4-gpus
-  timeout_in_minutes: 45
+  timeout_in_minutes: 60
  working_dir: "/vllm-workspace/tests"
  num_devices: 4
  source_file_dependencies:
@@ -3,7 +3,7 @@ depends_on:
  - image-build-cpu
 steps:
 - label: Docker Build Metadata
-  timeout_in_minutes: 20
+  timeout_in_minutes: 10
  device: cpu-small
  source_file_dependencies:
    - .buildkite/release-pipeline.yaml
@@ -4,7 +4,7 @@ depends_on:
 steps:
 - label: DeepSeek V2-Lite Sync EPLB Accuracy (4xH100)
  key: deepseek-v2-lite-sync-eplb-accuracy-4xh100
-  timeout_in_minutes: 25
+  timeout_in_minutes: 60
  device: h100
  optional: true
  num_devices: 4
@@ -14,7 +14,7 @@ steps:

 - label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy (4xH100)
  key: qwen3-30b-a3b-fp8-block-sync-eplb-accuracy-4xh100
-  timeout_in_minutes: 25
+  timeout_in_minutes: 60
  device: h100
  optional: true
  num_devices: 4
@@ -24,7 +24,7 @@ steps:

 - label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy (2xB200)
  key: qwen3-30b-a3b-fp8-block-sync-eplb-accuracy-2xb200
-  timeout_in_minutes: 20
+  timeout_in_minutes: 60
  device: b200-k8s
  optional: true
  num_devices: 2
@@ -34,7 +34,7 @@ steps:

 - label: Qwen3-30B-A3B-FP8 DP4 Async EPLB Accuracy
  key: qwen3-30b-a3b-fp8-dp4-async-eplb-accuracy
-  timeout_in_minutes: 25
+  timeout_in_minutes: 60
  device: h100
  optional: true
  num_devices: 4
@@ -44,7 +44,7 @@ steps:

 - label: DeepSeek V2-Lite Prefetch Offload Accuracy (H100)
  key: deepseek-v2-lite-prefetch-offload-accuracy-h100
-  timeout_in_minutes: 20
+  timeout_in_minutes: 60
  device: h100
  optional: true
  num_devices: 1
@@ -4,7 +4,7 @@ depends_on:
 steps:
 - label: Engine
  key: engine
-  timeout_in_minutes: 30
+  timeout_in_minutes: 15
  device: h200_18gb
  source_file_dependencies:
  - vllm/compilation/
@@ -28,15 +28,14 @@ steps:
  - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py test_jit_monitor.py
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 40
+      device: mi325_1
+      timeout_in_minutes: 60
      depends_on:
      - image-build-amd

 - label: Engine (1 GPU)
  key: engine-1-gpu
-  timeout_in_minutes: 45
+  timeout_in_minutes: 30
  source_file_dependencies:
    - vllm/v1/engine/
    - tests/v1/engine/
@@ -45,14 +44,14 @@ steps:
    - pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py
  mirror:
    amd:
-      device: mi250_1
-      timeout_in_minutes: 45
+      device: mi325_1
+      timeout_in_minutes: 40
      depends_on:
      - image-build-amd

 - label: e2e Scheduling (1 GPU)
  key: e2e-scheduling-1-gpu
-  timeout_in_minutes: 53
+  timeout_in_minutes: 30
  device: h200_18gb
  source_file_dependencies:
    - vllm/v1/
@@ -61,15 +60,15 @@ steps:
    - pytest -v -s v1/e2e/general/test_async_scheduling.py
  mirror:
    amd:
-      device: mi250_1
-      timeout_in_minutes: 55
+      device: mi325_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: 40
+  timeout_in_minutes: 30
  source_file_dependencies:
    - vllm/v1/
    - tests/v1/e2e/general/
@@ -77,8 +76,8 @@ steps:
    - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py
  mirror:
    amd:
-      device: mi250_1
-      timeout_in_minutes: 50
+      device: mi325_1
+      timeout_in_minutes: 35
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -88,7 +87,7 @@ steps:

 - label: V1 e2e (2 GPUs)
  key: v1-e2e-2-gpus
-  timeout_in_minutes: 25 # TODO: Fix timeout after we have more confidence in the test stability
+  timeout_in_minutes: 60 # TODO: Fix timeout after we have more confidence in the test stability
  optional: true
  num_devices: 2
  source_file_dependencies:
@@ -115,15 +114,13 @@ steps:
    - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "tensor_parallelism"
  mirror:
    amd:
-      dind: false
      device: mi300_2
-      timeout_in_minutes: 30
      depends_on:
      - image-build-amd

 - label: V1 e2e (4 GPUs)
  key: v1-e2e-4-gpus
-  timeout_in_minutes: 20 # TODO: Fix timeout after we have more confidence in the test stability
+  timeout_in_minutes: 60 # TODO: Fix timeout after we have more confidence in the test stability
  optional: true
  num_devices: 4
  source_file_dependencies:
@@ -151,7 +148,7 @@ steps:

 - label: V1 e2e (4xH100)
  key: v1-e2e-4xh100
-  timeout_in_minutes: 35
+  timeout_in_minutes: 60
  device: h100
  num_devices: 4
  optional: true
@@ -3,9 +3,8 @@ depends_on:
  - image-build
 steps:
 - label: Entrypoints Unit Tests
-  device: h200_35gb
  key: entrypoints-unit-tests
-  timeout_in_minutes: 25
+  timeout_in_minutes: 10
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/entrypoints
@@ -16,9 +15,8 @@ steps:
  - pytest -v -s entrypoints/weight_transfer

 - label: Entrypoints Integration (LLM)
-  device: h200_35gb
  key: entrypoints-integration-llm
-  timeout_in_minutes: 60
+  timeout_in_minutes: 40
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -30,16 +28,16 @@ steps:
  - pytest -v -s entrypoints/llm/offline_mode # Needs to avoid interference with other tests
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 55
+      device: mi325_1
+      # TODO(akaratza): Test after Torch >= 2.12 bump
+      soft_fail: true
      depends_on:
      - image-build-amd

 - label: Entrypoints Integration (API Server)
  key: entrypoints-integration-api-server
  device: h200_35gb
-  timeout_in_minutes: 75
+  timeout_in_minutes: 130
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -52,16 +50,13 @@ steps:
  - pytest -v -s entrypoints/scale_out
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 65
+      device: mi325_1
      depends_on:
      - image-build-amd

 - label: Entrypoints Integration (API Server OpenAI - Part 1)
-  device: h200_35gb
  key: entrypoints-integration-api-server-openai-part-1
-  timeout_in_minutes: 68
+  timeout_in_minutes: 50
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -72,16 +67,14 @@ steps:
  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 65
+      device: mi325_1
+      timeout_in_minutes: 80
      depends_on:
      - image-build-amd

 - label: Entrypoints Integration (API Server OpenAI - Part 2)
-  device: h200_35gb
  key: entrypoints-integration-api-server-openai-part-2
-  timeout_in_minutes: 83
+  timeout_in_minutes: 50
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -93,14 +86,12 @@ steps:
  - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 70
+      device: mi325_1
+      timeout_in_minutes: 80
      depends_on:
      - image-build-amd

 - label: Entrypoints Integration (API Server Generate)
-  device: h200_35gb
  key: entrypoints-integration-api-server-generate
  timeout_in_minutes: 50
  working_dir: "/vllm-workspace/tests"
@@ -117,14 +108,12 @@ steps:
  - pytest -v -s entrypoints/anthropic
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 65
+      device: mi325_1
+      timeout_in_minutes: 60
      depends_on:
      - image-build-amd

 - label: Entrypoints Integration (Responses API)
-  device: h200_35gb
  key: entrypoints-integration-responses-api
  timeout_in_minutes: 50
  working_dir: "/vllm-workspace/tests"
@@ -137,7 +126,7 @@ steps:
 - label: Entrypoints Integration (Speech to Text)
  device: h200_35gb
  key: entrypoints-integration-speech_to_text
-  timeout_in_minutes: 45
+  timeout_in_minutes: 50
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -149,7 +138,7 @@ steps:
 - label: Entrypoints Integration (Multimodal)
  device: h200_35gb
  key: entrypoints-integration-multimodal
-  timeout_in_minutes: 45
+  timeout_in_minutes: 50
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -159,9 +148,8 @@ steps:
  - pytest -v -s entrypoints/multimodal

 - label: Entrypoints Integration (Pooling)
-  device: h200_35gb
  key: entrypoints-integration-pooling
-  timeout_in_minutes: 75
+  timeout_in_minutes: 50
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -172,7 +160,7 @@ steps:

 - label: OpenAI API Correctness
  key: openai-api-correctness
-  timeout_in_minutes: 20
+  timeout_in_minutes: 30
  device: h200_18gb
  source_file_dependencies:
  - csrc/
@@ -181,9 +169,7 @@ steps:
  - pytest -s entrypoints/openai/correctness/
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 30
+      device: mi325_1
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -4,7 +4,7 @@ depends_on:
 steps:
 - label: EPLB Algorithm
  key: eplb-algorithm
-  timeout_in_minutes: 20
+  timeout_in_minutes: 15
  device: h200_18gb
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -16,9 +16,7 @@ steps:
  - pytest -v -s distributed/test_eplb_utils.py
  mirror:
    amd:
-      dind: false
      device: mi300_1
-      timeout_in_minutes: 30
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -29,7 +27,7 @@ steps:

 - label: EPLB Execution # 17min
  key: eplb-execution
-  timeout_in_minutes: 25
+  timeout_in_minutes: 27
  working_dir: "/vllm-workspace/tests"
  num_devices: 4
  source_file_dependencies:
@@ -41,7 +39,7 @@ steps:

 - label: Elastic EP Scaling Test
  key: elastic-ep-scaling-test
-  timeout_in_minutes: 30
+  timeout_in_minutes: 20
  device: h100
  working_dir: "/vllm-workspace/tests"
  num_devices: 4
@@ -4,7 +4,7 @@ depends_on:
 steps:
 - label: vLLM IR Tests
  key: vllm-ir-tests
-  timeout_in_minutes: 35
+  timeout_in_minutes: 10
  device: h200_18gb
  working_dir: "/vllm-workspace/"
  source_file_dependencies:
@@ -15,21 +15,19 @@ steps:
    - pytest -v -s tests/kernels/ir

 - label: Kernels Core Operation Test
-  device: h200_35gb
  key: kernels-core-operation-test
-  timeout_in_minutes: 120
+  timeout_in_minutes: 75
  source_file_dependencies:
  - csrc/
  - tests/kernels/core
  - tests/kernels/test_concat_mla_q.py
  - tests/kernels/test_fused_qk_norm_rope_gate.py
  commands:
-    - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_fused_qk_norm_rope_gate.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
-  parallelism: 3
+    - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py  kernels/test_concat_mla_q.py kernels/test_fused_qk_norm_rope_gate.py

 - label: Kernels MiniMax Reduce RMS Test (2 GPUs)
  key: kernels-minimax-reduce-rms-test-2-gpus
-  timeout_in_minutes: 20
+  timeout_in_minutes: 15
  num_devices: 2
  device: h100
  source_file_dependencies:
@@ -43,7 +41,7 @@ steps:

 - label: Deepseek V4 Kernel Test (H100)
  key: deepseek-v4-kernel-test-h100
-  timeout_in_minutes: 30
+  timeout_in_minutes: 15
  device: h100
  source_file_dependencies:
  - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
@@ -56,7 +54,7 @@ steps:

 - label: Deepseek V4 Kernel Test (B200)
  key: deepseek-v4-kernel-test-b200
-  timeout_in_minutes: 20
+  timeout_in_minutes: 15
  device: b200-k8s
  source_file_dependencies:
  - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
@@ -67,7 +65,7 @@ steps:

 - label: Kernels Attention Test %N
  key: kernels-attention-test
-  timeout_in_minutes: 65
+  timeout_in_minutes: 35
  source_file_dependencies:
  - csrc/attention/
  - vllm/v1/attention
@@ -80,9 +78,8 @@ steps:
  parallelism: 2
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 90
+      device: mi325_1
+      timeout_in_minutes: 55
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -109,7 +106,7 @@ steps:

 - label: Kernels Quantization Test %N
  key: kernels-quantization-test
-  timeout_in_minutes: 60
+  timeout_in_minutes: 90
  source_file_dependencies:
  - csrc/quantization/
  - vllm/model_executor/layers/quantization
@@ -119,9 +116,7 @@ steps:
  parallelism: 2
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 120
+      device: mi325_1
      source_file_dependencies:
      - csrc/quantization/
      - vllm/model_executor/layers/quantization
@@ -136,7 +131,7 @@ steps:

 - label: Kernels MoE Test %N
  key: kernels-moe-test
-  timeout_in_minutes: 50
+  timeout_in_minutes: 25
  source_file_dependencies:
  - csrc/quantization/cutlass_w8a8/moe/
  - csrc/moe/
@@ -151,9 +146,8 @@ steps:
  parallelism: 5
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 55
+      device: mi325_1
+      timeout_in_minutes: 50
      source_file_dependencies:
      - csrc/quantization/cutlass_w8a8/moe/
      - csrc/moe/
@@ -168,9 +162,8 @@ steps:
      - image-build-amd

 - label: Kernels Mamba Test
-  device: h200_35gb
  key: kernels-mamba-test
-  timeout_in_minutes: 40
+  timeout_in_minutes: 45
  source_file_dependencies:
  - csrc/mamba/
  - tests/kernels/mamba
@@ -179,19 +172,19 @@ steps:
    - pytest -v -s kernels/mamba

 - label: Kernels KDA Test
-  timeout_in_minutes: 25
+  timeout_in_minutes: 20
  device: h200_18gb
  source_file_dependencies:
-  - vllm/third_party/flash_linear_attention/ops/kda.py
-  - vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py
-  - vllm/third_party/flash_linear_attention/ops/l2norm.py
+  - vllm/model_executor/layers/fla/ops/kda.py
+  - vllm/model_executor/layers/fla/ops/chunk_delta_h.py
+  - vllm/model_executor/layers/fla/ops/l2norm.py
  - tests/kernels/test_kda.py
  commands:
    - pytest -v -s kernels/test_kda.py

 - label: Kernels DeepGEMM Test (H100)
  key: kernels-deepgemm-test-h100
-  timeout_in_minutes: 35
+  timeout_in_minutes: 45
  device: h100
  num_devices: 1
  source_file_dependencies:
@@ -218,7 +211,7 @@ steps:

 - label: Kernels (B200)
  key: kernels-b200
-  timeout_in_minutes: 80
+  timeout_in_minutes: 30
  working_dir: "/vllm-workspace/"
  device: b200-k8s
  # optional: true
@@ -238,15 +231,6 @@ steps:
  - vllm/v1/attention/backends/mla/flashinfer_mla.py
  - vllm/v1/attention/selector.py
  - vllm/platforms/cuda.py
-  - vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py
-  - vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py
-  - vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_splitk.py
-  - vllm/cute_utils/
-  - vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/
-  - vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_cutedsl.py
-  - tests/kernels/mamba/test_gdn_prefill_cutedsl.py
-  - tests/kernels/test_bf16x3_router_gemm_cutedsl.py
-  - tests/kernels/test_ll_bf16_gemm.py
  - tests/kernels/test_top_k_per_row.py
  commands:
    - nvidia-smi
@@ -275,28 +259,24 @@ steps:
    - pytest -v -s tests/kernels/moe/test_flashinfer_moe.py
    - pytest -v -s tests/kernels/moe/test_trtllm_nvfp4_moe.py
    - pytest -v -s tests/kernels/moe/test_cutedsl_moe.py
-    - pytest -v -s tests/kernels/mamba/test_gdn_prefill_cutedsl.py
-    - pytest -v -s tests/kernels/test_bf16x3_router_gemm_cutedsl.py
-    - pytest -v -s tests/kernels/test_ll_bf16_gemm.py
    # e2e
    - pytest -v -s tests/models/quantization/test_nvfp4.py

 - label: Kernels Helion Test
  key: kernels-helion-test
-  timeout_in_minutes: 115
+  timeout_in_minutes: 30
  device: h100
  source_file_dependencies:
  - vllm/utils/import_utils.py
  - tests/kernels/helion/
  commands:
    - pip install helion==1.1.0
-    - pytest -v -s kernels/helion/ --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
-  parallelism: 2
+    - pytest -v -s kernels/helion/

 
 - label: Kernels FP8 MoE Test (1xH100)
  key: kernels-fp8-moe-test-1xh100
-  timeout_in_minutes: 40
+  timeout_in_minutes: 90
  device: h100
  num_devices: 1
  optional: true
@@ -313,7 +293,7 @@ steps:

 - label: Kernels FP8 MoE Test (2xH100)
  key: kernels-fp8-moe-test-2xh100
-  timeout_in_minutes: 45
+  timeout_in_minutes: 90
  device: h100
  num_devices: 2
  optional: true
@@ -323,7 +303,7 @@ steps:

 - label: Kernels Fp4 MoE Test (B200)
  key: kernels-fp4-moe-test-b200
-  timeout_in_minutes: 25
+  timeout_in_minutes: 60
  device: b200-k8s
  num_devices: 1
  optional: true
@@ -336,7 +316,7 @@ steps:

 - label: Kernels FusedMoE Layer Test (2 H100s)
  key: kernels-fusedmoe-layer-test-2-h100s
-  timeout_in_minutes: 30
+  timeout_in_minutes: 90
  device: h100
  num_devices: 2
  source_file_dependencies:
@@ -5,7 +5,7 @@ steps:
 - label: LM Eval Small Models
  device: h200_35gb
  key: lm-eval-small-models
-  timeout_in_minutes: 45
+  timeout_in_minutes: 75
  source_file_dependencies:
  - csrc/
  - vllm/model_executor/layers/quantization
@@ -14,9 +14,8 @@ steps:
  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.txt
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 45
+      device: mi325_1
+      timeout_in_minutes: 55
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -57,7 +56,7 @@ steps:

 - label: LM Eval Small Models (1xB200)
  key: lm-eval-small-models-1xb200
-  timeout_in_minutes: 50
+  timeout_in_minutes: 120
  device: b200-k8s
  optional: true
  source_file_dependencies:
@@ -79,31 +78,9 @@ steps:
  commands:
  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small-tp.txt

- label: LM Eval PCP (4xB200)
-  key: lm-eval-pcp-4xb200
-  timeout_in_minutes: 360
-  device: b200-k8s
-  num_devices: 4
-  optional: true
-  source_file_dependencies:
-  - csrc/
-  - tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml
-  - tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-EP.yaml
-  - tests/evals/gsm8k/configs/models-pcp.txt
-  - vllm/model_executor/layers/quantization
-  - vllm/config/parallel.py
-  - vllm/distributed/parallel_state.py
-  - vllm/model_executor/layers/attention/mla_attention.py
-  - vllm/model_executor/layers/attention/pcp.py
-  - vllm/v1/worker/gpu/model_runner.py
-  - vllm/v1/worker/gpu/pcp_manager.py
-  autorun_on_main: true
-  commands:
-  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-pcp.txt
-
 - label: LM Eval Large Models EP (2xB200)
  key: lm-eval-large-models-ep-2xb200
-  timeout_in_minutes: 60
+  timeout_in_minutes: 120
  device: b200-k8s
  optional: true
  num_devices: 2
@@ -115,7 +92,7 @@ steps:

 - label: LM Eval Qwen3.5 Models (2xB200)
  key: lm-eval-qwen3-5-models-2xb200
-  timeout_in_minutes: 45
+  timeout_in_minutes: 120
  device: b200-k8s
  optional: true
  num_devices: 2
@@ -126,13 +103,13 @@ steps:
  - vllm/transformers_utils/configs/qwen3_5_moe.py
  - vllm/model_executor/models/qwen3_next.py
  - vllm/model_executor/models/qwen3_next_mtp.py
-  - vllm/third_party/flash_linear_attention/ops/
+  - vllm/model_executor/layers/fla/ops/
  commands:
  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-qwen35-blackwell.txt

 - label: LM Eval Large Models (8xH200)
  key: lm-eval-large-models-8xh200
-  timeout_in_minutes: 50
+  timeout_in_minutes: 60
  device: h200
  optional: true
  num_devices: 8
@@ -140,9 +117,8 @@ steps:
    - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-h200.txt
  mirror:
    amd:
-      dind: false
      device: mi300_8
-      timeout_in_minutes: 40
+      timeout_in_minutes: 180
      depends_on:
      - image-build-amd
      commands:
@@ -176,7 +152,7 @@ steps:

 - label: LM Eval Humming f16 (A100 - TEMPORARY)
  key: lm-eval-humming-f16-a100
-  timeout_in_minutes: 75
+  timeout_in_minutes: 120
  device: a100
  optional: true
  num_devices: 1
@@ -191,7 +167,7 @@ steps:

 - label: LM Eval Humming Act int8 (A100 - TEMPORARY)
  key: lm-eval-humming-act-a100
-  timeout_in_minutes: 45
+  timeout_in_minutes: 120
  device: a100
  optional: true
  num_devices: 1
@@ -206,7 +182,7 @@ steps:

 - label: LM Eval Humming f16 (H100 - TEMPORARY)
  key: lm-eval-humming-f16-h100
-  timeout_in_minutes: 70
+  timeout_in_minutes: 120
  device: h100
  optional: true
  num_devices: 1
@@ -221,7 +197,7 @@ steps:

 - label: LM Eval Humming Act fp8/int8 (H100 - TEMPORARY)
  key: lm-eval-humming-act-h100
-  timeout_in_minutes: 70
+  timeout_in_minutes: 120
  device: h100
  optional: true
  num_devices: 1
@@ -237,7 +213,7 @@ steps:

 - label: LM Eval Humming f16 (B200 - TEMPORARY)
  key: lm-eval-humming-f16-b200
-  timeout_in_minutes: 50
+  timeout_in_minutes: 120
  device: b200-k8s
  optional: true
  num_devices: 1
@@ -252,7 +228,7 @@ steps:

 - label: LM Eval Humming Act fp8/int8 (B200 - TEMPORARY)
  key: lm-eval-humming-act-b200
-  timeout_in_minutes: 50
+  timeout_in_minutes: 120
  device: b200-k8s
  optional: true
  num_devices: 1
@@ -268,7 +244,7 @@ steps:

 - label: LM Eval TurboQuant KV Cache
  key: lm-eval-turboquant-kv-cache
-  timeout_in_minutes: 55
+  timeout_in_minutes: 75
  device: h200_18gb
  source_file_dependencies:
  - vllm/model_executor/layers/quantization/turboquant/
@@ -280,7 +256,7 @@ steps:

 - label: GPQA Eval (GPT-OSS) (2xH100)
  key: gpqa-eval-gpt-oss-2xh100
-  timeout_in_minutes: 35
+  timeout_in_minutes: 120
  device: h100
  optional: true
  num_devices: 2
@@ -294,7 +270,7 @@ steps:

 - label: GPQA Eval (GPT-OSS) (2xB200)
  key: gpqa-eval-gpt-oss-2xb200
-  timeout_in_minutes: 30
+  timeout_in_minutes: 120
  device: b200-k8s
  optional: true
  num_devices: 2
@@ -308,7 +284,7 @@ steps:

 - label: GPQA Eval (GPT-OSS) (DGX Spark)
  key: gpqa-eval-gpt-oss-spark
-  timeout_in_minutes: 35
+  timeout_in_minutes: 120
  device: dgx-spark
  optional: true
  num_devices: 1
@@ -337,7 +313,7 @@ steps:

 - label: LM Eval KV-Offload (2xH100)
  key: kv-offload-medium
-  timeout_in_minutes: 30
+  timeout_in_minutes: 60
  device: h100
  num_devices: 2
  source_file_dependencies:
@@ -351,7 +327,7 @@ steps:

 - label: LM Eval KV-Offload (4xH100)
  key: kv-offload-large
-  timeout_in_minutes: 40
+  timeout_in_minutes: 60
  device: h100
  num_devices: 4
  source_file_dependencies:
@@ -365,7 +341,7 @@ steps:

 - label: MRCR Eval Small Models
  device: h200_35gb
-  timeout_in_minutes: 25
+  timeout_in_minutes: 30
  source_file_dependencies:
  - tests/evals/mrcr/
  commands:
@@ -5,7 +5,7 @@ steps:
 - label: LoRA %N
  device: h200_35gb
  key: lora
-  timeout_in_minutes: 40
+  timeout_in_minutes: 30
  source_file_dependencies:
  - vllm/lora
  - tests/lora
@@ -14,10 +14,9 @@ steps:
  parallelism: 4
  mirror:
    amd:
-      dind: false
-      device: mi300_1
+      device: mi325_1
      working_dir: "/vllm-workspace/tests"
-      timeout_in_minutes: 85
+      timeout_in_minutes: 60
      source_file_dependencies:
      - vllm/lora
      - tests/lora
@@ -28,7 +27,7 @@ steps:

 - label: LoRA TP (Distributed)
  key: lora-tp-distributed
-  timeout_in_minutes: 60
+  timeout_in_minutes: 30
  num_devices: 4
  source_file_dependencies:
  - vllm/lora
@@ -47,4 +46,4 @@ steps:
    - pytest -v -s -x lora/test_qwen3_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
+    - pytest -v -s -x lora/test_qwen35_densemodel_lora.py
@@ -5,7 +5,7 @@ steps:
 - label: V1 Spec Decode
  device: h200_35gb
  key: v1-spec-decode
-  timeout_in_minutes: 40
+  timeout_in_minutes: 30
  source_file_dependencies:
    - vllm/config/
    - vllm/distributed/
@@ -23,15 +23,14 @@ steps:
    - pytest -v -s -m 'not slow_test' v1/spec_decode
  mirror:
    amd:
-      dind: false
      device: mi300_1
-      timeout_in_minutes: 50
+      timeout_in_minutes: 65
      depends_on:
      - image-build-amd

 - label: V1 Sample + Logits
  key: v1-sample-logits
-  timeout_in_minutes: 83
+  timeout_in_minutes: 30
  device: h200_18gb
  source_file_dependencies:
    - vllm/config/
@@ -59,16 +58,13 @@ steps:
    - pytest -v -s v1/test_outputs.py
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 70
+      device: mi325_1
      depends_on:
      - image-build-amd

 - label: V1 Core + KV + Metrics
-  device: h200_35gb
  key: v1-core-kv-metrics
-  timeout_in_minutes: 80
+  timeout_in_minutes: 30
  source_file_dependencies:
    - vllm/config/
    - vllm/distributed/
@@ -92,9 +88,7 @@ steps:
    - tests/v1/kv_offload
    - tests/v1/simple_kv_offload
    - tests/v1/worker
-    - tests/v1/streaming_input
    - tests/v1/kv_connector/unit
-    - tests/v1/ec_connector/unit
    - tests/v1/metrics
    - tests/entrypoints/openai/correctness/test_lmeval.py
  commands:
@@ -106,18 +100,15 @@ steps:
    - pytest -v -s v1/kv_offload
    - pytest -v -s v1/simple_kv_offload
    - pytest -v -s v1/worker
-    - pytest -v -s v1/streaming_input
    - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
-    - pytest -v -s -m 'not cpu_test' v1/ec_connector/unit
    - pytest -v -s -m 'not cpu_test' v1/metrics
    # Integration test for streaming correctness (requires special branch).
    - pip install -U git+https://github.com/vllm-project/lm-evaluation-harness.git@streaming-api
    - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 65
+      device: mi325_1
+      timeout_in_minutes: 60
      depends_on:
      - image-build-amd

@@ -148,7 +139,6 @@ steps:
    - pytest -v -s -m 'cpu_test' v1/core
    - pytest -v -s v1/structured_output
    - pytest -v -s v1/test_serial_utils.py
-    - pytest -v -s v1/cudagraph/test_cudagraph_manager.py
    - pytest -v -s -m 'cpu_test' v1/kv_connector/unit
    - pytest -v -s -m 'cpu_test' v1/metrics

@@ -182,7 +172,7 @@ steps:

 - label: Regression
  key: regression
-  timeout_in_minutes: 30
+  timeout_in_minutes: 20
  device: h200_18gb
  source_file_dependencies:
  - vllm/config/
@@ -205,14 +195,14 @@ steps:
 - label: Examples
  device: h200_35gb
  key: examples
-  timeout_in_minutes: 40
+  timeout_in_minutes: 45
  working_dir: "/vllm-workspace/examples"
  source_file_dependencies:
  - vllm/entrypoints
  - vllm/multimodal
  - examples/
  commands:
-    - pip install --no-deps tensorizer # for tensorizer test
+    - pip install tensorizer # for tensorizer test
     # for basic
    - python3 basic/offline_inference/chat.py
    - python3 basic/offline_inference/generate.py --model facebook/opt-125m
@@ -236,9 +226,7 @@ steps:
    - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 75
+      device: mi325_1
      source_file_dependencies:
      - vllm/entrypoints
      - vllm/multimodal
@@ -249,7 +237,7 @@ steps:

 - label: Metrics, Tracing (2 GPUs)
  key: metrics-tracing-2-gpus
-  timeout_in_minutes: 25
+  timeout_in_minutes: 20
  num_devices: 2
  source_file_dependencies:
  - vllm/config/
@@ -274,11 +262,10 @@ steps:
  - pytest -v -s v1/tracing
  mirror:
    amd:
-      dind: false
-      device: mi300_2
-      timeout_in_minutes: 30
+      device: mi325_2
      depends_on:
      - image-build-amd
+      optional: true

 - label: Python-only Installation
  key: python-only-installation
@@ -293,8 +280,8 @@ steps:
  - bash standalone_tests/python_only_compile.sh
  mirror:
    amd:
-      device: mi250_1
-      timeout_in_minutes: 55
+      device: mi325_1
+      timeout_in_minutes: 20
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -305,7 +292,7 @@ steps:
 - label: Async Engine, Inputs, Utils, Worker
  device: h200_35gb
  key: async-engine-inputs-utils-worker
-  timeout_in_minutes: 25
+  timeout_in_minutes: 50
  source_file_dependencies:
  - vllm/assets/
  - vllm/config/
@@ -332,7 +319,7 @@ steps:
  key: async-engine-inputs-utils-worker-config-cpu
  depends_on:
  - image-build-cpu
-  timeout_in_minutes: 65
+  timeout_in_minutes: 30
  source_file_dependencies:
  - vllm/assets/
  - vllm/config/
@@ -394,7 +381,7 @@ steps:

 - label: Batch Invariance (A100)
  key: batch-invariance-a100
-  timeout_in_minutes: 40
+  timeout_in_minutes: 30
  device: a100
  source_file_dependencies:
    - vllm/v1/attention
@@ -408,7 +395,7 @@ steps:

 - label: Batch Invariance (H100)
  key: batch-invariance-h100
-  timeout_in_minutes: 40
+  timeout_in_minutes: 30
  device: h100
  source_file_dependencies:
    - vllm/v1/attention
@@ -424,7 +411,7 @@ steps:

 - label: Batch Invariance (B200)
  key: batch-invariance-b200
-  timeout_in_minutes: 35
+  timeout_in_minutes: 30
  device: b200-k8s
  source_file_dependencies:
    - vllm/v1/attention
@@ -443,7 +430,7 @@ steps:
 - label: Acceptance Length Test (Large Models) # optional
  device: h200_35gb
  key: acceptance-length-test-large-models
-  timeout_in_minutes: 20
+  timeout_in_minutes: 25
  gpu: h100
  optional: true
  num_gpus: 1
@@ -3,16 +3,13 @@ depends_on:
  - image-build
 steps:
 - label: Model Executor
-  device: h200_35gb
  key: model-executor
-  timeout_in_minutes: 60
+  timeout_in_minutes: 35
  source_file_dependencies:
  - vllm/engine/arg_utils.py
  - vllm/config/model.py
  - vllm/model_executor
-  - vllm/model_executor/warmup
  - tests/model_executor
-  - tests/model_executor/test_jit_warmup.py
  - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
  commands:
    - apt-get update && apt-get install -y curl libsodium23
@@ -28,18 +25,14 @@ steps:
    - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py --timeout=900 --timeout-method=thread
  mirror:
    amd:
-      dind: false
      device: mi300_1
-      timeout_in_minutes: 60
      depends_on:
      - image-build-amd
      source_file_dependencies:
      - vllm/engine/arg_utils.py
      - vllm/config/model.py
      - vllm/model_executor
-      - vllm/model_executor/warmup
      - tests/model_executor
-      - tests/model_executor/test_jit_warmup.py
      - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
      - vllm/_aiter_ops.py
      - vllm/platforms/rocm.py
@@ -5,7 +5,7 @@ steps:
 - label: Model Runner V2 Core Tests
  device: h200_35gb
  key: model-runner-v2-core-tests
-  timeout_in_minutes: 35
+  timeout_in_minutes: 45
  source_file_dependencies:
  - vllm/v1/worker/gpu/
  - vllm/v1/worker/gpu_worker.py
@@ -27,7 +27,7 @@ steps:
 - label: Model Runner V2 Examples
  device: h200_35gb
  key: model-runner-v2-examples
-  timeout_in_minutes: 35
+  timeout_in_minutes: 45
  working_dir: "/vllm-workspace/examples"
  source_file_dependencies:
    - vllm/v1/worker/gpu/
@@ -41,7 +41,7 @@ steps:
  commands:
    - set -x
    - export VLLM_USE_V2_MODEL_RUNNER=1
-    - pip install --no-deps tensorizer # for tensorizer test
+    - pip install tensorizer # for tensorizer test
    - python3 basic/offline_inference/chat.py # for basic
    - python3 basic/offline_inference/generate.py --model facebook/opt-125m
    #- python3 basic/offline_inference/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10  # TODO
@@ -63,7 +63,7 @@ steps:

 - label: Model Runner V2 Distributed (2 GPUs)
  key: model-runner-v2-distributed-2-gpus
-  timeout_in_minutes: 30
+  timeout_in_minutes: 45
  working_dir: "/vllm-workspace/tests"
  num_devices: 2
  source_file_dependencies:
@@ -84,7 +84,7 @@ steps:

 - label: Model Runner V2 Pipeline Parallelism (4 GPUs)
  key: model-runner-v2-pipeline-parallelism-4-gpus
-  timeout_in_minutes: 50
+  timeout_in_minutes: 60
  working_dir: "/vllm-workspace/tests"
  num_devices: 4
  source_file_dependencies:
@@ -4,7 +4,7 @@ depends_on:
 steps:
 - label: Basic Models Tests (Initialization)
  key: basic-models-tests-initialization
-  timeout_in_minutes: 25
+  timeout_in_minutes: 45
  device: h200_18gb
  source_file_dependencies:
  - vllm/
@@ -17,7 +17,7 @@ steps:
 - label: Basic Models Tests (Extra Initialization) %N
  device: h200_35gb
  key: basic-models-tests-extra-initialization
-  timeout_in_minutes: 100
+  timeout_in_minutes: 45
  source_file_dependencies:
  - vllm/model_executor/models/
  - tests/models/test_initialization.py
@@ -27,12 +27,12 @@ steps:
    # subset of supported models (the complement of the small subset in the above
    # test.) Also run if model initialization test file is modified
    - pytest -v -s models/test_initialization.py -k 'not test_can_initialize_small_subset' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
-  parallelism: 4
+  parallelism: 2

 - label: Basic Models Tests (Other)
  device: h200_35gb
  key: basic-models-tests-other
-  timeout_in_minutes: 35
+  timeout_in_minutes: 45
  source_file_dependencies:
  - vllm/
  - tests/models/test_terratorch.py
@@ -42,30 +42,15 @@ steps:
    - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 50
+      device: mi325_1
      depends_on:
      - image-build-amd

- label: Inkling Unit Tests (B200)
-  key: inkling-unit-tests-b200
-  timeout_in_minutes: 40
-  device: b200-k8s
-  source_file_dependencies:
-  - vllm/models/inkling/
-  - vllm/cute_utils/
-  - cmake/external_projects/tml_fa4.cmake
-  - tests/models/inkling/
-  commands:
-    # FA4 kernel tests require SM100; the suite skips them elsewhere.
-    - pytest -v -s models/inkling
-
 - label: Basic Models Test (Other CPU) # 5min
  key: basic-models-test-other-cpu
  depends_on:
  - image-build-cpu
-  timeout_in_minutes: 20
+  timeout_in_minutes: 10
  source_file_dependencies:
  - vllm/
  - tests/models/test_utils.py
@@ -4,7 +4,7 @@ depends_on:
 steps:
 - label: Distributed Model Tests (2 GPUs)
  key: distributed-model-tests-2-gpus
-  timeout_in_minutes: 60
+  timeout_in_minutes: 50
  working_dir: "/vllm-workspace/tests"
  num_devices: 2
  source_file_dependencies:
@@ -4,7 +4,7 @@ depends_on:
 steps:
 - label: Language Models Tests (Standard)
  key: language-models-tests-standard
-  timeout_in_minutes: 30
+  timeout_in_minutes: 25
  device: h200_18gb
  source_file_dependencies:
  - vllm/
@@ -15,16 +15,13 @@ steps:
    - pytest -v -s models/language -m 'core_model and (not slow_test)'
  mirror:
    amd:
-      dind: false
      device: mi300_1
-      timeout_in_minutes: 45
      depends_on:
      - image-build-amd

 - label: Language Models Tests (Extra Standard) %N
-  device: h200_35gb
  key: language-models-tests-extra-standard
-  timeout_in_minutes: 40
+  timeout_in_minutes: 45
  source_file_dependencies:
  - vllm/model_executor/models/
  - tests/models/language/pooling/test_embedding.py
@@ -38,9 +35,7 @@ steps:
  parallelism: 2
  mirror:
    amd:
-      dind: false
      device: mi300_1
-      timeout_in_minutes: 40
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -54,25 +49,25 @@ steps:
      - tests/models/language/pooling/test_classification.py
      - vllm/_aiter_ops.py
      - vllm/platforms/rocm.py
+
 - label: Language Models Tests (Hybrid) %N
-  device: h200_35gb
  key: language-models-tests-hybrid
-  timeout_in_minutes: 65
+  timeout_in_minutes: 75
  source_file_dependencies:
  - vllm/
  - tests/models/language/generation
  commands:
    # Install fast path packages for testing against transformers
+    # Note: also needed to run plamo2 model in vLLM
    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
    - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
-    # Shard the hybrid language model tests that are numerically stable on Hopper.
-    - pytest -v -s models/language/generation -m hybrid_model -k 'not granite-4.0-tiny-preview' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
+    # Shard hybrid language model tests
+    - pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
  parallelism: 2
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 60
+      device: mi325_1
+      timeout_in_minutes: 90
      depends_on:
      - image-build-amd
      commands:
@@ -80,37 +75,24 @@ steps:
      - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
      - pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB

-# Granite 4 hybrid generation is sensitive to hardware-specific Triton SSD
-# autotuning (https://github.com/vllm-project/vllm/issues/25194). Keep this one
-# correctness test on L4 until its H200 output matches the Transformers reference.
- label: Language Models Tests (Granite L4 Compatibility)
-  key: language-models-tests-granite-l4-compatibility
-  timeout_in_minutes: 65
-  source_file_dependencies:
-  - vllm/
-  - tests/models/language/generation
-  commands:
-    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
-    - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
-    - pytest -v -s models/language/generation -m hybrid_model -k 'granite-4.0-tiny-preview'
-
 - label: Language Models Test (Extended Generation) # 80min
  device: h200_35gb
  key: language-models-test-extended-generation
-  timeout_in_minutes: 65
+  timeout_in_minutes: 110
  optional: true
  source_file_dependencies:
  - vllm/
  - tests/models/language/generation
  commands:
    # Install fast path packages for testing against transformers
+    # Note: also needed to run plamo2 model in vLLM
    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
    - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
    - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'

 - label: Language Models Test (PPL)
  key: language-models-test-ppl
-  timeout_in_minutes: 30
+  timeout_in_minutes: 110
  device: h200_18gb
  optional: true
  source_file_dependencies:
@@ -119,10 +101,10 @@ steps:
  commands:
    - pytest -v -s models/language/generation_ppl_test

- label: Language Models Test (Extended Pooling)
+- label: Language Models Test (Extended Pooling)  # 36min
  device: h200_35gb
  key: language-models-test-extended-pooling
-  timeout_in_minutes: 120
+  timeout_in_minutes: 50
  optional: true
  source_file_dependencies:
  - vllm/
@@ -131,15 +113,14 @@ steps:
    - pytest -v -s models/language/pooling -m 'not core_model'
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 95
+      device: mi325_1
+      timeout_in_minutes: 100
      depends_on:
      - image-build-amd

 - label: Language Models Test (MTEB)
  key: language-models-test-mteb
-  timeout_in_minutes: 68
+  timeout_in_minutes: 110
  device: h200_18gb
  optional: true
  source_file_dependencies:
@@ -4,7 +4,7 @@ depends_on:
 steps:
 - label: "Multi-Modal Models (Standard) 1: qwen2"
  key: multi-modal-models-standard-1-qwen2
-  timeout_in_minutes: 68
+  timeout_in_minutes: 45
  device: h200_18gb
  source_file_dependencies:
  - vllm/
@@ -14,15 +14,13 @@ steps:
    - pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 65
+      device: mi325_1
      depends_on:
      - image-build-amd

 - label: "Multi-Modal Models (Standard) 2: qwen3 + gemma"
  key: multi-modal-models-standard-2-qwen3-gemma
-  timeout_in_minutes: 75
+  timeout_in_minutes: 45
  device: h200_18gb
  source_file_dependencies:
  - vllm/
@@ -33,16 +31,14 @@ steps:
    - pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 55
+      device: mi325_1
      depends_on:
      - 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: 40
+  timeout_in_minutes: 45
  source_file_dependencies:
  - vllm/
  - tests/models/multimodal
@@ -51,15 +47,14 @@ steps:
    - pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model
  mirror:
    amd:
-      device: mi250_1
-      timeout_in_minutes: 55
+      device: mi325_1
      depends_on:
      - image-build-amd

 - label: "Multi-Modal Models (Standard) 4: other + whisper"
  device: h200_35gb
  key: multi-modal-models-standard-4-other-whisper
-  timeout_in_minutes: 75
+  timeout_in_minutes: 45
  source_file_dependencies:
  - vllm/
  - tests/models/multimodal
@@ -70,29 +65,26 @@ steps:
    - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model  # Otherwise, mp_method="spawn" doesn't work
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 50
+      device: mi325_1
      depends_on:
      - image-build-amd

- label: Multi-Modal Processor (CPU) %N
+- label: Multi-Modal Processor (CPU)
  key: multi-modal-processor-cpu
  depends_on:
  - image-build-cpu
-  timeout_in_minutes: 125
+  timeout_in_minutes: 60
  source_file_dependencies:
  - vllm/
  - tests/models/multimodal
  - tests/models/registry.py
  device: cpu-medium
  commands:
-    - 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
-  parallelism: 4
+    - pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py

 - label: Multi-Modal Processor # 44min
  key: multi-modal-processor
-  timeout_in_minutes: 98
+  timeout_in_minutes: 60
  device: h200_18gb
  source_file_dependencies:
  - vllm/
@@ -104,7 +96,7 @@ steps:
 - label: Multi-Modal Accuracy Eval (Small Models) # 50min
  device: h200_35gb
  key: multi-modal-accuracy-eval-small-models
-  timeout_in_minutes: 30
+  timeout_in_minutes: 70
  working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
  source_file_dependencies:
  - vllm/multimodal/
@@ -114,9 +106,7 @@ steps:
  - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-mm-small.txt --tp-size=1
  mirror:
    amd:
-      dind: false
      device: mi300_1
-      timeout_in_minutes: 35
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -127,7 +117,6 @@ steps:
      - vllm/model_executor/model_loader/

 - label: Multi-Modal Models (Extended Generation 1)
-  device: h200_35gb
  key: multi-modal-models-extended-generation-1
  optional: true
  source_file_dependencies:
@@ -139,9 +128,7 @@ steps:
    - pytest -v -s models/multimodal/test_mapping.py
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 90
+      device: mi325_1
      depends_on:
      - image-build-amd

@@ -176,8 +163,7 @@ steps:
    - pytest -v -s models/multimodal/pooling -m 'not core_model'
  mirror:
    amd:
-      dind: false
-      device: mi300_1
+      device: mi325_1
      timeout_in_minutes: 60
      depends_on:
      - image-build-amd
@@ -4,7 +4,7 @@ depends_on:
 steps:
 - label: Plugin Tests (2 GPUs)
  key: plugin-tests-2-gpus
-  timeout_in_minutes: 35
+  timeout_in_minutes: 60
  working_dir: "/vllm-workspace/tests"
  num_devices: 2
  source_file_dependencies:
@@ -5,7 +5,7 @@ steps:
 - label: PyTorch Compilation Unit Tests
  device: h200_35gb
  key: pytorch-compilation-unit-tests
-  timeout_in_minutes: 150
+  timeout_in_minutes: 10
  source_file_dependencies:
    - vllm/__init__.py
    - vllm/_aiter_ops.py
@@ -78,7 +78,7 @@ steps:

 - label: PyTorch Compilation Passes Unit Tests
  key: pytorch-compilation-passes-unit-tests
-  timeout_in_minutes: 45
+  timeout_in_minutes: 20
  source_file_dependencies:
    - vllm/__init__.py
    - vllm/_aiter_ops.py
@@ -107,11 +107,16 @@ steps:
    - tests/compile/passes
  commands:
  - pytest -s -v compile/passes --ignore compile/passes/distributed
+  mirror:
+    amd:
+      device: mi300_1
+      timeout_in_minutes: 180
+      depends_on:
+      - image-build-amd

 - label: PyTorch Fullgraph Smoke Test
-  device: h200_35gb
  key: pytorch-fullgraph-smoke-test
-  timeout_in_minutes: 90
+  timeout_in_minutes: 35
  source_file_dependencies:
  - vllm/__init__.py
  - vllm/_aiter_ops.py
@@ -143,46 +148,11 @@ steps:
  # as it is a heavy test that is covered in other steps.
  # Use `find` to launch multiple instances of pytest so that
  # they do not suffer from https://github.com/vllm-project/vllm/issues/28965
-  - "find compile/fullgraph/ -name 'test_*.py' -not -name 'test_full_cudagraph.py' -not -name 'test_full_graph.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
-
-# Hopper-only DeepSeek-V2-Lite cases in this file require two 29.3-GiB model
-# instances and cannot fit a 35GB MIG slice. L4 retains the original coverage:
-# those SM90 cases skip while the architecture-compatible cases still run.
- label: PyTorch Fullgraph CUDAGraph (L4 Compatibility)
-  key: pytorch-fullgraph-cudagraph-l4-compatibility
-  timeout_in_minutes: 60
-  source_file_dependencies:
-  - vllm/__init__.py
-  - vllm/_aiter_ops.py
-  - vllm/_custom_ops.py
-  - vllm/compilation/
-  - vllm/config/
-  - vllm/distributed/
-  - vllm/engine/
-  - vllm/env_override.py
-  - vllm/envs.py
-  - vllm/forward_context.py
-  - vllm/inputs/
-  - vllm/ir/
-  - vllm/kernels/
-  - vllm/logger.py
-  - vllm/model_executor/
-  - vllm/multimodal/
-  - vllm/platforms/
-  - vllm/plugins/
-  - vllm/sampling_params.py
-  - vllm/sequence.py
-  - vllm/transformers_utils/
-  - vllm/triton_utils/
-  - vllm/utils/
-  - vllm/v1/
-  - tests/compile
-  commands:
-  - pytest -s -v compile/fullgraph/test_full_cudagraph.py
+  - "find compile/fullgraph/ -name 'test_*.py' -not -name 'test_full_graph.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"

 - label: PyTorch Fullgraph
  key: pytorch-fullgraph
-  timeout_in_minutes: 40
+  timeout_in_minutes: 30
  device: h200_18gb
  source_file_dependencies:
  - vllm/__init__.py
@@ -227,9 +197,7 @@ steps:
  - bash standalone_tests/pytorch_nightly_dependency.sh
  mirror:
    amd:
-      dind: false
      device: mi300_1
-      timeout_in_minutes: 30
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -3,11 +3,8 @@ depends_on:
  - image-build
 steps:
 - label: Quantization
-  device: h200_35gb
  key: quantization
-  timeout_in_minutes: 75
-  env:
-    VLLM_USE_V2_MODEL_RUNNER: "0"
+  timeout_in_minutes: 90
  source_file_dependencies:
  - csrc/
  - vllm/model_executor/layers/quantization
@@ -22,14 +19,11 @@ steps:
  # TODO(jerryzh168): resolve the above comment
  - uv pip install --system torchao==0.17.0 --index-url https://download.pytorch.org/whl/cu130
  - uv pip install --system conch-triton-kernels
-  # The SM90-only checkpoint currently contains a removed weight_chan_scale
-  # parameter. It was not exercised by the previous L4 job.
-  - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py -k 'not test_compressed_tensors_w4a8_fp8'
+  - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py

 - label: Quantized Fusions
-  device: h200_35gb
  key: quantized-fusions
-  timeout_in_minutes: 20
+  timeout_in_minutes: 30
  source_file_dependencies:
  - tests/fusion
  - vllm/model_executor/layers/fusion
@@ -41,7 +35,7 @@ steps:

 - label: Quantized MoE Test (B200)
  key: quantized-moe-test-b200
-  timeout_in_minutes: 120
+  timeout_in_minutes: 60
  working_dir: "/vllm-workspace/"
  device: b200-k8s
  source_file_dependencies:
@@ -58,11 +52,8 @@ steps:
    - pytest -s -v tests/quantization/test_blackwell_moe.py

 - label: Quantized Models Test
-  device: h200_35gb
  key: quantized-models-test
-  timeout_in_minutes: 65
-  env:
-    VLLM_USE_V2_MODEL_RUNNER: "0"
+  timeout_in_minutes: 60
  source_file_dependencies:
  - vllm/model_executor/layers/quantization
  - tests/models/quantization
@@ -3,7 +3,7 @@ depends_on:
  - image-build
 steps:
 - label: Rust Frontend OpenAI Coverage
-  timeout_in_minutes: 30
+  timeout_in_minutes: 90
  device: h200_18gb
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -38,7 +38,7 @@ steps:
  - pytest -v -s v1/sample/test_logprobs_e2e.py -k "test_prompt_logprobs_e2e_server"

 - label: Rust Frontend Serve/Admin Coverage
-  timeout_in_minutes: 25
+  timeout_in_minutes: 60
  device: h200_18gb
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -67,7 +67,7 @@ steps:
  - pytest -v -s entrypoints/serve/tokenize/test_tokenization.py -k "not tokenizer_info"

 - label: Rust Frontend Core Correctness
-  timeout_in_minutes: 20
+  timeout_in_minutes: 30
  device: h200_18gb
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -81,8 +81,7 @@ steps:
  - pytest -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine

 - label: Rust Frontend Tool Use
-  device: h200_35gb
-  timeout_in_minutes: 25
+  timeout_in_minutes: 60
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - rust/
@@ -96,7 +95,7 @@ steps:
  - 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: 25
+  timeout_in_minutes: 30
  num_devices: 4
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -4,7 +4,7 @@ steps:
 - label: Rust Frontend Cargo Style + Clippy
  key: rust-frontend-cargo-style-clippy
  depends_on: []
-  timeout_in_minutes: 20
+  timeout_in_minutes: 30
  device: cpu-medium
  no_plugin: true
  source_file_dependencies:
@@ -18,7 +18,7 @@ steps:
 - label: Rust Frontend Cargo Tests
  key: rust-frontend-cargo-tests
  depends_on: []
-  timeout_in_minutes: 20
+  timeout_in_minutes: 30
  device: cpu-medium
  no_plugin: true
  source_file_dependencies:
@@ -5,7 +5,7 @@ steps:
 - label: Samplers Test
  device: h200_35gb
  key: samplers-test
-  timeout_in_minutes: 40
+  timeout_in_minutes: 75
  source_file_dependencies:
  - vllm/model_executor/layers
  - vllm/sampling_metadata.py
@@ -19,18 +19,8 @@ steps:
    - VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers
  mirror:
    amd:
-      device: mi250_1
-      timeout_in_minutes: 40
+      device: mi325_1
      depends_on:
      - image-build-amd
-      source_file_dependencies:
-      - vllm/model_executor/layers
-      - vllm/sampling_metadata.py
-      - vllm/v1/sample/
-      - vllm/entrypoints/generate/beam_search/
-      - tests/samplers
-      - tests/conftest.py
-      - vllm/_aiter_ops.py
-      - vllm/platforms/rocm.py
      commands:
      - pytest -v -s samplers
@@ -4,7 +4,7 @@ depends_on:
 steps:
 - label: Spec Decode Eagle
  key: spec-decode-eagle
-  timeout_in_minutes: 25
+  timeout_in_minutes: 30
  device: h200_18gb
  source_file_dependencies:
    - vllm/v1/spec_decode/
@@ -14,9 +14,8 @@ steps:
    - pytest -v -s v1/e2e/spec_decode -k "eagle_correctness"
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 55
+      device: mi325_1
+      timeout_in_minutes: 45
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -30,7 +29,7 @@ steps:

 - label: Spec Decode Eagle Nightly B200
  key: spec-decode-eagle-nightly-b200
-  timeout_in_minutes: 25
+  timeout_in_minutes: 30
  device: b200-k8s
  optional: true
  source_file_dependencies:
@@ -42,7 +41,7 @@ steps:

 - label: Spec Decode Speculators + MTP
  key: spec-decode-speculators-mtp
-  timeout_in_minutes: 20
+  timeout_in_minutes: 30
  device: h200_18gb
  source_file_dependencies:
    - vllm/v1/spec_decode/
@@ -54,9 +53,8 @@ steps:
    - pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 75
+      device: mi325_1
+      timeout_in_minutes: 65
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -84,7 +82,7 @@ steps:
  
 - label: Spec Decode Ngram + Suffix
  key: spec-decode-ngram-suffix
-  timeout_in_minutes: 20
+  timeout_in_minutes: 30
  device: h200_18gb
  source_file_dependencies:
    - vllm/v1/spec_decode/
@@ -94,9 +92,10 @@ steps:
    - pytest -v -s v1/e2e/spec_decode -k "ngram or suffix"
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 35
+      device: mi325_1
+      timeout_in_minutes: 65
+      # TODO(akaratza): Test after Torch >= 2.12 bump
+      soft_fail: true
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -110,7 +109,7 @@ steps:

 - label: Spec Decode Draft Model
  key: spec-decode-draft-model
-  timeout_in_minutes: 45
+  timeout_in_minutes: 30
  device: h200_18gb
  source_file_dependencies:
    - vllm/v1/spec_decode/
@@ -120,9 +119,8 @@ steps:
    - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
  mirror:
    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 55
+      device: mi325_1
+      timeout_in_minutes: 50
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -136,7 +134,7 @@ steps:

 - label: Spec Decode Draft Model Nightly B200
  key: spec-decode-draft-model-nightly-b200
-  timeout_in_minutes: 40
+  timeout_in_minutes: 30
  device: b200-k8s
  optional: true
  source_file_dependencies:
@@ -148,7 +146,7 @@ steps:

 - label: Speculators Correctness
  key: speculators-correctness
-  timeout_in_minutes: 30
+  timeout_in_minutes: 60
  device: h100
  optional: true
  num_devices: 1
@@ -161,7 +159,7 @@ steps:
    - pytest -v -s v1/spec_decode/test_speculators_correctness.py -m slow_test

 - label: Spec Decode MTP hybrid (B200)
-  timeout_in_minutes: 20
+  timeout_in_minutes: 30
  device: b200-k8s
  optional: true
  source_file_dependencies:
@@ -172,19 +170,3 @@ steps:
    - tests/v1/e2e/spec_decode/
  commands:
    - pytest -v -s v1/e2e/spec_decode -k "qwen3_5-hybrid"
-
- label: Spec Decode DeepSeek MTP Parallel Load (B200)
-  key: spec-decode-deepseek-mtp-parallel-load-b200
-  timeout_in_minutes: 30
-  device: b200-k8s
-  optional: true
-  num_devices: 2
-  source_file_dependencies:
-    - vllm/v1/spec_decode/llm_base_proposer.py
-    - vllm/v1/spec_decode/eagle.py
-    - vllm/v1/worker/gpu/spec_decode/eagle/
-    - vllm/model_executor/models/deepseek_mtp.py
-    - vllm/model_executor/models/deepseek_v2.py
-    - tests/v1/e2e/spec_decode/test_mtp_parallel_load.py
-  commands:
-    - pytest -v -s v1/e2e/spec_decode/test_mtp_parallel_load.py
@@ -4,7 +4,7 @@ depends_on:
 steps:
 - label: Weight Loading Multiple GPU  # 33min
  key: weight-loading-multiple-gpu
-  timeout_in_minutes: 50
+  timeout_in_minutes: 45
  working_dir: "/vllm-workspace/tests"
  num_devices: 2
  optional: true
@@ -15,9 +15,7 @@ steps:
    - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models.txt
  mirror:
    amd:
-      dind: false
      device: mi300_2
-      timeout_in_minutes: 35
      depends_on:
      - image-build-amd
      commands:
@@ -3,7 +3,6 @@
 dist
 vllm/*.so
 vllm/vllm-rs
-.git

 # Byte-compiled / optimized / DLL files
 __pycache__/
@@ -47,7 +47,6 @@

 # Rust Frontend
 /rust/ @BugenZhao @njhill
-/rust/src/bench @esmeetu
 /build_rust.sh @BugenZhao @njhill
 /rust-toolchain.toml @BugenZhao @njhill
 /.buildkite/test_areas/rust* @BugenZhao @njhill
@@ -173,7 +172,7 @@ mkdocs.yaml @hmellor
 # Kernels
 /vllm/v1/attention/ops/chunked_prefill_paged_decode.py @tdoublep
 /vllm/v1/attention/ops/triton_unified_attention.py @tdoublep
-/vllm/third_party/flash_linear_attention @ZJY0516 @vadiklyutiy
+/vllm/model_executor/layers/fla @ZJY0516 @vadiklyutiy

 # ROCm related: specify owner with write access to notify AMD folks for careful code review
 /vllm/**/*rocm* @tjtanaa @dllehr-amd
@@ -181,18 +181,6 @@ pull_request_rules:
      add:
        - performance

- name: label-quantization
-  description: Automatically apply quantization label
-  conditions:
-    - label != stale
-    - or:
-      - files~=^vllm/model_executor/layers/quantization/
-      - title~=(?i)quant
-  actions:
-    label:
-      add:
-        - quantization
-
 - name: label-qwen
  description: Automatically apply qwen label
  conditions:
@@ -130,47 +130,6 @@ jobs:
                  },
                ],
              },
-              quantization: {
-                keywords: [
-                  {
-                    term: "quantization",
-                    searchIn: "both"
-                  },
-                  {
-                    term: "quantized",
-                    searchIn: "both"
-                  },
-                ],
-              },
-              "intel-gpu": {
-                // Keyword search - matches whole words only (with word boundaries)
-                keywords: [
-                  {
-                    term: "B50",
-                    searchIn: "both"
-                  },
-                  {
-                    term: "B60",
-                    searchIn: "both"
-                  },
-                  {
-                    term: "B70",
-                    searchIn: "both"
-                  },
-                  {
-                    term: "intel gpu",
-                    searchIn: "both"
-                  },
-                  {
-                    term: "Arc GPU",
-                    searchIn: "both"
-                  },
-                  {
-                    term: "BMG",
-                    searchIn: "both"
-                  },
-                ],
-              },
              // Add more label configurations here as needed
              // example: {
              //   keywords: [...],
@@ -364,7 +323,7 @@ jobs:
            // {users} will be replaced with @mentions
            const ccConfig = {
              rocm: {
-                users: ['hongxiayang', 'tjtanaa', 'vllmellm', 'giuseppegrossi'],
+                users: ['hongxiayang', 'tjtanaa', 'vllmellm'],
                message: 'CC {users} for ROCm-related issue',
              },
              mistral: {
@@ -532,4 +491,4 @@ jobs:
              issue_number: context.issue.number,
              body: message,
            });
-            core.notice(`Requested missing ROCm info from @${author}: ${missing.map(m => m.name).join(', ')}`);
+            core.notice(`Requested missing ROCm info from @${author}: ${missing.map(m => m.name).join(', ')}`);
@@ -18,9 +18,6 @@ vllm/third_party/deep_gemm/
 # fmha_sm100 vendored package built from source
 vllm/third_party/fmha_sm100/

-# tml-fa4 vendored package built from source
-vllm/third_party/tml_fa4/
-
 # triton jit
 .triton

@@ -4,7 +4,7 @@ default_install_hook_types:
 default_stages:
  - pre-commit # Run locally
  - manual # Run in CI
-exclude: 'vllm/third_party/.*|vllm/models/kimi_k3/nvidia/ops/third_party/.*|vllm/models/kimi_k3/amd/ops/third_party/.*'
+exclude: 'vllm/third_party/.*'
 repos:
 - repo: https://github.com/astral-sh/ruff-pre-commit
  rev: v0.14.0
@@ -30,7 +30,7 @@ repos:
  - id: markdownlint-cli2
    language_version: lts
    args: [--fix]
-    exclude: (^|/)CLAUDE\.md$
+    exclude: ^CLAUDE\.md$
 - repo: https://github.com/rhysd/actionlint
  rev: v1.7.7
  hooks:
@@ -210,7 +210,7 @@ repos:
    name: Check SPDX headers
    entry: python tools/pre_commit/check_spdx_header.py
    language: python
-    types_or: [python, rust, proto]
+    types: [python]
  - id: check-root-lazy-imports
    name: Check root lazy imports
    entry: python tools/pre_commit/check_init_lazy_imports.py
@@ -0,0 +1,2 @@
+collect_env.py
+vllm/model_executor/layers/fla/ops/*.py
@@ -68,17 +68,8 @@ endif()
 # requirements.txt files and should be kept consistent.  The ROCm torch
 # versions are derived from docker/Dockerfile.rocm
 #
-set(TORCH_SUPPORTED_VERSION_CUDA "2.13.0")
-set(TORCH_SUPPORTED_VERSION_ROCM "2.13.0")
-# TORCH_NIGHTLY=1 builds run against unpinned nightly wheels, so the supported-
-# version check would always warn. Only treat it as a nightly build when the
-# value is exactly "1" (the bootstrap exports TORCH_NIGHTLY=0 by default, which
-# must NOT suppress the warning for normal builds).
-if (DEFINED ENV{TORCH_NIGHTLY} AND "$ENV{TORCH_NIGHTLY}" STREQUAL "1")
-  set(TORCH_NIGHTLY_BUILD TRUE)
-else()
-  set(TORCH_NIGHTLY_BUILD FALSE)
-endif()
+set(TORCH_SUPPORTED_VERSION_CUDA "2.11.0")
+set(TORCH_SUPPORTED_VERSION_ROCM "2.11.0")

 #
 # Try to find python package with an executable that exactly matches
@@ -114,11 +105,6 @@ find_package(Torch REQUIRED)
 # Supported NVIDIA architectures.
 # This check must happen after find_package(Torch) because that's when CMAKE_CUDA_COMPILER_VERSION gets defined
 if(DEFINED CMAKE_CUDA_COMPILER_VERSION AND
-   CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.4)
-  # Rubin (10.7) can run SM100 family code, but CUDA 13.4 also supports
-  # targeting it directly.
-  set(CUDA_SUPPORTED_ARCHS "7.5;8.0;8.6;8.7;8.9;9.0;10.0;10.7;11.0;12.0")
-elseif(DEFINED CMAKE_CUDA_COMPILER_VERSION AND
   CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.0)
   # starting from CUDA 12.9 and Blackwell (10.0), we use family-specific targets (10.0f, 12.0f, etc)
   # to support the whole generation without specifying all sub-architectures
@@ -189,7 +175,7 @@ endif()
 if (NOT HIP_FOUND AND NOT PYTORCH_FOUND_HIP AND CUDA_FOUND)
  set(VLLM_GPU_LANG "CUDA")

-  if (NOT TORCH_NIGHTLY_BUILD AND NOT Torch_VERSION VERSION_EQUAL ${TORCH_SUPPORTED_VERSION_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()
@@ -202,7 +188,7 @@ elseif(HIP_FOUND OR PYTORCH_FOUND_HIP)
  enable_language(HIP)

  # ROCm 5.X and 6.X
-  if (NOT TORCH_NIGHTLY_BUILD AND ROCM_VERSION_DEV_MAJOR GREATER_EQUAL 5 AND
+  if (ROCM_VERSION_DEV_MAJOR GREATER_EQUAL 5 AND
      Torch_VERSION VERSION_LESS ${TORCH_SUPPORTED_VERSION_ROCM})
    message(WARNING "Pytorch version >= ${TORCH_SUPPORTED_VERSION_ROCM} "
      "expected for ROCm build, saw ${Torch_VERSION} instead.")
@@ -395,7 +381,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
    "csrc/libtorch_stable/cuda_view.cu"
    "csrc/libtorch_stable/cuda_utils_kernels.cu"
    "csrc/libtorch_stable/activation_kernels.cu"
-    "csrc/libtorch_stable/ngram_embedding_kernels.cu"
    "csrc/libtorch_stable/quantization/activation_kernels.cu"
    "csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu"
    "csrc/libtorch_stable/quantization/w8a8/fp8/common.cu"
@@ -416,11 +401,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
    "csrc/libtorch_stable/mamba/selective_scan_fwd.cu"
    "csrc/libtorch_stable/cache_kernels.cu"
    "csrc/libtorch_stable/cache_kernels_fused.cu"
-    "csrc/libtorch_stable/custom_all_gather_reduce_scatter.cu"
-    "csrc/libtorch_stable/custom_all_gather_reduce_scatter_ops.cpp"
    "csrc/libtorch_stable/custom_all_reduce.cu"
-    "csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu"
-    "csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu")
+    "csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")

  if(VLLM_GPU_LANG STREQUAL "CUDA" AND
     DEFINED CMAKE_CUDA_COMPILER_VERSION AND
@@ -428,7 +410,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")

    if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
      cuda_archs_loose_intersection(COOPERATIVE_TOPK_ARCHS
-        "9.0a;10.0f;10.1f;10.3f;10.7f;11.0f;12.0f;12.1f" "${CUDA_ARCHS}")
+        "9.0a;10.0f;10.1f;10.3f;11.0f;12.0f;12.1f" "${CUDA_ARCHS}")
    else()
      cuda_archs_loose_intersection(COOPERATIVE_TOPK_ARCHS
        "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
@@ -703,7 +685,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")

  # DeepSeek V3 fused A GEMM kernel (requires SM 9.0+, Hopper and later)
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0f;10.7f;11.0f;12.0f" "${CUDA_ARCHS}")
+    cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0f;11.0f;12.0f" "${CUDA_ARCHS}")
  else()
    cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
  endif()
@@ -823,7 +805,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
  # The cutlass_scaled_mm kernels for Blackwell SM100 (c3x, i.e. CUTLASS 3.x)
  # require CUDA 12.8 or later
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;10.7f;11.0f" "${CUDA_ARCHS}")
+    cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
  else()
    cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
  endif()
@@ -907,7 +889,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
  endif()

  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;10.7f;11.0f" "${CUDA_ARCHS}")
+    cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
  else()
    cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
  endif()
@@ -932,7 +914,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")

  # moe_data.cu is used by all CUTLASS MoE kernels.
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0f;10.7f;11.0f;12.0f" "${CUDA_ARCHS}")
+    cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0f;11.0f;12.0f" "${CUDA_ARCHS}")
  else()
    cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
  endif()
@@ -989,7 +971,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
  # SM10x/11x FP4 kernels. MXFP4 experts quantization is currently compiled
  # only in this block; SM12x has separate NVFP4 matmul/MoE kernels above.
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(FP4_SM100_ARCHS "10.0f;10.7f;11.0f" "${CUDA_ARCHS}")
+    cuda_archs_loose_intersection(FP4_SM100_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
  else()
    cuda_archs_loose_intersection(FP4_SM100_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
  endif()
@@ -1055,7 +1037,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
  # Runtime dispatch is gated in
  # vllm/v1/attention/backends/mla/cutlass_mla.py.
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(MLA_ARCHS "10.0f;10.7f;11.0f" "${CUDA_ARCHS}")
+    cuda_archs_loose_intersection(MLA_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
  else()
    cuda_archs_loose_intersection(MLA_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
  endif()
@@ -1077,41 +1059,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
    set(MLA_ARCHS)
  endif()

-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(FUSED_KDA_DECODE_ARCHS
-      "9.0a;10.0f;12.0f" "${CUDA_ARCHS}")
-  endif()
-  if(FUSED_KDA_DECODE_ARCHS)
-    set(FUSED_KDA_DECODE_SRC
-      "csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel.cu")
-    set_gencode_flags_for_srcs(
-      SRCS "${FUSED_KDA_DECODE_SRC}"
-      CUDA_ARCHS "${FUSED_KDA_DECODE_ARCHS}")
-    set_property(SOURCE ${FUSED_KDA_DECODE_SRC} APPEND PROPERTY
-      COMPILE_OPTIONS "$<$<COMPILE_LANGUAGE:CUDA>:--use_fast_math>")
-    list(APPEND VLLM_STABLE_EXT_SRC "${FUSED_KDA_DECODE_SRC}")
-    message(STATUS
-      "Building fused KDA decode for archs: ${FUSED_KDA_DECODE_ARCHS}")
-  endif()
-
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(KIMI_K3_ATTN_RES_ARCHS
-      "10.0f" "${CUDA_ARCHS}")
-  endif()
-  if(KIMI_K3_ATTN_RES_ARCHS)
-    set(KIMI_K3_ATTN_RES_SRC
-      "csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu")
-    set_gencode_flags_for_srcs(
-      SRCS "${KIMI_K3_ATTN_RES_SRC}"
-      CUDA_ARCHS "${KIMI_K3_ATTN_RES_ARCHS}")
-    set_property(SOURCE ${KIMI_K3_ATTN_RES_SRC} APPEND PROPERTY
-      COMPILE_OPTIONS
-      "$<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr;--expt-extended-lambda;--use_fast_math>")
-    list(APPEND VLLM_STABLE_EXT_SRC "${KIMI_K3_ATTN_RES_SRC}")
-    message(STATUS
-      "Building Kimi K3 AttnRes for archs: ${KIMI_K3_ATTN_RES_ARCHS}")
-  endif()
-
  # Hadacore kernels
  cuda_archs_loose_intersection(HADACORE_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}")
  if(HADACORE_ARCHS)
@@ -1153,14 +1100,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
      target_compile_definitions(_C_stable_libtorch PRIVATE
        VLLM_ENABLE_COOPERATIVE_TOPK=1)
    endif()
-    if(FUSED_KDA_DECODE_ARCHS)
-      target_compile_definitions(_C_stable_libtorch PRIVATE
-        VLLM_ENABLE_FUSED_KDA_DECODE=1)
-    endif()
-    if(KIMI_K3_ATTN_RES_ARCHS)
-      target_compile_definitions(_C_stable_libtorch PRIVATE
-        VLLM_ENABLE_KIMI_K3_ATTN_RES=1)
-    endif()
    # Needed by CUTLASS kernels
    target_compile_definitions(_C_stable_libtorch PRIVATE
      CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
@@ -1415,7 +1354,6 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
  set(VLLM_ROCM_EXT_SRC
    "csrc/rocm/torch_bindings.cpp"
    "csrc/rocm/skinny_gemms.cu"
-    "csrc/rocm/skinny_gemms_int4.cu"
    "csrc/rocm/attention.cu")

  set(VLLM_ROCM_HAS_GFX1100 OFF)
@@ -1458,9 +1396,7 @@ if (VLLM_GPU_LANG STREQUAL "CUDA")
    include(cmake/external_projects/deepgemm.cmake)
    include(cmake/external_projects/fmha_sm100.cmake)
    include(cmake/external_projects/flashmla.cmake)
-    include(cmake/external_projects/flashkda.cmake)
    include(cmake/external_projects/qutlass.cmake)
-    include(cmake/external_projects/tml_fa4.cmake)

    # vllm-flash-attn should be last as it overwrites some CMake functions
    include(cmake/external_projects/vllm_flash_attn.cmake)
@@ -48,7 +48,7 @@ vLLM is flexible and easy to use with:
 - Tool calling and reasoning parsers
 - OpenAI-compatible API server, plus Anthropic Messages API and gRPC support
 - Efficient multi-LoRA support for dense and MoE layers
- Support for NVIDIA GPUs, AMD GPUs, Intel GPUs, and x86/ARM/PowerPC CPUs. Additionally, diverse hardware plugins such as Google TPUs, Intel Gaudi, IBM Spyre, Huawei Ascend, Rebellions NPU, Apple Silicon, MetaX GPU, and more.
+- Support for NVIDIA GPUs, AMD GPUs, and x86/ARM/PowerPC CPUs. Additionally, diverse hardware plugins such as Google TPUs, Intel Gaudi, IBM Spyre, Huawei Ascend, Rebellions NPU, Apple Silicon, MetaX GPU, and more.

 vLLM seamlessly supports 200+ model architectures on Hugging Face, including:

@@ -75,11 +75,7 @@ def run_mla_benchmark(config: BenchmarkConfig, **kwargs) -> BenchmarkResult:
    from mla_runner import run_mla_benchmark as run_mla

    return run_mla(
-        config.backend,
-        config,
-        prefill_backend=config.prefill_backend,
-        sparse_mla_force_mqa=config.sparse_mla_force_mqa,
-        **kwargs,
+        config.backend, config, prefill_backend=config.prefill_backend, **kwargs
    )


@@ -596,30 +592,6 @@ def main():
        default="profile",
        help="Output file name for ncu profile (default: 'profile').",
    )
-    parser.add_argument(
-        "--torch-profile",
-        action="store_true",
-        default=False,
-        help="Collect a PyTorch profiler Chrome trace for each benchmark run.",
-    )
-    parser.add_argument(
-        "--torch-profile-dir",
-        default=None,
-        help="Directory for PyTorch profiler traces.",
-    )
-    parser.add_argument(
-        "--torch-profile-iters",
-        type=int,
-        default=3,
-        help="Number of forward passes to record per PyTorch profiler trace.",
-    )
-    parser.add_argument(
-        "--sparse-mla-mha-variants",
-        nargs="+",
-        default=None,
-        choices=["dense_mha", "mqa"],
-        help="Sparse MLA variants to run in mha_vs_mqa mode. Defaults to both.",
-    )

    # Parameter sweep (use YAML config for advanced sweeps)
    parser.add_argument(
@@ -669,7 +641,6 @@ def main():

        # Prefill backends (e.g., ["fa3", "fa4"])
        args.prefill_backends = yaml_config.get("prefill_backends", None)
-        args.prefill_backend = yaml_config.get("prefill_backend", None)

        # FP8 output benchmark knobs; CLI wins.
        if args.fp8_output_scale is None:
@@ -712,9 +683,6 @@ def main():
            args.num_q_heads = model.get("num_q_heads", args.num_q_heads)
            args.num_kv_heads = model.get("num_kv_heads", args.num_kv_heads)
            args.block_size = model.get("block_size", args.block_size)
-            args.max_model_len = model.get(
-                "max_model_len", getattr(args, "max_model_len", None)
-            )
            # MLA-specific dimensions
            args.kv_lora_rank = model.get("kv_lora_rank", args.kv_lora_rank)
            args.qk_nope_head_dim = model.get("qk_nope_head_dim", args.qk_nope_head_dim)
@@ -733,21 +701,6 @@ def main():
            args.cuda_graphs = yaml_config["cuda_graphs"]
        if "ncu_profile" in yaml_config:
            args.ncu_profile = yaml_config["ncu_profile"]
-        if "torch_profile" in yaml_config:
-            args.torch_profile = yaml_config["torch_profile"]
-        if "torch_profile_dir" in yaml_config:
-            args.torch_profile_dir = yaml_config["torch_profile_dir"]
-        if "torch_profile_iters" in yaml_config:
-            args.torch_profile_iters = yaml_config["torch_profile_iters"]
-        args.sparse_mla_topk_pattern = yaml_config.get(
-            "sparse_mla_topk_pattern", "random"
-        )
-        args.sparse_mla_dense_mha_max_seq_len = yaml_config.get(
-            "sparse_mla_dense_mha_max_seq_len", None
-        )
-        args.sparse_mla_mha_variants = yaml_config.get(
-            "sparse_mla_mha_variants", args.sparse_mla_mha_variants
-        )

        # Parameter sweep configuration
        if "parameter_sweep" in yaml_config:
@@ -889,6 +842,8 @@ def main():
                        num_kv_heads=args.num_kv_heads,
                        block_size=args.block_size,
                        device=args.device,
+                        repeats=args.repeats,
+                        warmup_iters=args.warmup_iters,
                        profile_memory=args.profile_memory,
                        kv_cache_dtype=args.kv_cache_dtype,
                        use_cuda_graphs=args.cuda_graphs,
@@ -1108,133 +1063,6 @@ def main():
                    f"\n  [yellow]Prefill always faster for batch_size={bs}[/]"
                )

-    # Handle MHA vs MQA comparison mode for sparse MLA
-    elif hasattr(args, "mode") and args.mode == "mha_vs_mqa":
-        console.print("[yellow]Mode: MHA vs MQA comparison for sparse MLA[/]")
-
-        sparse_mla_topk_pattern = getattr(args, "sparse_mla_topk_pattern", "random")
-        dense_mha_max_seq_len = getattr(args, "sparse_mla_dense_mha_max_seq_len", None)
-        prefill_backend = getattr(args, "prefill_backend", None)
-        if prefill_backend:
-            console.print(f"Prefill backend: {prefill_backend}")
-        available_variants = [
-            ("dense_mha", False, "dense"),
-            ("mqa", True, "auto"),
-        ]
-        requested_variants = getattr(args, "sparse_mla_mha_variants", None)
-        if requested_variants is not None:
-            valid_variants = {label for label, _, _ in available_variants}
-            invalid_variants = sorted(set(requested_variants) - valid_variants)
-            if invalid_variants:
-                raise ValueError(
-                    "Invalid sparse_mla_mha_variants entries: "
-                    f"{invalid_variants}. Valid variants are: "
-                    f"{sorted(valid_variants)}"
-                )
-            requested_variant_set = set(requested_variants)
-            variants = [
-                variant
-                for variant in available_variants
-                if variant[0] in requested_variant_set
-            ]
-        else:
-            variants = available_variants
-        formatter = ResultsFormatter(console)
-        total = 0
-        for spec in args.batch_specs:
-            q_len = max(request.q_len for request in parse_batch_spec(spec))
-            for variant_label, _, _ in variants:
-                if (
-                    variant_label == "dense_mha"
-                    and dense_mha_max_seq_len is not None
-                    and q_len > dense_mha_max_seq_len
-                ):
-                    continue
-                total += len(backends)
-
-        with tqdm(total=total, desc="Benchmarking") as pbar:
-            for spec in args.batch_specs:
-                q_len = max(request.q_len for request in parse_batch_spec(spec))
-                for backend in backends:
-                    for variant_label, force_mqa, mha_mode in variants:
-                        if (
-                            variant_label == "dense_mha"
-                            and dense_mha_max_seq_len is not None
-                            and q_len > dense_mha_max_seq_len
-                        ):
-                            continue
-                        config = BenchmarkConfig(
-                            backend=f"{backend}_{variant_label}",
-                            batch_spec=spec,
-                            num_layers=args.num_layers,
-                            head_dim=args.head_dim,
-                            num_q_heads=args.num_q_heads,
-                            num_kv_heads=args.num_kv_heads,
-                            block_size=args.block_size,
-                            device=args.device,
-                            max_model_len=getattr(args, "max_model_len", None),
-                            kv_cache_dtype=args.kv_cache_dtype,
-                            profile_memory=args.profile_memory,
-                            use_cuda_graphs=args.cuda_graphs,
-                            ncu_profile=args.ncu_profile,
-                            torch_profile=args.torch_profile,
-                            torch_profile_dir=args.torch_profile_dir,
-                            torch_profile_iters=args.torch_profile_iters,
-                            warmup_ms=args.warmup_ms,
-                            kv_lora_rank=getattr(args, "kv_lora_rank", None),
-                            qk_nope_head_dim=getattr(args, "qk_nope_head_dim", None),
-                            qk_rope_head_dim=getattr(args, "qk_rope_head_dim", None),
-                            v_head_dim=getattr(args, "v_head_dim", None),
-                            sparse_mla_force_mqa=force_mqa,
-                            sparse_mla_mha_mode=mha_mode,
-                            sparse_mla_dense_mha_max_seq_len=dense_mha_max_seq_len,
-                            sparse_mla_topk_pattern=sparse_mla_topk_pattern,
-                            prefill_backend=prefill_backend,
-                        )
-
-                        # run_mla_benchmark needs the real backend name
-                        from mla_runner import run_mla_benchmark as run_mla
-
-                        run_label = f"{backend}_{variant_label} {spec}"
-                        pbar.set_postfix_str(run_label)
-
-                        try:
-                            result = run_mla(
-                                backend,
-                                config,
-                                prefill_backend=prefill_backend,
-                                sparse_mla_force_mqa=force_mqa,
-                            )
-                        except Exception as e:
-                            result = BenchmarkResult(
-                                config=config,
-                                mean_time=float("inf"),
-                                median_time=float("inf"),
-                                std_time=0,
-                                min_time=float("inf"),
-                                max_time=float("inf"),
-                                error=str(e),
-                            )
-
-                        all_results.append(result)
-                        if args.output_csv:
-                            formatter.save_csv(all_results, args.output_csv)
-                        if args.output_json:
-                            formatter.save_json(all_results, args.output_json)
-
-                        if not result.success:
-                            console.print(
-                                f"[red]Error {backend}_{variant_label} "
-                                f"{spec}: {result.error}[/]"
-                            )
-
-                        pbar.update(1)
-
-        # Display results with variant labels as separate "backends"
-        console.print("\n[bold green]MHA vs MQA Results:[/]")
-        variant_backends = [f"{b}_{v}" for b in backends for v, _, _ in variants]
-        formatter.print_table(all_results, variant_backends)
-
    # Handle model parameter sweep mode
    elif hasattr(args, "model_parameter_sweep") and args.model_parameter_sweep:
        # Model parameter sweep
@@ -4,10 +4,8 @@
 """Common utilities for attention benchmarking."""

 import csv
-import gc
 import json
 import math
-from collections.abc import Sequence
 from dataclasses import asdict, dataclass
 from pathlib import Path
 from typing import Any
@@ -46,13 +44,10 @@ def run_do_bench(
    kwargs: dict[str, Any] = {"return_mode": "all"}
    if use_cuda_graphs:
        result = triton.testing.do_bench_cudagraph(benchmark_fn, **kwargs)
-        gc.collect()
-        torch.accelerator.empty_cache()
    else:
        if warmup_ms is not None:
            kwargs["warmup"] = warmup_ms
        result = triton.testing.do_bench(benchmark_fn, **kwargs)
-    torch.accelerator.synchronize()
    return result


@@ -96,6 +91,42 @@ except ImportError:
    AttentionLayerBase = object  # Fallback


+class MockKVBProj:
+    """Mock KV projection layer for MLA prefill mode.
+
+    Mimics ColumnParallelLinear behavior for kv_b_proj in MLA backends.
+    Projects kv_c_normed to [qk_nope_head_dim + v_head_dim] per head.
+    """
+
+    def __init__(self, num_heads: int, qk_nope_head_dim: int, v_head_dim: int):
+        self.num_heads = num_heads
+        self.qk_nope_head_dim = qk_nope_head_dim
+        self.v_head_dim = v_head_dim
+        self.out_dim = qk_nope_head_dim + v_head_dim
+        self.weight = torch.empty(0, dtype=torch.bfloat16)
+
+    def __call__(self, x: torch.Tensor) -> tuple[torch.Tensor]:
+        """
+        Project kv_c_normed to output space.
+
+        Args:
+            x: Input tensor [num_tokens, kv_lora_rank]
+
+        Returns:
+            Tuple containing output tensor
+                [num_tokens, num_heads, qk_nope_head_dim + v_head_dim]
+        """
+        num_tokens = x.shape[0]
+        result = torch.randn(
+            num_tokens,
+            self.num_heads,
+            self.out_dim,
+            device=x.device,
+            dtype=x.dtype,
+        )
+        return (result,)  # Return as tuple to match ColumnParallelLinear API
+
+
 class MockIndexer:
    """Mock Indexer for sparse MLA backends.

@@ -127,60 +158,6 @@ class MockIndexer:
        )
        self.topk_indices_buffer[:num_tokens] = indices

-    def fill_indices(
-        self,
-        num_tokens: int,
-        max_kv_len: int,
-        pattern: str = "random",
-        requests: Sequence[Any] | None = None,
-    ):
-        if pattern == "random":
-            self.fill_random_indices(num_tokens, max_kv_len)
-            return
-        if pattern == "prefix":
-            indices = torch.arange(
-                self.topk_tokens,
-                dtype=torch.int32,
-                device=self.topk_indices_buffer.device,
-            )
-            indices = (indices % max_kv_len).expand(num_tokens, -1)
-            self.topk_indices_buffer[:num_tokens] = indices
-            return
-        if pattern == "sliding_window":
-            if requests is None:
-                start = max(max_kv_len - self.topk_tokens, 0)
-                indices = torch.arange(
-                    start,
-                    start + self.topk_tokens,
-                    dtype=torch.int32,
-                    device=self.topk_indices_buffer.device,
-                )
-                indices = indices.clamp(max=max_kv_len - 1).expand(num_tokens, -1)
-                self.topk_indices_buffer[:num_tokens] = indices
-                return
-
-            rows = []
-            offsets = torch.arange(
-                self.topk_tokens,
-                dtype=torch.int32,
-                device=self.topk_indices_buffer.device,
-            ) - (self.topk_tokens - 1)
-            for request in requests:
-                q_len = request.q_len
-                kv_len = request.kv_len
-                context_len = kv_len - q_len
-                positions = torch.arange(
-                    context_len,
-                    kv_len,
-                    dtype=torch.int32,
-                    device=self.topk_indices_buffer.device,
-                )
-                row_indices = positions[:, None] + offsets[None, :]
-                rows.append(row_indices.clamp(min=0, max=kv_len - 1))
-            self.topk_indices_buffer[:num_tokens] = torch.cat(rows, dim=0)
-            return
-        raise ValueError(f"Unknown sparse MLA topk pattern: {pattern}")
-

 class MockLayer(AttentionLayerBase):
    """Mock attention layer with scale parameters and impl.
@@ -275,14 +252,10 @@ class BenchmarkConfig:
    num_kv_heads: int
    block_size: int
    device: str
-    max_model_len: int | None = None
    dtype: torch.dtype = torch.float16
    profile_memory: bool = False
-    use_cuda_graphs: bool = True
+    use_cuda_graphs: bool = False
    ncu_profile: bool = False
-    torch_profile: bool = False
-    torch_profile_dir: str | None = None
-    torch_profile_iters: int = 3
    warmup_ms: int | None = None

    # "auto" or "fp8"
@@ -298,10 +271,6 @@ class BenchmarkConfig:
    # Backend-specific tuning
    num_kv_splits: int | None = None  # CUTLASS MLA
    reorder_batch_threshold: int | None = None  # FlashAttn MLA, FlashMLA
-    sparse_mla_force_mqa: bool = False  # Force MQA path for sparse MLA
-    sparse_mla_mha_mode: str = "auto"  # "auto" or "dense"
-    sparse_mla_dense_mha_max_seq_len: int | None = None
-    sparse_mla_topk_pattern: str = "random"  # "random", "prefix", "sliding_window"
    num_splits: int | None = None  # FlashAttention split-K (0=auto, 1=disabled)


@@ -1,474 +0,0 @@
-# Sparse MLA benchmark: forward_mha vs forward_mqa
-#
-# Usage:
-#   python benchmark.py --config configs/mla_sparse_mha_vs_mqa.yaml
-#
-# Heatmap grid:
-#   - batch_size: 1, 2, 4, 8, 16, 32
-#   - seq_len: 32, 64, 128, 256, 512, 1024, 2048
-#   - q_len: powers of two through seq_len
-#
-# Specs with q_len < seq_len include context; the q_len == seq_len diagonal
-# covers pure prefill.
-# The model shape below is the DP case. For the TP8 run, manually change
-# model.num_q_heads from 128 to 16 before rerunning this benchmark.
-
-mode: mha_vs_mqa
-
-model:
-  name: "deepseek-v3"
-  num_layers: 60
-  num_q_heads: 128
-  num_kv_heads: 1
-  head_dim: 576
-  kv_lora_rank: 512
-  qk_nope_head_dim: 128
-  qk_rope_head_dim: 64
-  v_head_dim: 128
-  block_size: 128
-  max_model_len: 2048
-
-batch_specs:
-  # Batch size 1
-  # seq_len = 32
-  - "1q1s32"
-  - "1q2s32"
-  - "1q4s32"
-  - "1q8s32"
-  - "1q16s32"
-  - "1q32"
-  # seq_len = 64
-  - "1q1s64"
-  - "1q2s64"
-  - "1q4s64"
-  - "1q8s64"
-  - "1q16s64"
-  - "1q32s64"
-  - "1q64"
-  # seq_len = 128
-  - "1q1s128"
-  - "1q2s128"
-  - "1q4s128"
-  - "1q8s128"
-  - "1q16s128"
-  - "1q32s128"
-  - "1q64s128"
-  - "1q128"
-  # seq_len = 256
-  - "1q1s256"
-  - "1q2s256"
-  - "1q4s256"
-  - "1q8s256"
-  - "1q16s256"
-  - "1q32s256"
-  - "1q64s256"
-  - "1q128s256"
-  - "1q256"
-  # seq_len = 512
-  - "1q1s512"
-  - "1q2s512"
-  - "1q4s512"
-  - "1q8s512"
-  - "1q16s512"
-  - "1q32s512"
-  - "1q64s512"
-  - "1q128s512"
-  - "1q256s512"
-  - "1q512"
-  # seq_len = 1024
-  - "1q1s1024"
-  - "1q2s1024"
-  - "1q4s1024"
-  - "1q8s1024"
-  - "1q16s1024"
-  - "1q32s1024"
-  - "1q64s1024"
-  - "1q128s1024"
-  - "1q256s1024"
-  - "1q512s1024"
-  - "1q1024"
-  # seq_len = 2048
-  - "1q1s2048"
-  - "1q2s2048"
-  - "1q4s2048"
-  - "1q8s2048"
-  - "1q16s2048"
-  - "1q32s2048"
-  - "1q64s2048"
-  - "1q128s2048"
-  - "1q256s2048"
-  - "1q512s2048"
-  - "1q1024s2048"
-  - "1q2048"
-
-  # Batch size 2
-  # seq_len = 32
-  - "2q1s32"
-  - "2q2s32"
-  - "2q4s32"
-  - "2q8s32"
-  - "2q16s32"
-  - "2q32"
-  # seq_len = 64
-  - "2q1s64"
-  - "2q2s64"
-  - "2q4s64"
-  - "2q8s64"
-  - "2q16s64"
-  - "2q32s64"
-  - "2q64"
-  # seq_len = 128
-  - "2q1s128"
-  - "2q2s128"
-  - "2q4s128"
-  - "2q8s128"
-  - "2q16s128"
-  - "2q32s128"
-  - "2q64s128"
-  - "2q128"
-  # seq_len = 256
-  - "2q1s256"
-  - "2q2s256"
-  - "2q4s256"
-  - "2q8s256"
-  - "2q16s256"
-  - "2q32s256"
-  - "2q64s256"
-  - "2q128s256"
-  - "2q256"
-  # seq_len = 512
-  - "2q1s512"
-  - "2q2s512"
-  - "2q4s512"
-  - "2q8s512"
-  - "2q16s512"
-  - "2q32s512"
-  - "2q64s512"
-  - "2q128s512"
-  - "2q256s512"
-  - "2q512"
-  # seq_len = 1024
-  - "2q1s1024"
-  - "2q2s1024"
-  - "2q4s1024"
-  - "2q8s1024"
-  - "2q16s1024"
-  - "2q32s1024"
-  - "2q64s1024"
-  - "2q128s1024"
-  - "2q256s1024"
-  - "2q512s1024"
-  - "2q1024"
-  # seq_len = 2048
-  - "2q1s2048"
-  - "2q2s2048"
-  - "2q4s2048"
-  - "2q8s2048"
-  - "2q16s2048"
-  - "2q32s2048"
-  - "2q64s2048"
-  - "2q128s2048"
-  - "2q256s2048"
-  - "2q512s2048"
-  - "2q1024s2048"
-  - "2q2048"
-
-  # Batch size 4
-  # seq_len = 32
-  - "4q1s32"
-  - "4q2s32"
-  - "4q4s32"
-  - "4q8s32"
-  - "4q16s32"
-  - "4q32"
-  # seq_len = 64
-  - "4q1s64"
-  - "4q2s64"
-  - "4q4s64"
-  - "4q8s64"
-  - "4q16s64"
-  - "4q32s64"
-  - "4q64"
-  # seq_len = 128
-  - "4q1s128"
-  - "4q2s128"
-  - "4q4s128"
-  - "4q8s128"
-  - "4q16s128"
-  - "4q32s128"
-  - "4q64s128"
-  - "4q128"
-  # seq_len = 256
-  - "4q1s256"
-  - "4q2s256"
-  - "4q4s256"
-  - "4q8s256"
-  - "4q16s256"
-  - "4q32s256"
-  - "4q64s256"
-  - "4q128s256"
-  - "4q256"
-  # seq_len = 512
-  - "4q1s512"
-  - "4q2s512"
-  - "4q4s512"
-  - "4q8s512"
-  - "4q16s512"
-  - "4q32s512"
-  - "4q64s512"
-  - "4q128s512"
-  - "4q256s512"
-  - "4q512"
-  # seq_len = 1024
-  - "4q1s1024"
-  - "4q2s1024"
-  - "4q4s1024"
-  - "4q8s1024"
-  - "4q16s1024"
-  - "4q32s1024"
-  - "4q64s1024"
-  - "4q128s1024"
-  - "4q256s1024"
-  - "4q512s1024"
-  - "4q1024"
-  # seq_len = 2048
-  - "4q1s2048"
-  - "4q2s2048"
-  - "4q4s2048"
-  - "4q8s2048"
-  - "4q16s2048"
-  - "4q32s2048"
-  - "4q64s2048"
-  - "4q128s2048"
-  - "4q256s2048"
-  - "4q512s2048"
-  - "4q1024s2048"
-  - "4q2048"
-
-  # Batch size 8
-  # seq_len = 32
-  - "8q1s32"
-  - "8q2s32"
-  - "8q4s32"
-  - "8q8s32"
-  - "8q16s32"
-  - "8q32"
-  # seq_len = 64
-  - "8q1s64"
-  - "8q2s64"
-  - "8q4s64"
-  - "8q8s64"
-  - "8q16s64"
-  - "8q32s64"
-  - "8q64"
-  # seq_len = 128
-  - "8q1s128"
-  - "8q2s128"
-  - "8q4s128"
-  - "8q8s128"
-  - "8q16s128"
-  - "8q32s128"
-  - "8q64s128"
-  - "8q128"
-  # seq_len = 256
-  - "8q1s256"
-  - "8q2s256"
-  - "8q4s256"
-  - "8q8s256"
-  - "8q16s256"
-  - "8q32s256"
-  - "8q64s256"
-  - "8q128s256"
-  - "8q256"
-  # seq_len = 512
-  - "8q1s512"
-  - "8q2s512"
-  - "8q4s512"
-  - "8q8s512"
-  - "8q16s512"
-  - "8q32s512"
-  - "8q64s512"
-  - "8q128s512"
-  - "8q256s512"
-  - "8q512"
-  # seq_len = 1024
-  - "8q1s1024"
-  - "8q2s1024"
-  - "8q4s1024"
-  - "8q8s1024"
-  - "8q16s1024"
-  - "8q32s1024"
-  - "8q64s1024"
-  - "8q128s1024"
-  - "8q256s1024"
-  - "8q512s1024"
-  - "8q1024"
-  # seq_len = 2048
-  - "8q1s2048"
-  - "8q2s2048"
-  - "8q4s2048"
-  - "8q8s2048"
-  - "8q16s2048"
-  - "8q32s2048"
-  - "8q64s2048"
-  - "8q128s2048"
-  - "8q256s2048"
-  - "8q512s2048"
-  - "8q1024s2048"
-  - "8q2048"
-
-  # Batch size 16
-  # seq_len = 32
-  - "16q1s32"
-  - "16q2s32"
-  - "16q4s32"
-  - "16q8s32"
-  - "16q16s32"
-  - "16q32"
-  # seq_len = 64
-  - "16q1s64"
-  - "16q2s64"
-  - "16q4s64"
-  - "16q8s64"
-  - "16q16s64"
-  - "16q32s64"
-  - "16q64"
-  # seq_len = 128
-  - "16q1s128"
-  - "16q2s128"
-  - "16q4s128"
-  - "16q8s128"
-  - "16q16s128"
-  - "16q32s128"
-  - "16q64s128"
-  - "16q128"
-  # seq_len = 256
-  - "16q1s256"
-  - "16q2s256"
-  - "16q4s256"
-  - "16q8s256"
-  - "16q16s256"
-  - "16q32s256"
-  - "16q64s256"
-  - "16q128s256"
-  - "16q256"
-  # seq_len = 512
-  - "16q1s512"
-  - "16q2s512"
-  - "16q4s512"
-  - "16q8s512"
-  - "16q16s512"
-  - "16q32s512"
-  - "16q64s512"
-  - "16q128s512"
-  - "16q256s512"
-  - "16q512"
-  # seq_len = 1024
-  - "16q1s1024"
-  - "16q2s1024"
-  - "16q4s1024"
-  - "16q8s1024"
-  - "16q16s1024"
-  - "16q32s1024"
-  - "16q64s1024"
-  - "16q128s1024"
-  - "16q256s1024"
-  - "16q512s1024"
-  - "16q1024"
-  # seq_len = 2048
-  - "16q1s2048"
-  - "16q2s2048"
-  - "16q4s2048"
-  - "16q8s2048"
-  - "16q16s2048"
-  - "16q32s2048"
-  - "16q64s2048"
-  - "16q128s2048"
-  - "16q256s2048"
-  - "16q512s2048"
-  - "16q1024s2048"
-  - "16q2048"
-
-  # Batch size 32
-  # seq_len = 32
-  - "32q1s32"
-  - "32q2s32"
-  - "32q4s32"
-  - "32q8s32"
-  - "32q16s32"
-  - "32q32"
-  # seq_len = 64
-  - "32q1s64"
-  - "32q2s64"
-  - "32q4s64"
-  - "32q8s64"
-  - "32q16s64"
-  - "32q32s64"
-  - "32q64"
-  # seq_len = 128
-  - "32q1s128"
-  - "32q2s128"
-  - "32q4s128"
-  - "32q8s128"
-  - "32q16s128"
-  - "32q32s128"
-  - "32q64s128"
-  - "32q128"
-  # seq_len = 256
-  - "32q1s256"
-  - "32q2s256"
-  - "32q4s256"
-  - "32q8s256"
-  - "32q16s256"
-  - "32q32s256"
-  - "32q64s256"
-  - "32q128s256"
-  - "32q256"
-  # seq_len = 512
-  - "32q1s512"
-  - "32q2s512"
-  - "32q4s512"
-  - "32q8s512"
-  - "32q16s512"
-  - "32q32s512"
-  - "32q64s512"
-  - "32q128s512"
-  - "32q256s512"
-  - "32q512"
-  # seq_len = 1024
-  - "32q1s1024"
-  - "32q2s1024"
-  - "32q4s1024"
-  - "32q8s1024"
-  - "32q16s1024"
-  - "32q32s1024"
-  - "32q64s1024"
-  - "32q128s1024"
-  - "32q256s1024"
-  - "32q512s1024"
-  - "32q1024"
-  # seq_len = 2048
-  - "32q1s2048"
-  - "32q2s2048"
-  - "32q4s2048"
-  - "32q8s2048"
-  - "32q16s2048"
-  - "32q32s2048"
-  - "32q64s2048"
-  - "32q128s2048"
-  - "32q256s2048"
-  - "32q512s2048"
-  - "32q1024s2048"
-  - "32q2048"
-
-backends:
-  - FLASHMLA_SPARSE
-
-device: "cuda:0"
-profile_memory: false
-sparse_mla_dense_mha_max_seq_len: 2048
-sparse_mla_topk_pattern: "random"
-
-output:
-  csv: "benchmark_output/mla_sparse_mha_vs_mqa.csv"
-  json: "benchmark_output/mla_sparse_mha_vs_mqa.json"
@@ -9,8 +9,6 @@ needing full VllmConfig integration.
 """

 import statistics
-import tempfile
-from pathlib import Path

 import numpy as np
 import torch
@@ -19,6 +17,7 @@ from common import (
    BenchmarkResult,
    MockHfConfig,
    MockIndexer,
+    MockKVBProj,
    MockLayer,
    run_do_bench,
    run_ncu_profile,
@@ -34,59 +33,8 @@ from vllm.config import (
    VllmConfig,
    set_current_vllm_config,
 )
-from vllm.model_executor.layers.linear import ColumnParallelLinear
 from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum

-
-def _safe_profile_name(value: str) -> str:
-    return "".join(c if c.isalnum() or c in "._-" else "_" for c in value)
-
-
-def _create_kv_b_proj(
-    mla_dims: dict,
-    device: torch.device,
-):
-    kv_b_proj = ColumnParallelLinear(
-        mla_dims["kv_lora_rank"],
-        mla_dims["num_q_heads"]
-        * (mla_dims["qk_nope_head_dim"] + mla_dims["v_head_dim"]),
-        bias=False,
-        params_dtype=torch.bfloat16,
-        quant_config=None,
-        prefix="benchmark.kv_b_proj",
-    ).to(device)
-    with torch.no_grad():
-        kv_b_proj.weight.copy_(torch.randn_like(kv_b_proj.weight))
-    return kv_b_proj
-
-
-def _ensure_single_rank_model_parallel() -> None:
-    import torch.distributed as dist
-
-    from vllm.distributed import (
-        ensure_model_parallel_initialized,
-        init_distributed_environment,
-        model_parallel_is_initialized,
-    )
-
-    if not dist.is_available():
-        return
-    if not dist.is_initialized():
-        with tempfile.NamedTemporaryFile(
-            prefix="vllm_bench_dist_", delete=False
-        ) as init_file:
-            distributed_init_method = f"file://{init_file.name}"
-        init_distributed_environment(
-            world_size=1,
-            rank=0,
-            distributed_init_method=distributed_init_method,
-            local_rank=0,
-            backend="nccl",
-        )
-    if not model_parallel_is_initialized():
-        ensure_model_parallel_initialized(1, 1)
-
-
 # ============================================================================
 # VllmConfig Creation
 # ============================================================================
@@ -118,12 +66,10 @@ def create_minimal_vllm_config(
    block_size: int = 128,
    max_num_seqs: int = 256,
    max_num_batched_tokens: int = 8192,
-    max_model_len: int = 32768,
    mla_dims: dict | None = None,
    index_topk: int | None = None,
    prefill_backend: str | None = None,
    kv_cache_dtype: str = "auto",
-    sparse_mla_force_mqa: bool = False,
 ) -> VllmConfig:
    """
    Create minimal VllmConfig for MLA benchmarks.
@@ -140,8 +86,6 @@ def create_minimal_vllm_config(
        prefill_backend: Prefill backend name (e.g., "fa3", "fa4", "flashinfer",
                        "trtllm"). Configures the attention config to force
                        the specified prefill backend.
-        sparse_mla_force_mqa: If True, forces all sparse MLA tokens through
-                    forward_mqa (even prefill tokens).

    Returns:
        VllmConfig for benchmarking
@@ -187,7 +131,7 @@ def create_minimal_vllm_config(
            trust_remote_code=True,
            dtype="bfloat16",
            seed=0,
-            max_model_len=max_model_len,
+            max_model_len=32768,
            quantization=None,
            enforce_eager=False,
            max_logprobs=20,
@@ -219,7 +163,7 @@ def create_minimal_vllm_config(
    scheduler_config = SchedulerConfig(
        max_num_seqs=max_num_seqs,
        max_num_batched_tokens=max(max_num_batched_tokens, max_num_seqs),
-        max_model_len=max_model_len,
+        max_model_len=32768,
        is_encoder_decoder=False,
        enable_chunked_prefill=True,
    )
@@ -248,9 +192,6 @@ def create_minimal_vllm_config(
                "flash_attn_version"
            ]

-    if sparse_mla_force_mqa:
-        vllm_config.attention_config.sparse_mla_force_mqa = True
-
    return vllm_config


@@ -607,7 +548,12 @@ def _create_backend_impl(
    # Calculate scale
    scale = 1.0 / np.sqrt(mla_dims["qk_nope_head_dim"] + mla_dims["qk_rope_head_dim"])

-    kv_b_proj = _create_kv_b_proj(mla_dims, device)
+    # Create mock kv_b_proj layer for prefill mode
+    mock_kv_b_proj = MockKVBProj(
+        num_heads=mla_dims["num_q_heads"],
+        qk_nope_head_dim=mla_dims["qk_nope_head_dim"],
+        v_head_dim=mla_dims["v_head_dim"],
+    )

    # Create indexer for sparse backends
    indexer = None
@@ -638,7 +584,7 @@ def _create_backend_impl(
        "qk_rope_head_dim": mla_dims["qk_rope_head_dim"],
        "qk_head_dim": mla_dims["qk_nope_head_dim"] + mla_dims["qk_rope_head_dim"],
        "v_head_dim": mla_dims["v_head_dim"],
-        "kv_b_proj": kv_b_proj,
+        "kv_b_proj": mock_kv_b_proj,
    }

    # Add indexer for sparse backends
@@ -839,35 +785,14 @@ def _run_single_benchmark(
    # Fill indexer with random indices for sparse backends
    is_sparse = backend_cfg.get("is_sparse", False)
    if is_sparse and indexer is not None:
-        indexer.fill_indices(
-            total_q,
-            max_kv_len,
-            getattr(config, "sparse_mla_topk_pattern", "random"),
-        )
+        indexer.fill_random_indices(total_q, max_kv_len)

    # Determine which forward methods to use based on metadata.
-    # Non-sparse backends use .decode/.prefill sub-objects.
-    # Sparse backends use num_decode_tokens/num_prefills directly.
-    #
-    # sparse_mla_force_mqa overrides: even for prefill metadata, use MQA.
-    force_mqa = getattr(config, "sparse_mla_force_mqa", False)
-    force_dense_mha = getattr(config, "sparse_mla_mha_mode", "auto") == "dense"
-    if force_mqa:
-        has_decode = True
-        has_prefill = False
-    elif is_sparse:
-        has_decode = metadata.num_decode_tokens > 0
-        has_prefill = metadata.num_prefills > 0
-    else:
-        has_decode = metadata.decode is not None
-        has_prefill = metadata.prefill is not None
+    # Sparse MLA backends always use forward_mqa
+    has_decode = is_sparse or getattr(metadata, "decode", None) is not None
+    has_prefill = not is_sparse and getattr(metadata, "prefill", None) is not None
    if not has_decode and not has_prefill:
        raise RuntimeError("Metadata has neither decode nor prefill metadata")
-    if is_sparse and force_dense_mha and not has_prefill:
-        raise RuntimeError(
-            "Sparse MLA dense_mha benchmark did not produce prefill metadata. "
-            "Check reorder_batch_threshold/path forcing."
-        )

    num_decode = (
        metadata.num_decode_tokens
@@ -946,6 +871,7 @@ def _run_single_benchmark(
                metadata,
                prefill_inputs["k_scale"],
                prefill_fp8_output if fused_output else prefill_inputs["output"],
+                prefill_output_scale if fused_output else None,
            )
            if fused_output:
                out = prefill_fp8_output
@@ -972,48 +898,6 @@ def _run_single_benchmark(
            throughput_tokens_per_sec=0.0,
        )

-    if config.torch_profile:
-        profile_dir = Path(
-            config.torch_profile_dir or "benchmark_outputs/torch_profiles"
-        )
-        profile_dir.mkdir(parents=True, exist_ok=True)
-        trace_name = _safe_profile_name(f"{config.backend}_{config.batch_spec}")
-        trace_path = profile_dir / f"{trace_name}.json"
-        iters = max(config.torch_profile_iters, 1)
-
-        forward_fn()
-        torch.accelerator.synchronize()
-        with torch.profiler.profile(
-            activities=[
-                torch.profiler.ProfilerActivity.CPU,
-                torch.profiler.ProfilerActivity.CUDA,
-            ],
-            record_shapes=True,
-            profile_memory=True,
-            with_stack=False,
-        ) as prof:
-            for _ in range(iters):
-                forward_fn()
-                torch.accelerator.synchronize()
-                prof.step()
-        prof.export_chrome_trace(str(trace_path))
-        print(f"Saved PyTorch profiler trace to {trace_path}")
-        print(
-            prof.key_averages().table(
-                sort_by="cuda_time_total",
-                row_limit=25,
-            )
-        )
-        return BenchmarkResult(
-            config=config,
-            mean_time=0.0,
-            median_time=0.0,
-            std_time=0.0,
-            min_time=0.0,
-            max_time=0.0,
-            throughput_tokens_per_sec=0.0,
-        )
-
    all_ms = run_do_bench(benchmark_fn, config.use_cuda_graphs, config.warmup_ms)

    # Convert ms to seconds per layer
@@ -1036,7 +920,6 @@ def _run_mla_benchmark_batched(
    configs_with_params: list[tuple],  # [(config, threshold, num_splits), ...]
    index_topk: int = 2048,
    prefill_backend: str | None = None,
-    sparse_mla_force_mqa: bool = False,
    output_scale: float | None = None,
    fuse_quant_op: bool = False,
 ) -> list[BenchmarkResult]:
@@ -1057,8 +940,6 @@ def _run_mla_benchmark_batched(
        index_topk: Topk value for sparse MLA backends (default 2048)
        prefill_backend: Prefill backend name (e.g., "fa3", "fa4").
            When set, forces the specified FlashAttention version for prefill.
-        sparse_mla_force_mqa: If True, forces all sparse MLA tokens through
-            forward_mqa (even prefill tokens).

    Returns:
        List of BenchmarkResult objects
@@ -1099,41 +980,21 @@ def _run_mla_benchmark_batched(
        sum(r.q_len for r in parse_batch_spec(cfg.batch_spec))
        for cfg, *_ in configs_with_params
    )
-    max_model_len = max(
-        max_total_q,
-        max(
-            getattr(cfg, "max_model_len", None) or 32768
-            for cfg, *_ in configs_with_params
-        ),
-    )

    # Create and set vLLM config for MLA (reused across all benchmarks)
    vllm_config = create_minimal_vllm_config(
        model_name="deepseek-v3",  # Used only for model path
        block_size=block_size,
        max_num_batched_tokens=max_total_q,
-        max_model_len=max_model_len,
        mla_dims=mla_dims,  # Use custom dims from config or default
        index_topk=index_topk if is_sparse else None,
        prefill_backend=prefill_backend,
        kv_cache_dtype=kv_cache_dtype,
-        sparse_mla_force_mqa=sparse_mla_force_mqa,
    )

    results = []

-    # Initialize workspace manager (needed by metadata builders)
-    from vllm.v1.worker.workspace import (
-        init_workspace_manager,
-        is_workspace_manager_initialized,
-    )
-
-    if not is_workspace_manager_initialized():
-        init_workspace_manager(device)
-
    with set_current_vllm_config(vllm_config):
-        _ensure_single_rank_model_parallel()
-
        # Create backend impl, layer, builder, and indexer (reused across benchmarks)
        impl, layer, builder_instance, indexer = _create_backend_impl(
            backend_cfg,
@@ -1179,20 +1040,9 @@ def _run_mla_benchmark_batched(
        for config, threshold, num_splits in configs_with_params:
            # Set threshold for this benchmark (FlashAttn/FlashMLA only)
            original_threshold = None
-            effective_threshold = threshold
-            force_dense_mha = (
-                is_sparse
-                and getattr(config, "sparse_mla_mha_mode", "auto") == "dense"
-                and not getattr(config, "sparse_mla_force_mqa", False)
-            )
-            if force_dense_mha:
-                # Sparse MLA normally treats q_len <= 1 as decode. Use an
-                # impossible threshold so dense_mha benchmarks actually run
-                # the prefill/MHA path, including q_len=1 short extends.
-                effective_threshold = -1
-            if effective_threshold is not None and builder_instance:
+            if threshold is not None and builder_instance:
                original_threshold = builder_instance.reorder_batch_threshold
-                builder_instance.reorder_batch_threshold = effective_threshold
+                builder_instance.reorder_batch_threshold = threshold

            # Set num_splits for CUTLASS
            original_num_splits = None
@@ -1240,7 +1090,6 @@ def run_mla_benchmark(
    num_kv_splits: int | None = None,
    index_topk: int = 2048,
    prefill_backend: str | None = None,
-    sparse_mla_force_mqa: bool = False,
    output_scale: float | None = None,
    fuse_quant_op: bool = False,
 ) -> BenchmarkResult | list[BenchmarkResult]:
@@ -1262,8 +1111,6 @@ def run_mla_benchmark(
        index_topk: Topk value for sparse MLA backends (default 2048)
        prefill_backend: Prefill backend name (e.g., "fa3", "fa4").
            When set, forces the specified FlashAttention version for prefill.
-        sparse_mla_force_mqa: If True, forces all sparse MLA tokens through
-            forward_mqa (even prefill tokens).
        output_scale: Static per-tensor FP8 scale for prefill output (None = bf16).
        fuse_quant_op: With output_scale set, fuse the FP8 write into the prefill
            kernel vs a standalone post-quant kernel. See _run_single_benchmark.
@@ -1295,7 +1142,6 @@ def run_mla_benchmark(
        configs_with_params,
        index_topk,
        prefill_backend=prefill_backend,
-        sparse_mla_force_mqa=sparse_mla_force_mqa,
        output_scale=output_scale,
        fuse_quant_op=fuse_quant_op,
    )
@@ -69,11 +69,12 @@ def make_inputs(total_tokens, num_reqs, block_size):
    # Output workspace
    dst = torch.zeros(total_tokens, HEAD_DIM, dtype=torch.bfloat16, device="cuda")

+    seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device="cuda")
    workspace_starts_t = torch.tensor(
        workspace_starts, dtype=torch.int32, device="cuda"
    )

-    return cache, dst, block_table, workspace_starts_t
+    return cache, dst, block_table, seq_lens_t, workspace_starts_t


 def bench_scenario(label, num_reqs, total_tokens_list, save_path):
@@ -93,7 +94,7 @@ def bench_scenario(label, num_reqs, total_tokens_list, save_path):
        )
    )
    def bench_fn(total_tokens, provider, num_reqs):
-        cache, dst, block_table, ws_starts = make_inputs(
+        cache, dst, block_table, seq_lens_t, ws_starts = make_inputs(
            total_tokens, num_reqs, BLOCK_SIZE
        )

@@ -101,7 +102,7 @@ def bench_scenario(label, num_reqs, total_tokens_list, save_path):

        ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
            lambda: ops.cp_gather_and_upconvert_fp8_kv_cache(
-                cache, dst, block_table, ws_starts, num_reqs
+                cache, dst, block_table, seq_lens_t, ws_starts, num_reqs
            ),
            quantiles=quantiles,
            rep=500,
@@ -1,367 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Benchmark the Kimi-K3 latent MoE addmm against CuTe residual GEMM.
-
-The benchmark covers ``BF16[M, 3584] @ BF16[7168, 3584].T + BF16[M, 7168]``
-with FP32 accumulation and BF16 output. Both backends execute through CUDA
-Graph replay. Weights and residuals rotate across buffers exceeding L2 so the
-comparison models the full latent MoE projection-and-add path.
-"""
-
-from __future__ import annotations
-
-import argparse
-import dataclasses
-import importlib.util
-import json
-import math
-import statistics
-from collections.abc import Callable, Sequence
-from pathlib import Path
-from typing import Any
-
-import cutlass
-import cutlass.cute as cute
-import torch
-from cuda.bindings import driver as cuda
-from cuda.bindings.driver import CUstream
-from quack.compile_utils import make_fake_tensor
-
-N = 7168
-K = 3584
-
-
-@dataclasses.dataclass(frozen=True, slots=True)
-class Config:
-    block_size: int
-    outputs_per_block: int
-    k_unroll: int
-    vector_width: int = 8
-
-
-def parse_config(value: str) -> Config:
-    try:
-        parts = [int(part) for part in value.split(",")]
-    except ValueError as error:
-        raise argparse.ArgumentTypeError(
-            "config must be BLOCK,OUTPUTS,K_UNROLL[,VECTOR_WIDTH]"
-        ) from error
-    if len(parts) == 3:
-        return Config(*parts)
-    if len(parts) == 4:
-        return Config(*parts)
-    raise argparse.ArgumentTypeError(
-        "config must be BLOCK,OUTPUTS,K_UNROLL[,VECTOR_WIDTH]"
-    )
-
-
-def production_residual_config(m: int) -> Config | None:
-    """The measured Latent-MoE residual config for M, from the K3 table."""
-    from vllm.models.kimi_k3.nvidia.low_latency_gemm import KIMI_K3_PROJECTIONS
-
-    spec = KIMI_K3_PROJECTIONS.get((N, K))
-    config = spec.residual_config(m) if spec is not None else None
-    if config is None:
-        return None
-    return Config(
-        config.block_size,
-        config.outputs_per_block,
-        config.k_unroll,
-        config.vector_width,
-    )
-
-
-def candidate_configs(mode: str, selected: Config | None, m: int) -> list[Config]:
-    if mode == "selected":
-        if selected is not None:
-            return [selected]
-        # No explicit --config: fall back to the production table for this M.
-        config = production_residual_config(m)
-        return [config] if config is not None else []
-    if mode == "baseline":
-        return [Config(224, 4, 2)]
-    return [
-        Config(block_size, outputs_per_block, k_unroll, vector_width)
-        for vector_width in (4, 8)
-        for block_size in (32, 64, 128, 224, 448)
-        if block_size % 32 == 0 and K % (block_size * vector_width) == 0
-        for outputs_per_block in (1, 2, 4, 7, 8)
-        if N % outputs_per_block == 0
-        for k_unroll in (1, 2, 4)
-    ]
-
-
-def load_kernel_class(path: Path):
-    spec = importlib.util.spec_from_file_location("cute_skinny_device", path)
-    if spec is None or spec.loader is None:
-        raise RuntimeError(f"cannot load CuTe kernel from {path}")
-    module = importlib.util.module_from_spec(spec)
-    spec.loader.exec_module(module)
-    return module.CuteSkinnyGemm
-
-
-def stream() -> CUstream:
-    return CUstream(torch.cuda.current_stream().cuda_stream)
-
-
-def compile_kernel(kernel_class, m: int, config: Config, max_registers: int):
-    element_type = cutlass.BFloat16
-    n = cute.sym_int(divisibility=config.outputs_per_block)
-    k = cute.sym_int(divisibility=config.block_size * config.vector_width)
-    a = make_fake_tensor(element_type, (m, k), divisibility=config.vector_width)
-    b = make_fake_tensor(element_type, (n, k), divisibility=config.vector_width)
-    residual = make_fake_tensor(element_type, (m, n), divisibility=1)
-    c = make_fake_tensor(element_type, (m, n), divisibility=1)
-    kernel = kernel_class(
-        element_type=element_type,
-        num_rows=m,
-        block_size=config.block_size,
-        outputs_per_block=config.outputs_per_block,
-        vector_width=config.vector_width,
-        k_unroll=config.k_unroll,
-        has_residual=True,
-        use_pdl=True,
-    )
-    return cute.compile(
-        kernel,
-        a,
-        b,
-        residual,
-        c,
-        stream(),
-        options=(
-            "--enable-tvm-ffi --keep-cubin "
-            f"--ptxas-options -maxrregcount={max_registers} "
-            "--ptxas-options -lineinfo"
-        ),
-    )
-
-
-def resource_usage(compiled) -> dict[str, Any]:
-    executor = getattr(compiled, "_default_executor", None)
-    context = getattr(executor, "exec_context", None)
-    functions = getattr(context, "kernel_functions", None)
-    if not functions:
-        return {"resource_metrics_available": False}
-
-    def attribute(name, function) -> int:
-        error, value = cuda.cuFuncGetAttribute(name, function)
-        if error != cuda.CUresult.CUDA_SUCCESS:
-            raise RuntimeError(f"cuFuncGetAttribute failed with {error}")
-        return int(value)
-
-    registers = [
-        attribute(cuda.CUfunction_attribute.CU_FUNC_ATTRIBUTE_NUM_REGS, function)
-        for function in functions
-    ]
-    local_bytes = [
-        attribute(
-            cuda.CUfunction_attribute.CU_FUNC_ATTRIBUTE_LOCAL_SIZE_BYTES,
-            function,
-        )
-        for function in functions
-    ]
-    return {
-        "resource_metrics_available": True,
-        "registers_per_thread": max(registers, default=0),
-        "spill_bytes": max(local_bytes, default=0),
-    }
-
-
-def rotating_buffer_count(m: int, multiplier: float, limit: int) -> int:
-    properties = torch.cuda.get_device_properties(0)
-    bytes_per_pair = (N * K + m * N) * 2
-    target = math.ceil(multiplier * properties.L2_cache_size)
-    return max(2, min(limit, math.ceil(target / bytes_per_pair)))
-
-
-def graph_samples(
-    launch: Callable[[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], None],
-    activation: torch.Tensor,
-    weights: Sequence[torch.Tensor],
-    residuals: Sequence[torch.Tensor],
-    repeats: int,
-    replays: int,
-) -> tuple[list[float], list[torch.Tensor]]:
-    outputs = [torch.empty_like(residual) for residual in residuals]
-    for weight, residual, output in zip(weights, residuals, outputs):
-        launch(activation, weight, residual, output)
-    torch.cuda.synchronize()
-
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph):
-        for weight, residual, output in zip(weights, residuals, outputs):
-            launch(activation, weight, residual, output)
-    for _ in range(20):
-        graph.replay()
-    torch.cuda.synchronize()
-
-    samples = []
-    for _ in range(repeats):
-        start = torch.cuda.Event(enable_timing=True)
-        end = torch.cuda.Event(enable_timing=True)
-        start.record()
-        for _ in range(replays):
-            graph.replay()
-        end.record()
-        end.synchronize()
-        samples.append(start.elapsed_time(end) * 1000.0 / (replays * len(weights)))
-    return samples, outputs
-
-
-def summarize(samples: Sequence[float]) -> dict[str, Any]:
-    ordered = sorted(samples)
-
-    def percentile(fraction: float) -> float:
-        position = fraction * (len(ordered) - 1)
-        lower = math.floor(position)
-        upper = math.ceil(position)
-        if lower == upper:
-            return ordered[lower]
-        weight = position - lower
-        return ordered[lower] * (1.0 - weight) + ordered[upper] * weight
-
-    mean = statistics.mean(samples)
-    return {
-        "median_us": statistics.median(samples),
-        "p10_us": percentile(0.1),
-        "p90_us": percentile(0.9),
-        "mean_us": mean,
-        "cv_pct": statistics.pstdev(samples) / mean * 100.0,
-        "samples_us": list(samples),
-    }
-
-
-def correctness(
-    output: torch.Tensor,
-    activation: torch.Tensor,
-    weight: torch.Tensor,
-    residual: torch.Tensor,
-) -> dict[str, Any]:
-    actual = output.float()
-    reference = activation.float() @ weight.float().t() + residual.float()
-    error = (actual - reference).abs()
-    scaled_error = error / (reference.abs() + 1.0)
-    cosine = torch.nn.functional.cosine_similarity(
-        actual.flatten(), reference.flatten(), dim=0
-    ).item()
-    return {
-        "valid": cosine > 0.999,
-        "cosine": cosine,
-        "max_abs_error": error.max().item(),
-        "max_scaled_error": scaled_error.max().item(),
-    }
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument("--kernel", type=Path, required=True)
-    parser.add_argument("--output", type=Path, required=True)
-    parser.add_argument(
-        "--mode", choices=("baseline", "sweep", "selected"), default="baseline"
-    )
-    parser.add_argument("--config", type=parse_config)
-    parser.add_argument("--m", type=int, action="append")
-    parser.add_argument("--config-shard", type=int, default=0)
-    parser.add_argument("--num-config-shards", type=int, default=1)
-    parser.add_argument("--repeats", type=int, default=21)
-    parser.add_argument("--replays", type=int, default=200)
-    parser.add_argument("--cache-multiplier", type=float, default=3.0)
-    parser.add_argument("--max-buffers", type=int, default=32)
-    parser.add_argument("--max-registers", type=int, default=64)
-    args = parser.parse_args()
-
-    token_counts = args.m or list(range(1, 17))
-    if any(not 1 <= m <= 16 for m in token_counts):
-        raise ValueError("expected 1 <= M <= 16")
-    if not 0 <= args.config_shard < args.num_config_shards:
-        raise ValueError("config shard must be in [0, num_config_shards)")
-    torch.cuda.set_device(0)
-    if torch.cuda.get_device_capability() != (10, 3):
-        raise RuntimeError("this benchmark requires SM103")
-
-    kernel_class = load_kernel_class(args.kernel)
-    properties = torch.cuda.get_device_properties(0)
-    metadata = {
-        "device": properties.name,
-        "compute_capability": list(torch.cuda.get_device_capability()),
-        "torch_version": torch.__version__,
-        "cuda_version": torch.version.cuda,
-    }
-    args.output.parent.mkdir(parents=True, exist_ok=True)
-    with args.output.open("w", encoding="utf-8") as output_file:
-        for m in token_counts:
-            configs = candidate_configs(args.mode, args.config, m)
-            torch.manual_seed(20260722 + m)
-            count = rotating_buffer_count(m, args.cache_multiplier, args.max_buffers)
-            activation = torch.randn((m, K), device="cuda", dtype=torch.bfloat16)
-            weights = [
-                torch.randn((N, K), device="cuda", dtype=torch.bfloat16)
-                for _ in range(count)
-            ]
-            residuals = [
-                torch.randn((m, N), device="cuda", dtype=torch.bfloat16)
-                for _ in range(count)
-            ]
-            candidates: list[tuple[str, Config | None]] = [("cublas_addmm", None)]
-            candidates.extend(
-                ("cute_residual", config)
-                for index, config in enumerate(configs)
-                if index % args.num_config_shards == args.config_shard
-            )
-            for backend, config in candidates:
-                row: dict[str, Any] = {
-                    "m": m,
-                    "n": N,
-                    "k": K,
-                    "backend": backend,
-                    "mode": args.mode,
-                    "config": dataclasses.asdict(config) if config else {},
-                    "num_buffers": count,
-                    "cache_multiplier": args.cache_multiplier,
-                    **metadata,
-                }
-                try:
-                    if backend == "cublas_addmm":
-                        launch = lambda a, b, residual, c: torch.addmm(
-                            residual, a, b.t(), out=c
-                        )
-                    else:
-                        if config is None:
-                            raise AssertionError("missing CuTe config")
-                        compiled = compile_kernel(
-                            kernel_class, m, config, args.max_registers
-                        )
-                        launch = lambda a, b, residual, c, fn=compiled: fn(
-                            a, b, residual, c, stream()
-                        )
-                        row.update(resource_usage(compiled))
-                    samples, outputs = graph_samples(
-                        launch,
-                        activation,
-                        weights,
-                        residuals,
-                        args.repeats,
-                        args.replays,
-                    )
-                    row.update(
-                        correctness(outputs[0], activation, weights[0], residuals[0])
-                    )
-                    row.update(summarize(samples))
-                except Exception as error:  # noqa: BLE001
-                    row.update(
-                        {
-                            "valid": False,
-                            "error": f"{type(error).__name__}: {error}",
-                        }
-                    )
-                output_file.write(json.dumps(row, sort_keys=True) + "\n")
-                output_file.flush()
-                print(json.dumps(row, sort_keys=True), flush=True)
-
-            del activation, weights, residuals
-            torch.cuda.empty_cache()
-
-
-if __name__ == "__main__":
-    main()
@@ -1,806 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Benchmark the Kimi K3 latent-MoE tail and its up-projection kernels.
-
-The ``up-projection`` subcommand isolates the TP-local dynamic and static-M
-skinny GEMMs. It rotates weights through a working set larger than L2 to model
-successive model layers.
-
-The ``whole-tail`` subcommand measures the distributed operator. Its reference
-path includes two AllReduces, RMSNorm, the replicated up-projection, and the
-final add. CUDA-event samples report the slowest rank so cross-rank skew is
-included.
-
-Examples:
-
-.. code-block:: console
-
-    .venv/bin/python \
-      benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py up-projection
-
-    torchrun --nproc-per-node=8 \
-      benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py whole-tail
-
-For multi-node runs, launch one ``torchrun`` agent per node and use a shared
-rendezvous endpoint.
-"""
-
-from __future__ import annotations
-
-import argparse
-import json
-import math
-import os
-import statistics
-from collections.abc import Callable, Sequence
-from dataclasses import asdict
-from pathlib import Path
-from typing import Any
-
-import cutlass
-import cutlass.utils as utils
-import torch
-import torch.distributed as dist
-import torch.nn.functional as F
-from cuda.bindings import driver as cuda
-
-from vllm.distributed import get_tp_group
-from vllm.distributed.parallel_state import (
-    init_distributed_environment,
-    initialize_model_parallel,
-    set_custom_all_reduce,
-)
-from vllm.model_executor.warmup.cutedsl_warmup import cutedsl_warmup
-from vllm.models.kimi_k3.nvidia.ops import latent_moe_tail
-from vllm.models.kimi_k3.nvidia.ops.cute_dsl.latent_moe_tail import (
-    fused_add_multicast_gemm,
-    fused_add_multicast_skinny_gemm,
-)
-
-HIDDEN_SIZE = 7168
-LATENT_SIZE = 3584
-RMS_EPS = 0.1
-MAX_NUM_TOKENS = 16
-MMA_TILER_MN = (64, 32)
-CLUSTER_SHAPE_MN = (1, 8)
-B_PRIME_STAGES = 2
-
-
-def parse_up_projection_config(
-    value: str,
-) -> fused_add_multicast_skinny_gemm.SkinnyConfig:
-    try:
-        values = [int(part) for part in value.split(",")]
-    except ValueError as error:
-        raise argparse.ArgumentTypeError(
-            "config must be BLOCK,OUTPUTS,K_UNROLL[,VECTOR_WIDTH[,PREFETCH_B]]"
-        ) from error
-    if len(values) in (3, 4):
-        return fused_add_multicast_skinny_gemm.SkinnyConfig(*values)
-    if len(values) == 5 and values[4] in (0, 1):
-        return fused_add_multicast_skinny_gemm.SkinnyConfig(
-            *values[:4],
-            prefetch_b_before_pdl=bool(values[4]),
-        )
-    raise argparse.ArgumentTypeError(
-        "config must be BLOCK,OUTPUTS,K_UNROLL"
-        "[,VECTOR_WIDTH[,PREFETCH_B]], where PREFETCH_B is 0 or 1"
-    )
-
-
-def parse_tail_skinny_config(
-    value: str,
-) -> tuple[int, fused_add_multicast_skinny_gemm.SkinnyConfig]:
-    try:
-        values = [int(part) for part in value.split(",")]
-    except ValueError as error:
-        raise argparse.ArgumentTypeError(
-            "config must be M,BLOCK,OUTPUTS,K_UNROLL[,VECTOR_WIDTH[,PREFETCH_B]]"
-        ) from error
-    if len(values) == 4:
-        num_tokens, *config = values
-        return num_tokens, fused_add_multicast_skinny_gemm.SkinnyConfig(*config)
-    if len(values) == 5:
-        num_tokens, *config = values
-        return num_tokens, fused_add_multicast_skinny_gemm.SkinnyConfig(*config)
-    if len(values) == 6 and values[5] in (0, 1):
-        num_tokens, block, outputs, unroll, vector_width, prefetch = values
-        return num_tokens, fused_add_multicast_skinny_gemm.SkinnyConfig(
-            block,
-            outputs,
-            unroll,
-            vector_width,
-            bool(prefetch),
-        )
-    raise argparse.ArgumentTypeError(
-        "config must be M,BLOCK,OUTPUTS,K_UNROLL"
-        "[,VECTOR_WIDTH[,PREFETCH_B]], where PREFETCH_B is 0 or 1"
-    )
-
-
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description=__doc__)
-    subparsers = parser.add_subparsers(dest="scope", required=True)
-
-    up_projection = subparsers.add_parser(
-        "up-projection",
-        help="Benchmark the isolated TP-local up-projection kernels.",
-    )
-    up_projection.add_argument(
-        "--backend",
-        choices=("dynamic", "skinny", "both"),
-        default="both",
-    )
-    up_projection.add_argument("--tp-size", type=int, default=16)
-    up_projection.add_argument(
-        "--num-tokens",
-        type=int,
-        nargs="+",
-        default=[*range(1, 9), 16],
-    )
-    up_projection.add_argument(
-        "--skinny-config",
-        type=parse_up_projection_config,
-        action="append",
-        help="Benchmark a static-M config for every selected token count.",
-    )
-    up_projection.add_argument("--cache-multiplier", type=float, default=2.0)
-    up_projection.add_argument("--max-weights", type=int, default=64)
-    up_projection.add_argument("--warmup-replays", type=int, default=10)
-    up_projection.add_argument("--samples", type=int, default=31)
-    up_projection.add_argument("--output", type=Path)
-
-    whole_tail = subparsers.add_parser(
-        "whole-tail",
-        help="Benchmark the distributed latent-MoE tail operator.",
-    )
-    whole_tail.add_argument(
-        "--backend",
-        choices=("reference", "fused", "both"),
-        default="both",
-    )
-    whole_tail.add_argument(
-        "--num-tokens",
-        type=int,
-        nargs="+",
-        default=[1, 5, 8, 16],
-    )
-    whole_tail.add_argument("--warmup-replays", type=int, default=20)
-    whole_tail.add_argument("--samples", type=int, default=51)
-    whole_tail.add_argument(
-        "--skinny-max-num-tokens",
-        type=int,
-        nargs="+",
-        help="Override the fused operator's static-M cutoff; use 0 for dynamic-only.",
-    )
-    whole_tail.add_argument(
-        "--skinny-config",
-        type=parse_tail_skinny_config,
-        action="append",
-        help="Override one static-M config for tuning.",
-    )
-    whole_tail.add_argument("--output", type=Path)
-    return parser.parse_args()
-
-
-def percentile(samples: Sequence[float], fraction: float) -> float:
-    ordered = sorted(samples)
-    position = fraction * (len(ordered) - 1)
-    lower = math.floor(position)
-    upper = math.ceil(position)
-    if lower == upper:
-        return ordered[lower]
-    upper_weight = position - lower
-    return ordered[lower] * (1.0 - upper_weight) + ordered[upper] * upper_weight
-
-
-def summarize(samples_us: Sequence[float]) -> dict[str, Any]:
-    mean_us = statistics.mean(samples_us)
-    return {
-        "median_us": statistics.median(samples_us),
-        "p10_us": percentile(samples_us, 0.1),
-        "p90_us": percentile(samples_us, 0.9),
-        "mean_us": mean_us,
-        "cv_pct": statistics.pstdev(samples_us) / mean_us * 100.0,
-        "samples_us": list(samples_us),
-    }
-
-
-def rotating_weight_count(
-    shard_size: int,
-    cache_multiplier: float,
-    limit: int,
-) -> int:
-    properties = torch.cuda.get_device_properties(
-        torch.accelerator.current_device_index()
-    )
-    weight_bytes = shard_size * LATENT_SIZE * 2
-    target_bytes = math.ceil(properties.L2_cache_size * cache_multiplier)
-    return max(2, min(limit, math.ceil(target_bytes / weight_bytes)))
-
-
-def capture_up_projection_graph(
-    launches: Sequence[Callable[[], None]],
-) -> torch.cuda.CUDAGraph:
-    for launch in launches:
-        launch()
-    torch.accelerator.synchronize()
-
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph):
-        for launch in launches:
-            launch()
-    torch.accelerator.synchronize()
-    return graph
-
-
-def benchmark_up_projection_graph(
-    graph: torch.cuda.CUDAGraph,
-    *,
-    operations_per_replay: int,
-    warmup_replays: int,
-    samples: int,
-) -> dict[str, Any]:
-    for _ in range(warmup_replays):
-        graph.replay()
-    torch.accelerator.synchronize()
-
-    samples_us = []
-    start = torch.cuda.Event(enable_timing=True)
-    end = torch.cuda.Event(enable_timing=True)
-    for _ in range(samples):
-        start.record()
-        graph.replay()
-        end.record()
-        end.synchronize()
-        samples_us.append(start.elapsed_time(end) * 1000.0 / operations_per_replay)
-    return summarize(samples_us)
-
-
-class DynamicKernel:
-    def __init__(
-        self,
-        shard_size: int,
-        mailbox: torch.Tensor,
-        shared_shard: torch.Tensor,
-    ) -> None:
-        self.shard_size = shard_size
-        self.mailbox = mailbox
-        self.mailbox_c = fused_add_multicast_gemm._as_cute(mailbox)
-        compile_latent = torch.empty(
-            (1, MAX_NUM_TOKENS, LATENT_SIZE),
-            dtype=torch.bfloat16,
-            device=mailbox.device,
-        )
-        compile_weight = torch.empty(
-            (1, shard_size, LATENT_SIZE),
-            dtype=torch.bfloat16,
-            device=mailbox.device,
-        )
-        cluster_size = math.prod(CLUSTER_SHAPE_MN)
-        max_active_clusters = utils.HardwareInfo().get_max_active_clusters(cluster_size)
-        self.compiled = fused_add_multicast_gemm.compile_kernel(
-            (MAX_NUM_TOKENS, shard_size, LATENT_SIZE, 1),
-            fused_add_multicast_gemm._as_cute(
-                compile_latent,
-                dynamic_m=True,
-            ),
-            fused_add_multicast_gemm._as_cute(compile_weight),
-            self.mailbox_c,
-            fused_add_multicast_gemm._as_cute(shared_shard),
-            HIDDEN_SIZE,
-            shard_size,
-            MMA_TILER_MN,
-            CLUSTER_SHAPE_MN,
-            max_active_clusters,
-            B_PRIME_STAGES,
-        )
-
-    def launch(
-        self,
-        latent: torch.Tensor,
-        weight: torch.Tensor,
-        shared_shard: torch.Tensor,
-    ) -> None:
-        stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
-        self.compiled(
-            fused_add_multicast_gemm._as_cute(
-                latent.unsqueeze(0),
-                dynamic_m=True,
-            ),
-            fused_add_multicast_gemm._as_cute(weight.unsqueeze(0)),
-            self.mailbox_c,
-            fused_add_multicast_gemm._as_cute(shared_shard),
-            cutlass.Int64(latent.shape[0]),
-            cutlass.Int64(self.mailbox.data_ptr()),
-            stream,
-        )
-
-
-class SkinnyKernel:
-    def __init__(
-        self,
-        num_tokens: int,
-        shard_size: int,
-        config: fused_add_multicast_skinny_gemm.SkinnyConfig,
-    ) -> None:
-        self.compiled = fused_add_multicast_skinny_gemm.compile_kernel(
-            num_rows=num_tokens,
-            latent_dim=LATENT_SIZE,
-            hidden_dim=HIDDEN_SIZE,
-            shard_dim=shard_size,
-            config=config,
-        )
-
-    def launch(
-        self,
-        latent: torch.Tensor,
-        weight: torch.Tensor,
-        shared_shard: torch.Tensor,
-        mailbox: torch.Tensor,
-    ) -> None:
-        self.compiled(
-            fused_add_multicast_skinny_gemm._as_cute(latent),
-            fused_add_multicast_skinny_gemm._as_cute(weight),
-            fused_add_multicast_skinny_gemm._as_cute(shared_shard),
-            cutlass.Int64(mailbox.data_ptr()),
-            cuda.CUstream(torch.cuda.current_stream().cuda_stream),
-        )
-
-
-def check_up_projection_output(
-    actual: torch.Tensor,
-    latent: torch.Tensor,
-    weight: torch.Tensor,
-    shared_shard: torch.Tensor,
-) -> None:
-    gemm = F.linear(latent.float(), weight.float()).to(torch.bfloat16)
-    expected = (gemm.float() + shared_shard.float()).to(torch.bfloat16)
-    torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2)
-
-
-def make_up_projection_launches(
-    launch: Callable[[torch.Tensor, torch.Tensor, torch.Tensor], None],
-    latent: torch.Tensor,
-    weights: Sequence[torch.Tensor],
-    shared_shard: torch.Tensor,
-) -> list[Callable[[], None]]:
-    return [
-        lambda weight=weight: launch(latent, weight, shared_shard) for weight in weights
-    ]
-
-
-def benchmark_up_projection(args: argparse.Namespace) -> None:
-    if args.tp_size <= 0 or HIDDEN_SIZE % args.tp_size:
-        raise ValueError("TP size must be positive and divide the hidden size")
-    if any(not 1 <= num_tokens <= MAX_NUM_TOKENS for num_tokens in args.num_tokens):
-        raise ValueError("--num-tokens values must be in [1, 16]")
-    if args.cache_multiplier <= 0 or args.max_weights <= 0:
-        raise ValueError("cache multiplier and max weights must be positive")
-    if args.warmup_replays < 0 or args.samples <= 0:
-        raise ValueError("warmup replays must be nonnegative and samples positive")
-
-    torch.accelerator.set_device_index(0)
-    device = torch.device("cuda", 0)
-    if torch.cuda.get_device_capability(device)[0] != 10:
-        raise RuntimeError("Kimi K3 latent-MoE tail requires SM100")
-
-    shard_size = HIDDEN_SIZE // args.tp_size
-    weight_count = rotating_weight_count(
-        shard_size,
-        args.cache_multiplier,
-        args.max_weights,
-    )
-    torch.manual_seed(20260726)
-    weights = [
-        torch.randn(
-            (shard_size, LATENT_SIZE),
-            dtype=torch.bfloat16,
-            device=device,
-        )
-        / LATENT_SIZE**0.5
-        for _ in range(weight_count)
-    ]
-    mailbox = torch.empty(
-        (1, MAX_NUM_TOKENS, HIDDEN_SIZE),
-        dtype=torch.bfloat16,
-        device=device,
-    )
-    shared = torch.randn(
-        (MAX_NUM_TOKENS, HIDDEN_SIZE),
-        dtype=torch.bfloat16,
-        device=device,
-    )
-    shared_shard = shared[:, :shard_size]
-    use_dynamic = args.backend in ("dynamic", "both")
-    use_skinny = args.backend in ("skinny", "both")
-    dynamic_kernel = (
-        DynamicKernel(shard_size, mailbox, shared_shard) if use_dynamic else None
-    )
-
-    results = []
-    for num_tokens in args.num_tokens:
-        latent = torch.randn(
-            (num_tokens, LATENT_SIZE),
-            dtype=torch.bfloat16,
-            device=device,
-        )
-        result: dict[str, Any] = {"num_tokens": num_tokens}
-        if dynamic_kernel is not None:
-            launches = make_up_projection_launches(
-                dynamic_kernel.launch,
-                latent,
-                weights,
-                shared_shard,
-            )
-            graph = capture_up_projection_graph(launches)
-            result["dynamic"] = benchmark_up_projection_graph(
-                graph,
-                operations_per_replay=len(launches),
-                warmup_replays=args.warmup_replays,
-                samples=args.samples,
-            )
-            check_up_projection_output(
-                mailbox[0, :num_tokens, :shard_size],
-                latent,
-                weights[-1],
-                shared_shard[:num_tokens],
-            )
-        if use_skinny:
-            configs = args.skinny_config or [
-                fused_add_multicast_skinny_gemm.config_for_m(
-                    num_tokens,
-                    shard_size,
-                )
-            ]
-            skinny_results = []
-            for config in configs:
-                skinny_kernel = SkinnyKernel(num_tokens, shard_size, config)
-
-                def launch_skinny(
-                    latent: torch.Tensor,
-                    weight: torch.Tensor,
-                    shared_shard: torch.Tensor,
-                    *,
-                    skinny_kernel: SkinnyKernel = skinny_kernel,
-                    num_tokens: int = num_tokens,
-                ) -> None:
-                    skinny_kernel.launch(
-                        latent,
-                        weight,
-                        shared_shard[:num_tokens],
-                        mailbox,
-                    )
-
-                launches = make_up_projection_launches(
-                    launch_skinny,
-                    latent,
-                    weights,
-                    shared_shard,
-                )
-                graph = capture_up_projection_graph(launches)
-                timing = benchmark_up_projection_graph(
-                    graph,
-                    operations_per_replay=len(launches),
-                    warmup_replays=args.warmup_replays,
-                    samples=args.samples,
-                )
-                check_up_projection_output(
-                    mailbox[0, :num_tokens, :shard_size],
-                    latent,
-                    weights[-1],
-                    shared_shard[:num_tokens],
-                )
-                skinny_results.append(
-                    {
-                        "config": asdict(config),
-                        **timing,
-                    }
-                )
-            result["skinny"] = skinny_results
-        results.append(result)
-
-    properties = torch.cuda.get_device_properties(device)
-    report = {
-        "scope": "up-projection",
-        "device": properties.name,
-        "compute_capability": list(torch.cuda.get_device_capability(device)),
-        "tp_size": args.tp_size,
-        "shard_size": shard_size,
-        "weight_count": weight_count,
-        "cache_multiplier": args.cache_multiplier,
-        "warmup_replays": args.warmup_replays,
-        "samples": args.samples,
-        "results": results,
-    }
-    rendered = json.dumps(report, indent=2)
-    print(rendered, flush=True)
-    if args.output is not None:
-        args.output.parent.mkdir(parents=True, exist_ok=True)
-        args.output.write_text(rendered + "\n", encoding="utf-8")
-
-
-def capture_tail_graph(
-    operation: Callable[[], torch.Tensor],
-    cpu_group: dist.ProcessGroup,
-) -> tuple[torch.cuda.CUDAGraph, torch.Tensor]:
-    for _ in range(3):
-        dist.barrier(group=cpu_group)
-        output = operation()
-    torch.accelerator.synchronize()
-
-    dist.barrier(group=cpu_group)
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph):
-        output = operation()
-    torch.accelerator.synchronize()
-    return graph, output
-
-
-def benchmark_tail_graph(
-    graph: torch.cuda.CUDAGraph,
-    *,
-    warmup_replays: int,
-    samples: int,
-    device_group: dist.ProcessGroup,
-    cpu_group: dist.ProcessGroup,
-) -> dict[str, Any]:
-    for _ in range(warmup_replays):
-        graph.replay()
-    torch.accelerator.synchronize()
-
-    dist.barrier(group=cpu_group)
-    starts = [torch.cuda.Event(enable_timing=True) for _ in range(samples + 1)]
-    ends = [torch.cuda.Event(enable_timing=True) for _ in range(samples + 1)]
-    for start, end in zip(starts, ends):
-        start.record()
-        graph.replay()
-        end.record()
-    torch.accelerator.synchronize()
-
-    samples_us = torch.tensor(
-        [start.elapsed_time(end) * 1000.0 for start, end in zip(starts, ends)],
-        dtype=torch.float64,
-        device=torch.accelerator.current_device_index(),
-    )
-    dist.all_reduce(samples_us, op=dist.ReduceOp.MAX, group=device_group)
-    return summarize(samples_us[1:].tolist())
-
-
-def make_inputs(
-    num_tokens: int,
-    rank: int,
-    device: torch.device,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    torch.manual_seed(20260726 + 100 * num_tokens + rank)
-    routed = torch.randn(
-        (num_tokens, LATENT_SIZE),
-        dtype=torch.bfloat16,
-        device=device,
-    ).mul_(0.01)
-    shared = torch.randn(
-        (num_tokens, HIDDEN_SIZE),
-        dtype=torch.bfloat16,
-        device=device,
-    )
-    return routed, shared
-
-
-def make_reference(
-    routed: torch.Tensor,
-    shared: torch.Tensor,
-    rms_weight: torch.Tensor,
-    up_weight: torch.Tensor,
-    device_group: dist.ProcessGroup,
-) -> Callable[[], torch.Tensor]:
-    routed_workspace = torch.empty_like(routed)
-    shared_workspace = torch.empty_like(shared)
-
-    def reference() -> torch.Tensor:
-        routed_workspace.copy_(routed)
-        dist.all_reduce(routed_workspace, group=device_group)
-        normalized = F.rms_norm(
-            routed_workspace,
-            (LATENT_SIZE,),
-            rms_weight,
-            RMS_EPS,
-        )
-        projected = F.linear(normalized, up_weight)
-        shared_workspace.copy_(shared)
-        dist.all_reduce(shared_workspace, group=device_group)
-        return projected.add(shared_workspace)
-
-    return reference
-
-
-def check_fused_output(
-    fused_output: torch.Tensor,
-    reference: Callable[[], torch.Tensor],
-    cpu_group: dist.ProcessGroup,
-) -> None:
-    dist.barrier(group=cpu_group)
-    expected = reference()
-    torch.testing.assert_close(fused_output, expected, atol=8e-2, rtol=3e-2)
-
-
-def benchmark_whole_tail(args: argparse.Namespace) -> None:
-    if any(not 1 <= num_tokens <= 16 for num_tokens in args.num_tokens):
-        raise ValueError("--num-tokens values must be in [1, 16]")
-    if args.warmup_replays < 0 or args.samples <= 0:
-        raise ValueError("warmup replays must be nonnegative and samples positive")
-    if args.skinny_max_num_tokens is not None and any(
-        not 0 <= cutoff <= 8 for cutoff in args.skinny_max_num_tokens
-    ):
-        raise ValueError("--skinny-max-num-tokens must be in [0, 8]")
-    skinny_configs = dict(args.skinny_config or ())
-    if len(skinny_configs) != len(args.skinny_config or ()):
-        raise ValueError("--skinny-config must not repeat an M value")
-    if any(not 1 <= num_tokens <= 8 for num_tokens in skinny_configs):
-        raise ValueError("--skinny-config M values must be in [1, 8]")
-    if not {"RANK", "WORLD_SIZE", "LOCAL_RANK"} <= os.environ.keys():
-        raise RuntimeError("launch this benchmark with torchrun")
-
-    rank = int(os.environ["RANK"])
-    world_size = int(os.environ["WORLD_SIZE"])
-    local_rank = int(os.environ["LOCAL_RANK"])
-    device = torch.device("cuda", local_rank)
-    torch.accelerator.set_device_index(device)
-    init_distributed_environment()
-    if world_size > 8:
-        set_custom_all_reduce(False)
-    initialize_model_parallel(tensor_model_parallel_size=world_size)
-    device_group = get_tp_group().device_group
-    cpu_group = dist.new_group(backend="gloo")
-
-    if torch.cuda.get_device_capability(device)[0] != 10:
-        raise RuntimeError("Kimi K3 latent-MoE tail requires SM100")
-
-    torch.manual_seed(20260726)
-    rms_weight = 1 + 0.1 * torch.randn(
-        LATENT_SIZE,
-        dtype=torch.bfloat16,
-        device=device,
-    )
-    up_weight = (
-        torch.randn(
-            (HIDDEN_SIZE, LATENT_SIZE),
-            dtype=torch.bfloat16,
-            device=device,
-        )
-        / LATENT_SIZE**0.5
-    )
-
-    use_reference = args.backend in ("reference", "both")
-    use_fused = args.backend in ("fused", "both")
-    fused_ops = []
-    if use_fused:
-        production_config_for_m = fused_add_multicast_skinny_gemm.config_for_m
-
-        def config_for_m(
-            num_rows: int,
-            shard_dim: int = 896,
-        ) -> fused_add_multicast_skinny_gemm.SkinnyConfig:
-            config = skinny_configs.get(num_rows)
-            if config is not None:
-                return config
-            return production_config_for_m(num_rows, shard_dim)
-
-        fused_add_multicast_skinny_gemm.config_for_m = config_for_m
-        cutoffs = args.skinny_max_num_tokens or [latent_moe_tail._SKINNY_MAX_NUM_TOKENS]
-        for cutoff in cutoffs:
-            latent_moe_tail._SKINNY_MAX_NUM_TOKENS = cutoff
-            latent_moe_tail.KimiK3LatentMoETailOp._instances.clear()
-            fused_ops.append(
-                (
-                    cutoff,
-                    latent_moe_tail.KimiK3LatentMoETailOp.initialize(
-                        hidden_size=HIDDEN_SIZE,
-                        latent_size=LATENT_SIZE,
-                        dtype=torch.bfloat16,
-                        device=device,
-                        rms_eps=RMS_EPS,
-                    ),
-                )
-            )
-        cutedsl_warmup()
-
-    results = []
-    for num_tokens in args.num_tokens:
-        routed, shared = make_inputs(num_tokens, rank, device)
-        reference = make_reference(
-            routed,
-            shared,
-            rms_weight,
-            up_weight,
-            device_group,
-        )
-        result: dict[str, Any] = {"num_tokens": num_tokens}
-        if use_reference:
-            reference_graph, _ = capture_tail_graph(reference, cpu_group)
-            result["reference"] = benchmark_tail_graph(
-                reference_graph,
-                warmup_replays=args.warmup_replays,
-                samples=args.samples,
-                device_group=device_group,
-                cpu_group=cpu_group,
-            )
-        for cutoff, fused_op in fused_ops:
-
-            def fused(
-                routed: torch.Tensor = routed,
-                shared: torch.Tensor = shared,
-                fused_op: latent_moe_tail.KimiK3LatentMoETailOp = fused_op,
-            ) -> torch.Tensor:
-                return fused_op(routed, shared, rms_weight, up_weight)
-
-            fused_graph, fused_output = capture_tail_graph(fused, cpu_group)
-            fused_key = "fused" if len(fused_ops) == 1 else f"fused_skinny_max_{cutoff}"
-            result[fused_key] = benchmark_tail_graph(
-                fused_graph,
-                warmup_replays=args.warmup_replays,
-                samples=args.samples,
-                device_group=device_group,
-                cpu_group=cpu_group,
-            )
-            check_fused_output(fused_output, reference, cpu_group)
-            if "reference" in result:
-                speedup = (
-                    result["reference"]["median_us"] / result[fused_key]["median_us"]
-                )
-                if len(fused_ops) == 1:
-                    result["speedup"] = speedup
-                else:
-                    result[f"{fused_key}_speedup"] = speedup
-        results.append(result)
-
-    properties = torch.cuda.get_device_properties(device)
-    report = {
-        "scope": "whole-tail",
-        "device": properties.name,
-        "compute_capability": list(torch.cuda.get_device_capability(device)),
-        "world_size": world_size,
-        "torch_version": torch.__version__,
-        "cuda_version": torch.version.cuda,
-        "warmup_replays": args.warmup_replays,
-        "samples": args.samples,
-        "skinny_max_num_tokens": [cutoff for cutoff, _ in fused_ops],
-        "skinny_configs": {
-            str(num_tokens): asdict(config)
-            for num_tokens, config in skinny_configs.items()
-        },
-        "timing_scope": {
-            "reference": (
-                "two input copies, two AllReduces, RMSNorm, full replicated "
-                "up-projection GEMM, and final add"
-            ),
-            "fused": (
-                "routed AllReduce/RMSNorm plus shared ReduceScatter, sharded "
-                "up-projection/multicast, and Lamport copy"
-            ),
-        },
-        "results": results,
-    }
-    if rank == 0:
-        rendered = json.dumps(report, indent=2)
-        print(rendered, flush=True)
-        if args.output is not None:
-            args.output.parent.mkdir(parents=True, exist_ok=True)
-            args.output.write_text(rendered + "\n", encoding="utf-8")
-
-    dist.barrier(group=cpu_group)
-
-
-def main() -> None:
-    args = parse_args()
-    if args.scope == "up-projection":
-        benchmark_up_projection(args)
-        return
-
-    from vllm.config import VllmConfig, set_current_vllm_config
-
-    with set_current_vllm_config(VllmConfig()):
-        benchmark_whole_tail(args)
-
-
-if __name__ == "__main__":
-    main()
@@ -1,239 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import argparse
-import json
-import os
-import statistics
-from collections.abc import Callable
-
-import torch
-import torch.distributed as dist
-
-import vllm._custom_ops as ops
-from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce
-
-
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--tokens", type=int, nargs="+", default=[8, 32, 128, 1024])
-    parser.add_argument("--hidden-size", type=int, default=7168)
-    parser.add_argument("--graph-repeats", type=int, default=20)
-    parser.add_argument("--warmup-replays", type=int, default=5)
-    parser.add_argument("--samples", type=int, default=15)
-    return parser.parse_args()
-
-
-def capture_graph(op: Callable[[], None], repeats: int) -> torch.cuda.CUDAGraph:
-    stream = torch.cuda.Stream()
-    stream.wait_stream(torch.cuda.current_stream())
-    with torch.cuda.stream(stream):
-        for _ in range(3):
-            op()
-    stream.synchronize()
-
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph, stream=stream):
-        for _ in range(repeats):
-            op()
-    torch.cuda.current_stream().wait_stream(stream)
-    return graph
-
-
-def max_rank_graph_time(
-    graph: torch.cuda.CUDAGraph,
-    repeats: int,
-    warmup_replays: int,
-    samples: int,
-    device_group: dist.ProcessGroup,
-    cpu_group: dist.ProcessGroup,
-) -> float:
-    for _ in range(warmup_replays):
-        graph.replay()
-    torch.accelerator.synchronize()
-
-    timings = []
-    start = torch.cuda.Event(enable_timing=True)
-    end = torch.cuda.Event(enable_timing=True)
-    for _ in range(samples):
-        dist.barrier(group=cpu_group)
-        start.record()
-        graph.replay()
-        end.record()
-        end.synchronize()
-        elapsed = torch.tensor(
-            start.elapsed_time(end) / repeats,
-            dtype=torch.float64,
-            device=torch.accelerator.current_device_index(),
-        )
-        dist.all_reduce(elapsed, op=dist.ReduceOp.MAX, group=device_group)
-        timings.append(elapsed.item())
-    return statistics.median(timings)
-
-
-def check_outputs(
-    comm: CustomAllreduce,
-    local: torch.Tensor,
-    reduce_input: torch.Tensor,
-    device_group: dist.ProcessGroup,
-) -> None:
-    expected_gather = torch.empty(
-        (local.shape[0] * dist.get_world_size(), local.shape[1]),
-        dtype=local.dtype,
-        device=local.device,
-    )
-    dist.all_gather_into_tensor(expected_gather, local, group=device_group)
-    gathered = comm.custom_all_gather(local)
-    assert gathered is not None
-    torch.testing.assert_close(gathered, expected_gather)
-
-    expected_scatter = torch.empty_like(local)
-    dist.reduce_scatter_tensor(
-        expected_scatter,
-        reduce_input.clone(),
-        group=device_group,
-    )
-    scattered = comm.custom_reduce_scatter(reduce_input)
-    assert scattered is not None
-    torch.testing.assert_close(scattered, expected_scatter)
-
-
-def benchmark_shape(
-    comm: CustomAllreduce,
-    global_tokens: int,
-    hidden_size: int,
-    graph_repeats: int,
-    warmup_replays: int,
-    samples: int,
-    device_group: dist.ProcessGroup,
-    cpu_group: dist.ProcessGroup,
-) -> dict[str, float | int]:
-    world_size = dist.get_world_size()
-    rank = dist.get_rank()
-    padded_tokens = (global_tokens + world_size - 1) // world_size * world_size
-    local_tokens = padded_tokens // world_size
-    local = torch.full(
-        (local_tokens, hidden_size),
-        rank + 1,
-        dtype=torch.bfloat16,
-        device=torch.accelerator.current_device_index(),
-    )
-    reduce_input = torch.full(
-        (padded_tokens, hidden_size),
-        rank + 1,
-        dtype=torch.bfloat16,
-        device=local.device,
-    )
-    check_outputs(comm, local, reduce_input, device_group)
-
-    custom_gather_out = torch.empty(
-        (padded_tokens, hidden_size),
-        dtype=local.dtype,
-        device=local.device,
-    )
-    custom_scatter_out = torch.empty_like(local)
-    nccl_gather_out = torch.empty_like(custom_gather_out)
-    nccl_scatter_out = torch.empty_like(local)
-
-    def custom_ag() -> None:
-        ops.mnnvl_lamport_all_gather(
-            comm._ptr,
-            local,
-            custom_gather_out,
-            comm.mnnvl_lamport_ag_local_ptr,
-            comm.mnnvl_lamport_ag_multicast_ptr,
-            comm.mnnvl_lamport_ag_epoch_ptr,
-            comm.mnnvl_buffer_size,
-        )
-
-    def custom_rs() -> None:
-        ops.mnnvl_lamport_reduce_scatter(
-            comm._ptr,
-            reduce_input,
-            custom_scatter_out,
-            comm.mnnvl_lamport_rs_local_ptr,
-            comm.mnnvl_lamport_rs_epoch_ptr,
-            comm.mnnvl_buffer_size,
-        )
-
-    def nccl_ag() -> None:
-        dist.all_gather_into_tensor(nccl_gather_out, local, group=device_group)
-
-    def nccl_rs() -> None:
-        dist.reduce_scatter_tensor(
-            nccl_scatter_out,
-            reduce_input,
-            group=device_group,
-        )
-
-    graphs = {
-        "custom_ag_us": capture_graph(custom_ag, graph_repeats),
-        "nccl_ag_us": capture_graph(nccl_ag, graph_repeats),
-        "custom_rs_us": capture_graph(custom_rs, graph_repeats),
-        "nccl_rs_us": capture_graph(nccl_rs, graph_repeats),
-    }
-    times = {
-        name: max_rank_graph_time(
-            graph,
-            graph_repeats,
-            warmup_replays,
-            samples,
-            device_group,
-            cpu_group,
-        )
-        * 1000
-        for name, graph in graphs.items()
-    }
-    torch.testing.assert_close(custom_gather_out, nccl_gather_out)
-    torch.testing.assert_close(custom_scatter_out, nccl_scatter_out)
-    return {
-        "global_tokens": global_tokens,
-        "padded_tokens": padded_tokens,
-        "local_bytes": local.nbytes,
-        "full_bytes": reduce_input.nbytes,
-        **times,
-        "ag_speedup": times["nccl_ag_us"] / times["custom_ag_us"],
-        "rs_speedup": times["nccl_rs_us"] / times["custom_rs_us"],
-    }
-
-
-def main() -> None:
-    args = parse_args()
-    local_rank = int(os.environ["LOCAL_RANK"])
-    torch.accelerator.set_device_index(local_rank)
-    dist.init_process_group("nccl")
-    device_group = dist.group.WORLD
-    cpu_group = dist.new_group(backend="gloo")
-
-    comm = CustomAllreduce(
-        group=cpu_group,
-        device=torch.device("cuda", local_rank),
-    )
-    assert not comm.disabled
-    assert comm.world_size == 16
-    assert comm.mnnvl_only
-    assert comm.mnnvl_multicast_ptr
-
-    results = [
-        benchmark_shape(
-            comm,
-            tokens,
-            args.hidden_size,
-            args.graph_repeats,
-            args.warmup_replays,
-            args.samples,
-            device_group,
-            cpu_group,
-        )
-        for tokens in args.tokens
-    ]
-    if dist.get_rank() == 0:
-        print(json.dumps(results, indent=2), flush=True)
-
-    comm.close()
-    dist.destroy_process_group(cpu_group)
-    dist.destroy_process_group()
-
-
-if __name__ == "__main__":
-    main()
@@ -1,201 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Benchmark the RDNAHybridW4A16LinearKernel across decode and prefill shapes.
-
-Usage:
-    python benchmark_int4_gemm.py
-    python benchmark_int4_gemm.py --models Qwen/Qwen3-4B
-    python benchmark_int4_gemm.py --group-size 128
-"""
-
-import argparse
-import copy
-import itertools
-import os
-
-import torch
-
-from vllm.triton_utils import triton
-
-# ---------------------------------------------------------------------------
-# Weight shapes: [K, N], TP_SPLIT_DIM
-# ---------------------------------------------------------------------------
-WEIGHT_SHAPES = {
-    "Qwen/Qwen3-4B": [
-        ([2560, 3840], 1),  # qkv_proj
-        ([2560, 2560], 0),  # o_proj
-        ([2560, 19456], 1),  # gate_up_proj
-        ([9728, 2560], 0),  # down_proj
-    ],
-    "Qwen/Qwen2.5-7B-Instruct": [
-        ([3584, 4608], 1),
-        ([3584, 3584], 0),
-        ([3584, 37888], 1),
-        ([18944, 3584], 0),
-    ],
-    "trymirai/SmolLM2-1.7B-Instruct-AWQ": [
-        ([2048, 6144], 1),  # qkv_proj
-        ([2048, 2048], 0),  # o_proj
-        ([2048, 16384], 1),  # gate_up_proj
-        ([8192, 2048], 0),  # down_proj
-    ],
-    "RedHatAI/Qwen3-8B-quantized.w4a16": [
-        ([4096, 6144], 1),  # qkv_proj
-        ([4096, 4096], 0),  # o_proj
-        ([4096, 24576], 1),  # gate_up_proj
-        ([12288, 4096], 0),  # down_proj
-    ],
-}
-
-
-# ---------------------------------------------------------------------------
-# Weight packing
-# ---------------------------------------------------------------------------
-def prepare_hybrid_weights(K, N, group_size, device="cuda"):
-    """Create random weights for benchmarking.
-
-    Returns (w_q_skinny, w_s_skinny, w_fp16, w_zp). The triton path derives
-    its int32 view from w_q_skinny, so no separate int32 buffer is returned.
-    """
-    num_groups = K // group_size
-
-    # Random packed weights — actual values don't matter for throughput
-    w_q_skinny_i32 = torch.randint(
-        0, 2**31, (N, K // 8), dtype=torch.int32, device=device
-    )
-    w_q_skinny = w_q_skinny_i32.view(torch.int8).contiguous()
-    w_s_skinny = torch.randn(N, num_groups, dtype=torch.float16, device=device) * 0.01
-
-    # Raw per-group zero-points for asymmetric benchmarks
-    w_zp = torch.randint(0, 16, (N, num_groups), dtype=torch.int32, device=device).to(
-        torch.float16
-    )
-
-    # FP16 baseline for F.linear
-    w_fp16 = torch.randn(N, K, dtype=torch.float16, device=device) * 0.01
-
-    return w_q_skinny, w_s_skinny, w_fp16, w_zp
-
-
-# ---------------------------------------------------------------------------
-# Benchmark
-# ---------------------------------------------------------------------------
-PROVIDERS = ["torch-fp16", "hybrid-w4a16", "hybrid-w4a16-zp"]
-
-
-@triton.testing.perf_report(
-    triton.testing.Benchmark(
-        x_names=["batch_size"],
-        x_vals=[1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096],
-        x_log=False,
-        line_arg="provider",
-        line_vals=PROVIDERS,
-        line_names=PROVIDERS,
-        ylabel="TFLOP/s (larger is better)",
-        plot_name="FP16 vs Hybrid W4A16",
-        args={},
-    )
-)
-def benchmark(batch_size, provider, N, K, group_size, weights):
-    M = batch_size
-    device = "cuda"
-    dtype = torch.float16
-    a = torch.randn((M, K), device=device, dtype=dtype)
-
-    quantiles = [0.5, 0.2, 0.8]
-
-    if provider == "torch-fp16":
-        w_fp16 = weights["w_fp16"]
-        ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
-            lambda: torch.nn.functional.linear(a, w_fp16),
-            quantiles=quantiles,
-        )
-    elif provider in ("hybrid-w4a16", "hybrid-w4a16-zp"):
-        from vllm.model_executor.kernels.linear.mixed_precision import (
-            rdna_hybrid_w4a16 as _k,
-        )
-
-        _rdna_hybrid_w4a16_apply_impl = _k._rdna_hybrid_w4a16_apply_impl
-        from vllm.utils.platform_utils import num_compute_units
-
-        w = weights
-        cu_count = num_compute_units()
-        use_zp = provider == "hybrid-w4a16-zp"
-
-        def run():
-            return _rdna_hybrid_w4a16_apply_impl(
-                a,
-                w["w_q_skinny"],
-                w["w_s_skinny"],
-                w["w_zp"] if use_zp else None,
-                None,  # bias
-                cu_count,
-                group_size,
-            )
-
-        ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
-            run,
-            quantiles=quantiles,
-        )
-    else:
-        return 0.0, 0.0, 0.0
-
-    to_tflops = lambda t_ms: (2 * M * N * K) * 1e-12 / (t_ms * 1e-3)
-    return to_tflops(ms), to_tflops(max_ms), to_tflops(min_ms)
-
-
-def prepare_shapes(args):
-    KN_model_names = []
-    for model, tp_size in itertools.product(args.models, args.tp_sizes):
-        for KN, tp_dim in copy.deepcopy(WEIGHT_SHAPES[model]):
-            KN[tp_dim] //= tp_size
-            KN.append(model)
-            KN_model_names.append(KN)
-    return KN_model_names
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(
-        description="Benchmark RDNAHybridW4A16LinearKernel"
-    )
-    parser.add_argument(
-        "--models",
-        nargs="+",
-        type=str,
-        default=["Qwen/Qwen3-4B"],
-        choices=list(WEIGHT_SHAPES.keys()),
-    )
-    parser.add_argument("--tp-sizes", nargs="+", type=int, default=[1])
-    parser.add_argument("--group-size", type=int, default=128)
-    parser.add_argument("--save-path", type=str, default=None)
-    args = parser.parse_args()
-
-    for K, N, model in prepare_shapes(args):
-        group_size = args.group_size
-        print(f"\n{'=' * 70}")
-        print(f"{model}, N={N} K={K}, group_size={group_size}")
-        print(f"{'=' * 70}")
-
-        w_q_skinny, w_s_skinny, w_fp16, w_zp = prepare_hybrid_weights(K, N, group_size)
-
-        weights = {
-            "w_q_skinny": w_q_skinny,
-            "w_s_skinny": w_s_skinny,
-            "w_fp16": w_fp16,
-            "w_zp": w_zp,
-        }
-
-        save_path = args.save_path or f"bench_int4_res_n{N}_k{K}"
-        os.makedirs(save_path, exist_ok=True)
-        benchmark.run(
-            print_data=True,
-            show_plots=False,
-            save_path=save_path,
-            N=N,
-            K=K,
-            group_size=group_size,
-            weights=weights,
-        )
-
-    print("\nBenchmark finished!")
@@ -1,108 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-# Benchmark ReLUSquaredActivation: custom CUDA kernel vs forward_native, both
-# eager and under torch.compile (Inductor fuses relu+square into one kernel).
-
-import itertools
-
-import torch
-import torch.nn.functional as F
-
-import vllm.model_executor.layers.activation  # noqa: F401
-from vllm.benchmarks.lib.utils import default_vllm_config
-from vllm.triton_utils import triton
-from vllm.utils.argparse_utils import FlexibleArgumentParser
-from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, set_random_seed
-
-# Capped so the largest tensor stays under 2**31 elements: the shared activation
-# kernel computes the per-token pointer offset (blockIdx.x * d) in 32-bit, which
-# overflows for tensors with >2**32 elements. Realistic token counts are well
-# below this; the kernel-vs-native gap is already clear at these sizes.
-batch_size_range = [1, 16, 128]
-seq_len_range = [1, 16, 64, 1024]
-intermediate_size = [3072, 9728, 12288]
-configs = list(itertools.product(batch_size_range, seq_len_range, intermediate_size))
-
-
-@default_vllm_config()
-def benchmark_relu_squared(
-    batch_size: int,
-    seq_len: int,
-    intermediate_size: int,
-    provider: str,
-    dtype: torch.dtype,
-):
-    device = "cuda"
-    num_tokens = batch_size * seq_len
-    set_random_seed(42)
-    torch.set_default_device(device)
-
-    x = torch.randn(num_tokens, intermediate_size, dtype=dtype, device=device)
-    out = torch.empty_like(x)
-
-    def native(x: torch.Tensor) -> torch.Tensor:
-        return torch.square(F.relu(x))
-
-    # Verify the custom kernel matches the native implementation before timing.
-    ref = native(x)
-    torch.ops._C.relu_squared(out, x)
-    torch.testing.assert_close(out, ref)
-
-    if provider == "custom":
-        # Custom CUDA kernel — single fused kernel.
-        fn = lambda: torch.ops._C.relu_squared(out, x)
-    elif provider == "native":
-        # forward_native, eager — relu and square as separate ops.
-        fn = lambda: native(x)
-    elif provider == "native_compiled":
-        # forward_native under torch.compile — Inductor fuses relu+square.
-        # This is the real production baseline (custom ops are off when
-        # Inductor is enabled), so it is the comparison reviewers care about.
-        compiled = torch.compile(native)
-        compiled(x)  # warm up / trigger compilation before timing
-        fn = lambda: compiled(x)
-
-    ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
-        fn, quantiles=[0.5, 0.2, 0.8]
-    )
-    return ms, max_ms, min_ms
-
-
-if __name__ == "__main__":
-    parser = FlexibleArgumentParser(
-        description="Benchmark ReLUSquaredActivation: custom kernel vs native."
-    )
-    parser.add_argument(
-        "--dtype",
-        type=str,
-        choices=["half", "bfloat16", "float"],
-        default="bfloat16",
-    )
-    args = parser.parse_args()
-
-    dtype = STR_DTYPE_TO_TORCH_DTYPE[args.dtype]
-
-    perf_report = triton.testing.perf_report(
-        triton.testing.Benchmark(
-            x_names=["batch_size", "seq_len", "intermediate_size"],
-            x_vals=configs,
-            line_arg="provider",
-            line_vals=["custom", "native_compiled", "native"],
-            line_names=[
-                "Custom Kernel",
-                "Native (torch.compile)",
-                "Native (eager)",
-            ],
-            styles=[("blue", "-"), ("green", "-"), ("red", "-")],
-            ylabel="ms",
-            plot_name="relu_squared-eager-performance",
-            args={},
-        )
-    )
-
-    perf_report(
-        lambda batch_size, seq_len, intermediate_size, provider: benchmark_relu_squared(
-            batch_size, seq_len, intermediate_size, provider, dtype
-        )
-    ).run(print_data=True)
@@ -1,267 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""End-to-end autoregressive decode benchmark: ReplaySSM vs the standard SSM kernel.
-
-Loads a hybrid Mamba2 model, replicates one prompt across the batch, and times a
-long greedy decode (CUDA graphs on) once with the standard kernel and once with
-ReplaySSM, then reports the per-step / throughput speedup. The two modes run in
-separate subprocesses so each gets a clean CUDA context.
-
-The FlashInfer FP4-MoE autotuner is disabled by default (it is unstable under
-CUDA-graph capture on the pre-release Blackwell FP4 path); pass
--no-disable-flashinfer-autotune for non-FP4 models.
-
-Examples:
-    python e2e_decode_speedup.py --model-id nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16
-    python e2e_decode_speedup.py --dtype auto --buffer-len 16 \
-        --model-id nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4   # B300 NVFP4
-"""
-
-import argparse
-import json
-import os
-import subprocess
-import sys
-import time
-
-DEFAULT_PROMPT = "My cat wrote all this CUDA code for a new language model and"
-
-MODE_LABEL = {"standard": "standard", "replayssm": "ReplaySSM"}
-
-
-def parse_args():
-    p = argparse.ArgumentParser(
-        description="E2E decode speedup: ReplaySSM vs the standard SSM kernel."
-    )
-    p.add_argument("--model-id", default="nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16")
-    p.add_argument("--prompt", default=DEFAULT_PROMPT)
-    p.add_argument("--batch-size", type=int, default=256)
-    p.add_argument("--num-steps", type=int, default=1000)
-    p.add_argument("--warmup-steps", type=int, default=128)
-    p.add_argument("--repeats", type=int, default=1)
-    p.add_argument(
-        "--buffer-len", type=int, default=16, help="ReplaySSM input-buffer length."
-    )
-    p.add_argument(
-        "--dtype",
-        default="bfloat16",
-        choices=["bfloat16", "float16", "float32", "auto"],
-    )
-    p.add_argument("--gpu-memory-utilization", type=float, default=0.9)
-    p.add_argument("--max-model-len", type=int, default=None)
-    p.add_argument(
-        "--disable-flashinfer-autotune",
-        action=argparse.BooleanOptionalAction,
-        default=True,
-        help="Disable the FlashInfer FP4-MoE autotuner (default: on). "
-        "It is unstable under CUDA-graph capture on the "
-        "pre-release Blackwell FP4 path; pass "
-        "--no-disable-flashinfer-autotune for non-FP4 models.",
-    )
-    p.add_argument(
-        "--mamba-ssm-cache-dtype",
-        default="auto",
-        choices=["auto", "float32", "float16", "bfloat16"],
-        help="SSM state dtype (both modes). 'auto' = config-driven; "
-        "'float32' = fp32 state, 'bfloat16' = s16 state.",
-    )
-    p.add_argument(
-        "--baseline-ssm-config",
-        default="",
-        help="Pin the STANDARD baseline's SSM launch config as "
-        "'bsm,nw' via override_ssm_config (forces the in-process "
-        "engine so the override reaches the kernel). Empty = off.",
-    )
-    p.add_argument(
-        "--worker",
-        choices=["standard", "replayssm"],
-        default=None,
-        help=argparse.SUPPRESS,
-    )
-    return p.parse_args()
-
-
-def resolve_max_model_len(args) -> int:
-    if args.max_model_len is not None:
-        return args.max_model_len
-    return args.num_steps + 256
-
-
-def run_worker(args):
-    # override_ssm_config is a module global; it only reaches the model if the
-    # engine runs in-process (default V1 spawns a separate EngineCore). Force it.
-    if args.worker == "standard" and args.baseline_ssm_config:
-        os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
-
-    import torch
-
-    from vllm import LLM, SamplingParams
-
-    mode = args.worker
-    max_model_len = resolve_max_model_len(args)
-
-    llm_kwargs = dict(
-        model=args.model_id,
-        tensor_parallel_size=1,
-        dtype=args.dtype,
-        max_model_len=max_model_len,
-        trust_remote_code=True,
-        enable_prefix_caching=False,
-        enable_chunked_prefill=False,
-        max_num_seqs=args.batch_size,
-        max_num_batched_tokens=max(max_model_len, args.batch_size * 64),
-        enforce_eager=False,
-        disable_log_stats=True,
-        gpu_memory_utilization=args.gpu_memory_utilization,
-        # SSM state dtype (applies to both standard and ReplaySSM).
-        mamba_ssm_cache_dtype=args.mamba_ssm_cache_dtype,
-    )
-    if args.disable_flashinfer_autotune:
-        # FP4-MoE autotuner is unstable under CUDA-graph capture on Blackwell;
-        # re-enable (--no-disable-flashinfer-autotune) only for non-FP4 models.
-        llm_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False}
-    if mode == "replayssm":
-        llm_kwargs.update(use_replayssm=True, replayssm_buffer_len=args.buffer_len)
-
-    _ssm_cm = None
-    if mode == "standard" and args.baseline_ssm_config:
-        from vllm.model_executor.layers.mamba.ops.mamba_ssm import override_ssm_config
-
-        _bsm, _nw = (int(x) for x in args.baseline_ssm_config.split(","))
-        _ssm_cm = override_ssm_config((_bsm, _nw))
-        _ssm_cm.__enter__()  # active through LLM() graph capture + decode
-        print(
-            f"[{mode}] override_ssm_config -> (BLOCK_SIZE_M={_bsm}, num_warps={_nw})",
-            flush=True,
-        )
-
-    llm = LLM(**llm_kwargs)
-    prompts = [args.prompt] * args.batch_size
-
-    def timed_generate(n_tokens):
-        sp = SamplingParams(
-            n=1,
-            temperature=0.0,
-            ignore_eos=True,
-            min_tokens=n_tokens,
-            max_tokens=n_tokens,
-        )
-        if torch.accelerator.is_available():
-            torch.accelerator.synchronize()
-        t0 = time.perf_counter()
-        outs = llm.generate(prompts, sp, use_tqdm=False)
-        if torch.accelerator.is_available():
-            torch.accelerator.synchronize()
-        elapsed = time.perf_counter() - t0
-        produced = min(len(o.outputs[0].token_ids) for o in outs)
-        assert produced == n_tokens, f"expected {n_tokens} tokens, got {produced}"
-        return elapsed
-
-    timed_generate(args.warmup_steps)
-
-    best = None
-    for _ in range(args.repeats):
-        elapsed = timed_generate(args.num_steps)
-        tok_s = args.batch_size * args.num_steps / elapsed
-        per_step_ms = elapsed / args.num_steps * 1e3
-        print(
-            f"[{mode}] {elapsed:.3f}s  {tok_s:,.0f} tok/s  {per_step_ms:.3f} ms/step",
-            flush=True,
-        )
-        if best is None or elapsed < best["elapsed_s"]:
-            best = {
-                "mode": mode,
-                "elapsed_s": elapsed,
-                "tok_s": tok_s,
-                "per_step_ms": per_step_ms,
-            }
-
-    print("RESULT_JSON " + json.dumps(best), flush=True)
-    if _ssm_cm is not None:
-        _ssm_cm.__exit__(None, None, None)
-
-
-def run_one_mode(args, mode) -> dict:
-    cmd = [
-        sys.executable,
-        __file__,
-        "--worker",
-        mode,
-        "--model-id",
-        args.model_id,
-        "--prompt",
-        args.prompt,
-        "--batch-size",
-        str(args.batch_size),
-        "--num-steps",
-        str(args.num_steps),
-        "--warmup-steps",
-        str(args.warmup_steps),
-        "--repeats",
-        str(args.repeats),
-        "--buffer-len",
-        str(args.buffer_len),
-        "--dtype",
-        args.dtype,
-        "--gpu-memory-utilization",
-        str(args.gpu_memory_utilization),
-        "--mamba-ssm-cache-dtype",
-        args.mamba_ssm_cache_dtype,
-        "--baseline-ssm-config",
-        args.baseline_ssm_config,
-    ]
-    cmd.append(
-        "--disable-flashinfer-autotune"
-        if args.disable_flashinfer_autotune
-        else "--no-disable-flashinfer-autotune"
-    )
-    if args.max_model_len is not None:
-        cmd += ["--max-model-len", str(args.max_model_len)]
-
-    result = None
-    proc = subprocess.Popen(
-        cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1
-    )
-    for line in proc.stdout:
-        sys.stdout.write(line)
-        sys.stdout.flush()
-        if line.startswith("RESULT_JSON "):
-            result = json.loads(line[len("RESULT_JSON ") :])
-    proc.wait()
-    if proc.returncode != 0:
-        raise RuntimeError(f"mode '{mode}' worker exited with {proc.returncode}")
-    if result is None:
-        raise RuntimeError(f"mode '{mode}' produced no RESULT_JSON line")
-    return result
-
-
-def main():
-    args = parse_args()
-    if args.worker is not None:
-        run_worker(args)
-        return
-
-    print(
-        f"model={args.model_id}  batch_size={args.batch_size}  "
-        f"steps={args.num_steps}  buffer_len={args.buffer_len}  dtype={args.dtype}"
-    )
-
-    std = run_one_mode(args, "standard")
-    fla = run_one_mode(args, "replayssm")
-    speedup = std["per_step_ms"] / fla["per_step_ms"]
-
-    print()
-    header = f"{'mode':<10}{'ms/step':>12}{'tok/s':>16}{'wall (s)':>12}"
-    print(header)
-    print("-" * len(header))
-    for r in (std, fla):
-        print(
-            f"{MODE_LABEL[r['mode']]:<10}{r['per_step_ms']:>12.3f}"
-            f"{r['tok_s']:>16,.0f}{r['elapsed_s']:>12.3f}"
-        )
-    print("-" * len(header))
-    print(f"speedup (standard / ReplaySSM, per step): {speedup:.3f}x")
-
-
-if __name__ == "__main__":
-    main()
@@ -430,7 +430,6 @@ set(VLLM_EXT_SRC
    "csrc/cpu/layernorm.cpp"
    "csrc/cpu/mla_decode.cpp"
    "csrc/cpu/pos_encoding.cpp"
-    "csrc/cpu/mamba_cpu.cpp"
    "csrc/moe/dynamic_4bit_int_moe_cpu.cpp"
    "csrc/cpu/cpu_attn.cpp"
    "csrc/cpu/torch_bindings.cpp")
@@ -490,7 +489,6 @@ if (ENABLE_X86_ISA)
        "csrc/cpu/spec_decode_utils.cpp"
        "csrc/cpu/cpu_attn.cpp"
        "csrc/cpu/dnnl_kernels.cpp"
-        "csrc/cpu/mamba_cpu.cpp"
        "csrc/cpu/torch_bindings.cpp"
        # TODO: Remove these files
        "csrc/cpu/activation.cpp"
@@ -504,7 +502,6 @@ if (ENABLE_X86_ISA)
        "csrc/cpu/utils.cpp"
        "csrc/cpu/spec_decode_utils.cpp"
        "csrc/cpu/cpu_attn.cpp"
-        "csrc/cpu/mamba_cpu.cpp"
        "csrc/cpu/dnnl_kernels.cpp"
        "csrc/cpu/torch_bindings.cpp"
        # TODO: Remove these files
@@ -28,9 +28,9 @@ if(DEEPGEMM_SRC_DIR)
  message(STATUS "DeepGEMM using local DEEPGEMM_SRC_DIR: ${deepgemm_SOURCE_DIR}")
 else()
  # Keep in sync with tools/install_deepgemm.sh
-  set(_DEEPGEMM_UPSTREAM_REPO "git@github.com:Inferact/DeepGEMM.git")
+  set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/deepseek-ai/DeepGEMM.git")
  # NOTE: This is currently targeting nv-dev branch due to sm120 support
-  set(_DEEPGEMM_UPSTREAM_TAG "f5a76426fa084087169693fd0cd815223576d6e9")
+  set(_DEEPGEMM_UPSTREAM_TAG "a6b593d2826719dcf4892609af7b84ee23aaf32a")

  set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}")
  if(NOT _deepgemm_fc_root)
@@ -68,9 +68,6 @@ endif()
 if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9)
    list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0f")
-    if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.4)
-      list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.7f")
-    endif()
  else()
    list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0a")
  endif()
@@ -1,74 +0,0 @@
-include(FetchContent)
-
-if(DEFINED ENV{FLASH_KDA_SRC_DIR})
-  set(FLASH_KDA_SRC_DIR $ENV{FLASH_KDA_SRC_DIR})
-endif()
-
-if(FLASH_KDA_SRC_DIR)
-  FetchContent_Declare(
-    flashkda
-    SOURCE_DIR ${FLASH_KDA_SRC_DIR}
-  )
-else()
-  FetchContent_Declare(
-    flashkda
-    GIT_REPOSITORY git@github.com:Inferact/FlashKDA.git
-    GIT_TAG a3e42bbbece3bb38f7c426b880315294a336e82f
-    GIT_PROGRESS TRUE
-    GIT_SUBMODULES cutlass
-  )
-endif()
-
-FetchContent_MakeAvailable(flashkda)
-message(STATUS "FlashKDA is available at ${flashkda_SOURCE_DIR}")
-
-set(FLASH_KDA_SUPPORT_ARCHS)
-if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0)
-  list(APPEND FLASH_KDA_SUPPORT_ARCHS "9.0a")
-endif()
-if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-  list(APPEND FLASH_KDA_SUPPORT_ARCHS "10.0f" "12.0f")
-elseif(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9)
-  list(APPEND FLASH_KDA_SUPPORT_ARCHS "10.0a" "10.3a" "12.0a")
-endif()
-
-cuda_archs_loose_intersection(
-  FLASH_KDA_ARCHS "${FLASH_KDA_SUPPORT_ARCHS}" "${CUDA_ARCHS}")
-
-if(FLASH_KDA_ARCHS)
-  message(STATUS "FlashKDA CUDA architectures: ${FLASH_KDA_ARCHS}")
-
-  set(FLASH_KDA_SOURCES
-    csrc/flashkda_registration.cpp
-    ${flashkda_SOURCE_DIR}/csrc/flash_kda.cpp
-    ${flashkda_SOURCE_DIR}/csrc/smxx/fwd_launch.cu)
-  set(FLASH_KDA_INCLUDES
-    ${flashkda_SOURCE_DIR}/csrc
-    ${flashkda_SOURCE_DIR}/cutlass/include
-    ${flashkda_SOURCE_DIR}/cutlass/examples/common
-    ${flashkda_SOURCE_DIR}/cutlass/tools/util/include)
-
-  set_gencode_flags_for_srcs(
-    SRCS "${FLASH_KDA_SOURCES}"
-    CUDA_ARCHS "${FLASH_KDA_ARCHS}")
-
-  define_extension_target(
-    _flashkda_C
-    DESTINATION vllm
-    LANGUAGE ${VLLM_GPU_LANG}
-    SOURCES ${FLASH_KDA_SOURCES}
-    COMPILE_FLAGS ${VLLM_GPU_FLAGS}
-    ARCHITECTURES ${VLLM_GPU_ARCHES}
-    INCLUDE_DIRECTORIES ${FLASH_KDA_INCLUDES}
-    USE_SABI 3
-    WITH_SOABI)
-
-  target_compile_options(_flashkda_C PRIVATE
-    $<$<COMPILE_LANGUAGE:CUDA>:-UPy_LIMITED_API --expt-relaxed-constexpr --expt-extended-lambda --use_fast_math -O3>
-    $<$<COMPILE_LANGUAGE:CXX>:-UPy_LIMITED_API>)
-else()
-  message(STATUS
-    "FlashKDA will not compile: CUDA >=12.0 and a supported architecture "
-    "(SM90, SM10x, or SM12x) are required")
-  add_custom_target(_flashkda_C)
-endif()
@@ -19,7 +19,7 @@ else()
  FetchContent_Declare(
        flashmla
        GIT_REPOSITORY https://github.com/vllm-project/FlashMLA
-        GIT_TAG a8f794d1251cbfd88a5011445dd5582289c727e4
+        GIT_TAG a6ec2ba7bd0a7dff98b3f4d3e6b52b159c48d78b
        GIT_PROGRESS TRUE
        CONFIGURE_COMMAND ""
        BUILD_COMMAND ""
@@ -35,7 +35,7 @@ set(FLASHMLA_VENDOR_DIR "${CMAKE_SOURCE_DIR}/vllm/third_party/flashmla")
 file(MAKE_DIRECTORY "${FLASHMLA_VENDOR_DIR}")
 file(READ "${flashmla_SOURCE_DIR}/flash_mla/flash_mla_interface.py"
     FLASHMLA_INTERFACE_CONTENT)
-string(REPLACE "flash_mla_cuda = torch.ops._flashmla_C"
+string(REPLACE "import flash_mla.cuda as flash_mla_cuda"
               "import vllm._flashmla_C\nflash_mla_cuda = torch.ops._flashmla_C"
               FLASHMLA_INTERFACE_CONTENT
               "${FLASHMLA_INTERFACE_CONTENT}")
@@ -60,9 +60,6 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9)
    # CUDA 12.9 has introduced "Family-Specific Architecture Features"
    # this supports all compute_10x family
    list(APPEND SUPPORT_ARCHS "10.0f")
-    if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.4)
-        list(APPEND SUPPORT_ARCHS "10.7f")
-    endif()
 elseif(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
    list(APPEND SUPPORT_ARCHS "10.0a")
 endif()
@@ -75,7 +72,7 @@ if(FLASH_MLA_ARCHS)
    list(APPEND VLLM_FLASHMLA_GPU_FLAGS "--expt-relaxed-constexpr" "--expt-extended-lambda" "--use_fast_math")

    set(FlashMLA_SOURCES
-        ${flashmla_SOURCE_DIR}/csrc/api/api.cpp
+        ${flashmla_SOURCE_DIR}/csrc/torch_api.cpp

        # Misc kernels for decoding
        ${flashmla_SOURCE_DIR}/csrc/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu
@@ -131,7 +128,6 @@ if(FLASH_MLA_ARCHS)

    set(FlashMLA_Extension_INCLUDES
        ${flashmla_SOURCE_DIR}/csrc
-        ${flashmla_SOURCE_DIR}/csrc/kerutils/include
        ${flashmla_SOURCE_DIR}/csrc/extension/sm90/dense_fp8/
        ${flashmla_SOURCE_DIR}/csrc/cutlass/include
        ${flashmla_SOURCE_DIR}/csrc/cutlass/tools/util/include
@@ -156,18 +152,15 @@ if(FLASH_MLA_ARCHS)
        USE_SABI 3
        WITH_SOABI)

-    # Enable C++20 for the FlashMLA sources (required for std::span, requires, etc.)
+    # Keep Stable ABI for the module, but *not* for CUDA/C++ files.
+    # This prevents Py_LIMITED_API from affecting nvcc and C++ compiles.
+    # Also enable C++20 for the FlashMLA sources (required for std::span, requires, etc.)
    target_compile_options(_flashmla_C PRIVATE
+        $<$<COMPILE_LANGUAGE:CUDA>:-UPy_LIMITED_API>
+        $<$<COMPILE_LANGUAGE:CXX>:-UPy_LIMITED_API>
        $<$<COMPILE_LANGUAGE:CXX>:-std=c++20>
        $<$<COMPILE_LANGUAGE:CUDA>:-std=c++20>)

-    # _flashmla_C is now ABI-stable torch 2.11+
-    target_compile_definitions(_flashmla_C PRIVATE
-        TORCH_TARGET_VERSION=0x020B000000000000ULL)
-    if(VLLM_GPU_LANG STREQUAL "CUDA")
-        target_compile_definitions(_flashmla_C PRIVATE USE_CUDA)
-    endif()
-
    define_extension_target(
        _flashmla_extension_C
        DESTINATION vllm
@@ -179,15 +172,15 @@ if(FLASH_MLA_ARCHS)
        USE_SABI 3
        WITH_SOABI)

-    # _flashmla_extension_C is now ABI-stable w/ torch 2.11+
-    target_compile_definitions(_flashmla_extension_C PRIVATE
-        TORCH_TARGET_VERSION=0x020B000000000000ULL)
-    if(VLLM_GPU_LANG STREQUAL "CUDA")
-        target_compile_definitions(_flashmla_extension_C PRIVATE USE_CUDA)
-    endif()
+    # Keep Stable ABI for the module, but *not* for CUDA/C++ files.
+    # This prevents Py_LIMITED_API from affecting nvcc and C++ compiles.
+    target_compile_options(_flashmla_extension_C PRIVATE
+        $<$<COMPILE_LANGUAGE:CUDA>:-UPy_LIMITED_API>
+        $<$<COMPILE_LANGUAGE:CXX>:-UPy_LIMITED_API>)
 else()
    message(STATUS "FlashMLA will not compile: unsupported CUDA architecture ${CUDA_ARCHS}")
    # Create empty targets for setup.py on unsupported systems
    add_custom_target(_flashmla_C)
    add_custom_target(_flashmla_extension_C)
 endif()
+
@@ -17,7 +17,7 @@ else()
  FetchContent_Declare(
    fmha_sm100
    GIT_REPOSITORY https://github.com/vllm-project/MSA.git
-    GIT_TAG 890aaa1a37a598ad17ccff0827fea21540d381fa
+    GIT_TAG 2e63ec37a0fc29bc20f39cd1a52e0f5affc33a73
    GIT_PROGRESS TRUE
    CONFIGURE_COMMAND ""
    BUILD_COMMAND ""
@@ -22,7 +22,7 @@ if(QUTLASS_SRC_DIR)
  set(qutlass_BINARY_DIR "${CMAKE_BINARY_DIR}/qutlass-binary-dir-unused")
 else()
  set(_QUTLASS_UPSTREAM_REPO "https://github.com/IST-DASLab/qutlass.git")
-  set(_QUTLASS_UPSTREAM_TAG "e74319e3405ce6d71965732880f5dc1f52371f64")
+  set(_QUTLASS_UPSTREAM_TAG "830d2c4537c7396e14a02a46fbddd18b5d107c65")

  set(_qutlass_fc_root "${FETCHCONTENT_BASE_DIR}")
  if(NOT _qutlass_fc_root)
@@ -55,11 +55,7 @@ message(STATUS "[QUTLASS] QuTLASS is available at ${qutlass_SOURCE_DIR}")

 if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
  cuda_archs_loose_intersection(QUTLASS_SM120_ARCHS "12.0f" "${CUDA_ARCHS}")
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.4)
-    cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0f;10.7f" "${CUDA_ARCHS}")
-  else()
-    cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0f" "${CUDA_ARCHS}")
-  endif()
+  cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0f" "${CUDA_ARCHS}")
 else()
  cuda_archs_loose_intersection(QUTLASS_SM120_ARCHS "12.0a;12.1a" "${CUDA_ARCHS}")
  cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0a;10.3a" "${CUDA_ARCHS}")
@@ -129,6 +125,8 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
    CUDA_ARCHS "${QUTLASS_ARCHS}"
  )

+  # QuTLASS uses legacy ATen headers and cannot be built with TORCH_TARGET_VERSION.
+  # Keep it as its own extension (registers torch.ops._qutlass_C).
  define_extension_target(
    _qutlass_C
    DESTINATION vllm
@@ -141,11 +139,9 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
    WITH_SOABI)

  target_compile_definitions(_qutlass_C PRIVATE
-    QUTLASS_MINIMAL_BUILD=1
+    QUTLASS_DISABLE_PYBIND=1
    TARGET_CUDA_ARCH=${QUTLASS_TARGET_CC}
-    CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1
-    TORCH_TARGET_VERSION=0x020B000000000000ULL
-    USE_CUDA)
+    CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)

  set_property(SOURCE ${QUTLASS_SOURCES} APPEND PROPERTY COMPILE_OPTIONS
    $<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr --use_fast_math -O3>
@@ -1,50 +0,0 @@
-include(FetchContent)
-
-if(DEFINED ENV{TML_FA4_SRC_DIR})
-  set(TML_FA4_SRC_DIR $ENV{TML_FA4_SRC_DIR})
-endif()
-
-if(TML_FA4_SRC_DIR)
-  FetchContent_Declare(
-    tml_fa4
-    SOURCE_DIR ${TML_FA4_SRC_DIR}
-    CONFIGURE_COMMAND ""
-    BUILD_COMMAND "")
-else()
-  FetchContent_Declare(
-    tml_fa4
-    GIT_REPOSITORY https://github.com/vllm-project/tml-fa4.git
-    GIT_TAG b206834606ed5b5f21f8eed6b0683f528ea9cf7d
-    GIT_PROGRESS TRUE
-    CONFIGURE_COMMAND ""
-    BUILD_COMMAND "")
-endif()
-
-FetchContent_GetProperties(tml_fa4)
-if(NOT tml_fa4_POPULATED)
-  FetchContent_Populate(tml_fa4)
-endif()
-message(STATUS "tml-fa4 is available at ${tml_fa4_SOURCE_DIR}")
-
-add_custom_target(tml_fa4)
-
-# Install into a private namespace so this implementation cannot shadow the
-# flash_attn package used by vLLM's standard attention backends.
-install(CODE "
-  file(GLOB_RECURSE TML_FA4_PY_FILES
-    \"${tml_fa4_SOURCE_DIR}/flash_attn/cute/*.py\")
-  foreach(SRC_FILE \${TML_FA4_PY_FILES})
-    file(RELATIVE_PATH REL_PATH
-      \"${tml_fa4_SOURCE_DIR}/flash_attn/cute\" \${SRC_FILE})
-    set(DST_FILE
-      \"\${CMAKE_INSTALL_PREFIX}/vllm/third_party/tml_fa4/\${REL_PATH}\")
-    get_filename_component(DST_DIR \${DST_FILE} DIRECTORY)
-    file(MAKE_DIRECTORY \${DST_DIR})
-    file(READ \${SRC_FILE} FILE_CONTENTS)
-    string(REPLACE
-      \"flash_attn.cute\"
-      \"vllm.third_party.tml_fa4\"
-      FILE_CONTENTS \"\${FILE_CONTENTS}\")
-    file(WRITE \${DST_FILE} \"\${FILE_CONTENTS}\")
-  endforeach()
-" COMPONENT tml_fa4)
@@ -39,7 +39,7 @@ else()
  FetchContent_Declare(
          vllm-flash-attn
          GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
-          GIT_TAG ed4b7342bc8f0489dd9b649d5288867e35fc6a32
+          GIT_TAG b3964b1d8b95d8e8447435668ab169a2700bab65
          GIT_PROGRESS TRUE
          # Don't share the vllm-flash-attn build between build types
          BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
@@ -396,24 +396,14 @@ function(cuda_archs_loose_intersection OUT_CUDA_ARCHS SRC_CUDA_ARCHS TGT_CUDA_AR
  # match — e.g. SRC="12.0f" matches TGT="12.1a" since SM121 is in the SM12x
  # family. The output uses TGT's value to preserve the user's compilation flags.
  set(_CUDA_ARCHS)
-  # Resolve exact base matches before family fallbacks so a generic entry such
-  # as 10.0f cannot consume a 10.7 target that has a 10.7f source entry.
-  foreach(_arch ${_SRC_CUDA_ARCHS})
-    if(_arch MATCHES "[af]$")
-      string(REGEX REPLACE "[af]$" "" _base "${_arch}")
-      if("${_base}" IN_LIST _TGT_CUDA_ARCHS)
-        list(REMOVE_ITEM _SRC_CUDA_ARCHS "${_arch}")
-        list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_base}")
-        list(APPEND _CUDA_ARCHS "${_arch}")
-      endif()
-    endif()
-  endforeach()
-
  foreach(_arch ${_SRC_CUDA_ARCHS})
    if(_arch MATCHES "[af]$")
      list(REMOVE_ITEM _SRC_CUDA_ARCHS "${_arch}")
      string(REGEX REPLACE "[af]$" "" _base "${_arch}")
-      if("${_base}a" IN_LIST _TGT_CUDA_ARCHS)
+      if ("${_base}" IN_LIST TGT_CUDA_ARCHS)
+        list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_base}")
+        list(APPEND _CUDA_ARCHS "${_arch}")
+      elseif("${_base}a" IN_LIST _TGT_CUDA_ARCHS)
        list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_base}a")
        list(APPEND _CUDA_ARCHS "${_base}a")
      elseif("${_base}f" IN_LIST _TGT_CUDA_ARCHS)
@@ -497,7 +487,7 @@ endfunction()

 function(cuda_archs_sm90plus OUT_CUDA_ARCHS TGT_CUDA_ARCHS)
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(_archs "9.0a;10.0f;10.7f;11.0f;12.0f" "${TGT_CUDA_ARCHS}")
+    cuda_archs_loose_intersection(_archs "9.0a;10.0f;11.0f;12.0f" "${TGT_CUDA_ARCHS}")
  else()
    cuda_archs_loose_intersection(_archs "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${TGT_CUDA_ARCHS}")
  endif()
@@ -67,8 +67,9 @@ void cp_gather_and_upconvert_fp8_kv_cache(
    torch::Tensor const& src_cache,         // [NUM_BLOCKS, BLOCK_SIZE, 656]
    torch::Tensor const& dst,               // [TOT_TOKENS, 576]
    torch::Tensor const& block_table,       // [BATCH, BLOCK_INDICES]
+    torch::Tensor const& seq_lens,          // [BATCH]
    torch::Tensor const& workspace_starts,  // [BATCH]
-    int64_t batch_size, std::optional<torch::Tensor> seq_starts = std::nullopt);
+    int64_t batch_size);

 // Indexer K quantization and cache function
 void indexer_k_quant_and_cache(
@@ -102,9 +102,7 @@ class TileGemm82 {
    kv_cache_t* __restrict__ curr_b = b_tile;

    for (int32_t k = 0; k < dynamic_k_size; ++k) {
-      auto fp32_b_regs = load_b_pair_vec(curr_b);
-      auto fp32_b_0_reg = fp32_b_regs.first;
-      auto fp32_b_1_reg = fp32_b_regs.second;
+      auto [fp32_b_0_reg, fp32_b_1_reg] = load_b_pair_vec(curr_b);

      float* __restrict__ curr_m_a = curr_a;
      vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
@@ -336,14 +336,13 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
    reg.val[1] = fp16_to_fp32_bits(raw_lo);
  }
  float reduce_sum() const {
-    // VSX horizontal reduction: 3 vector ops instead of 8 scalar adds.
-    // Step 1: pairwise sum of the two 4-wide halves
-    __vector float s = vec_add(reg.val[0], reg.val[1]);
-    // Step 2: rotate by 8 bytes (2 floats) and add
-    s = vec_add(s, vec_sld(s, s, 8));
-    // Step 3: rotate by 4 bytes (1 float) and add  => all lanes hold total
-    s = vec_add(s, vec_sld(s, s, 4));
-    return vec_extract(s, 0);
+    AliasReg ar;
+    ar.reg = reg;
+    float result = 0;
+    unroll_loop<int, VEC_ELEM_NUM>(
+        [&result, &ar](int i) { result += ar.values[i]; });
+
+    return result;
  }
  FP32Vec8 exp() const {
    f32x4x2_t out;
@@ -1,285 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-//
-// CPU at::Tensor wrappers for Mamba decode-step kernels defined in
-// mamba_kernels.hpp.
-
-#include "cpu/mamba_kernels.hpp"
-
-#include <ATen/ATen.h>
-#include <torch/library.h>
-#include <c10/util/Optional.h>
-
-#include "cpu_types.hpp"
-
-// ---------------------------------------------------------------------------
-// causal_conv1d_update
-// ---------------------------------------------------------------------------
-at::Tensor causal_conv1d_update_cpu_impl(
-    at::Tensor& x, at::Tensor& conv_state, const at::Tensor& weight,
-    const c10::optional<at::Tensor>& bias,
-    const c10::optional<std::string>& activation,
-    const c10::optional<at::Tensor>& conv_state_indices,
-    const c10::optional<at::Tensor>& query_start_loc, int64_t pad_slot_id) {
-  bool do_silu = false;
-  if (activation.has_value()) {
-    const std::string& act = activation.value();
-    do_silu = (act == "silu" || act == "swish");
-  }
-
-  at::ScalarType dtype = x.scalar_type();
-
-  // Input x: contiguous in native dtype.
-  at::Tensor x_c = x.is_contiguous() ? x : x.contiguous();
-
-  // conv_state: NEVER copy the full paged tensor just for layout reasons.
-  // If the dtype matches we work directly on conv_state (contiguous or not)
-  // by extracting strides and passing them to the kernel.
-  // Only a dtype-conversion copy is made when types differ (rare for BF16).
-  bool state_type_ok = (conv_state.scalar_type() == dtype);
-  at::Tensor state_c = state_type_ok ? conv_state : conv_state.to(dtype);
-  // state_c and conv_state may be non-contiguous — that is intentional.
-
-  // Weight: coerce to same dtype if needed (should match in practice)
-  at::Tensor w_c =
-      (weight.scalar_type() != dtype)
-          ? weight.to(dtype).contiguous()
-          : (weight.is_contiguous() ? weight : weight.contiguous());
-
-  // Bias stays float32 (small scalar, used only for fp32 accumulation)
-  at::Tensor bias_f32;
-  if (bias.has_value() && bias.value().defined())
-    bias_f32 = bias.value().to(at::kFloat).contiguous();
-
-  int64_t batch = x_c.size(0);
-  int64_t dim = x_c.size(1);
-  int64_t seqlen = (x_c.dim() == 3) ? x_c.size(2) : 1;
-  int64_t width = w_c.size(1);
-  int64_t state_len = state_c.size(2);
-
-  // Extract strides — works for contiguous AND non-contiguous (transposed)
-  // state. stride(0): between cache slots (e.g. num_slots × dim × width-1 in
-  // contiguous) stride(1): between conv channels (dim stride) stride(2):
-  // between state elements (=1 when contiguous, =dim when transposed)
-  int64_t stride_s_slot = state_c.stride(0);
-  int64_t stride_s_dim = state_c.stride(1);
-  int64_t stride_s_state = state_c.stride(2);
-
-  at::Tensor out = x_c.clone();  // native dtype, no float32 alloc
-
-  const int32_t* cache_idx_ptr = nullptr;
-  at::Tensor cache_idx_int;
-  if (conv_state_indices.has_value()) {
-    cache_idx_int = conv_state_indices.value().to(at::kInt).contiguous();
-    cache_idx_ptr = cache_idx_int.data_ptr<int32_t>();
-  }
-
-  VLLM_DISPATCH_FLOATING_TYPES(dtype, "causal_conv1d_update", [&] {
-    mamba_cpu::causal_conv1d_update_kernel<scalar_t>(
-        x_c.data_ptr<scalar_t>(), state_c.data_ptr<scalar_t>(), stride_s_slot,
-        stride_s_dim, stride_s_state, w_c.data_ptr<scalar_t>(),
-        bias_f32.defined() ? bias_f32.data_ptr<float>() : nullptr,
-        out.data_ptr<scalar_t>(), cache_idx_ptr,
-        static_cast<int32_t>(pad_slot_id), batch, dim, seqlen, width, state_len,
-        do_silu);
-  });
-
-  // Write back only when a type-conversion copy was made.
-  // Layout-only non-contiguity is handled via strides above — no copy needed.
-  if (!state_type_ok) conv_state.copy_(state_c);
-
-  return out;
-}
-
-// ---------------------------------------------------------------------------
-// selective_state_update
-// ---------------------------------------------------------------------------
-void selective_state_update_cpu_impl(
-    at::Tensor& state,    // (nstates, nheads, dim, dstate)
-    const at::Tensor& x,  // (N, nheads, dim)
-    const at::Tensor& dt, const at::Tensor& A, const at::Tensor& B,
-    const at::Tensor& C, const c10::optional<at::Tensor>& D,
-    const c10::optional<at::Tensor>& z,
-    const c10::optional<at::Tensor>& dt_bias, bool dt_softplus,
-    const c10::optional<at::Tensor>& state_batch_indices,
-    const c10::optional<at::Tensor>& dst_state_batch_indices,
-    int64_t null_block_id, at::Tensor& out,
-    const c10::optional<at::Tensor>& num_accepted_tokens,
-    const c10::optional<at::Tensor>& cu_seqlens) {
-  at::ScalarType state_type = state.scalar_type();
-  at::ScalarType input_type = x.scalar_type();
-
-  // x, B, C must be contiguous and match input_type
-  auto ensure_input = [input_type](const at::Tensor& t) -> at::Tensor {
-    at::Tensor r = (t.scalar_type() != input_type) ? t.to(input_type) : t;
-    return r.is_contiguous() ? r : r.contiguous();
-  };
-  at::Tensor x_in = ensure_input(x);
-  at::Tensor B_in = ensure_input(B);
-  at::Tensor C_in = ensure_input(C);
-  at::Tensor z_in;
-  if (z.has_value() && z.value().defined()) z_in = ensure_input(z.value());
-
-  // A, D, dt_bias are float32 model parameters that arrive here as expanded
-  // tensors, e.g. A is (nheads, head_dim, dstate) with strides (1, 0, 0).
-  // We need just the scalar value per head as a (nheads,) 1-D array so that
-  // A_ptr[h] in the kernel correctly reads head h's value.
-  //
-  // Strategy: peel trailing expanded (stride=0) dims via .select(), which is
-  // a zero-copy view.  For A: (nheads, head_dim, dstate) strides (1,0,0)
-  //   → .select(2,0) → (nheads, head_dim) strides (1,0)
-  //   → .select(1,0) → (nheads,)          stride  (1,)   ← contiguous, free.
-  // No allocation, no type conversion (A is already float32).
-  auto to_per_head_1d_f32 = [](const at::Tensor& t) -> at::Tensor {
-    at::Tensor r = t;
-    // Peel trailing dimensions that are broadcast (stride=0 or size=1)
-    while (r.dim() > 1) r = r.select(r.dim() - 1, 0);
-    if (r.scalar_type() != at::kFloat) r = r.to(at::kFloat);
-    return r.is_contiguous() ? r : r.contiguous();
-  };
-
-  at::Tensor A_f32 = to_per_head_1d_f32(A);  // (nheads,) float32
-  at::Tensor D_f32, dt_bias_f32;
-  if (D.has_value() && D.value().defined())
-    D_f32 = to_per_head_1d_f32(D.value());
-  if (dt_bias.has_value() && dt_bias.value().defined())
-    dt_bias_f32 = to_per_head_1d_f32(dt_bias.value());
-
-  // dt: reduce (N, nheads, head_dim) expanded tensor → (N, nheads) BEFORE
-  // the type conversion so we convert head_dim x fewer elements.
-  at::Tensor dt_f32;
-  {
-    // If dt was expanded to (N, nheads, head_dim) with stride-0 in dim 2,
-    // take a zero-copy view of index 0 along that dim first.
-    at::Tensor t2 = (dt.dim() == 3) ? dt.select(2, 0) : dt;  // (N, nheads)
-    at::Tensor t3 = (t2.scalar_type() != at::kFloat) ? t2.to(at::kFloat) : t2;
-    dt_f32 = t3.is_contiguous() ? t3 : t3.contiguous();
-  }
-
-  int64_t nheads = state.size(1);
-  int64_t dim = state.size(2);
-  int64_t dstate = state.size(3);
-  int64_t N = (cu_seqlens.has_value() && cu_seqlens.value().defined())
-                  ? cu_seqlens.value().size(0) - 1
-                  : x_in.size(0);
-  int64_t ngroups = B_in.size(1);
-
-  // Strides
-  int64_t stride_state_n = state.stride(0);
-  int64_t stride_state_h = state.stride(1);
-  int64_t stride_state_d = state.stride(2);
-  int64_t stride_x_n = x_in.stride(0);
-  int64_t stride_x_h = x_in.stride(1);
-  int64_t stride_dt_n = dt_f32.stride(0);  // dt is (N, nheads)
-  int64_t stride_BC_n = B_in.stride(0);
-  int64_t stride_BC_g = B_in.stride(1);
-  int64_t stride_out_n = out.stride(0);
-  int64_t stride_out_h = out.stride(1);
-
-  // Optional index pointers
-  auto get_int32_ptr =
-      [](const c10::optional<at::Tensor>& opt) -> const int32_t* {
-    return (opt.has_value() && opt.value().defined())
-               ? opt.value().data_ptr<int32_t>()
-               : nullptr;
-  };
-  const int32_t* sbi_ptr = get_int32_ptr(state_batch_indices);
-  const int32_t* dsbi_ptr = get_int32_ptr(dst_state_batch_indices);
-  const int32_t* nat_ptr = get_int32_ptr(num_accepted_tokens);
-  const int32_t* csl_ptr = get_int32_ptr(cu_seqlens);
-
-  // Dispatch on (state_t, input_t, out_t): write directly into `out`
-  // without any intermediate float32 buffer.
-  VLLM_DISPATCH_FLOATING_TYPES(state_type, "ssu_state", [&] {
-    using state_t = scalar_t;
-    VLLM_DISPATCH_FLOATING_TYPES(input_type, "ssu_input", [&] {
-      using input_t = scalar_t;
-      VLLM_DISPATCH_FLOATING_TYPES(out.scalar_type(), "ssu_out", [&] {
-        using out_t = scalar_t;
-        mamba_cpu::selective_state_update_kernel<state_t, input_t, out_t>(
-            state.data_ptr<state_t>(), stride_state_n, stride_state_h,
-            stride_state_d, x_in.data_ptr<input_t>(), stride_x_n, stride_x_h,
-            dt_f32.data_ptr<float>(), stride_dt_n, A_f32.data_ptr<float>(),
-            B_in.data_ptr<input_t>(), C_in.data_ptr<input_t>(), stride_BC_n,
-            stride_BC_g, D_f32.defined() ? D_f32.data_ptr<float>() : nullptr,
-            z_in.defined() ? z_in.data_ptr<input_t>() : nullptr,
-            dt_bias_f32.defined() ? dt_bias_f32.data_ptr<float>() : nullptr,
-            out.data_ptr<out_t>(), stride_out_n, stride_out_h, sbi_ptr,
-            dsbi_ptr, static_cast<int32_t>(null_block_id), nat_ptr, csl_ptr, N,
-            nheads, ngroups, dim, dstate, dt_softplus);
-      });
-    });
-  });
-}
-
-// ---------------------------------------------------------------------------
-// mamba_chunk_scan_fwd_cpu
-// ---------------------------------------------------------------------------
-void mamba_chunk_scan_fwd_cpu_impl(
-    at::Tensor& out,  // [seqlen, nheads, headdim] — pre-allocated by caller
-    at::Tensor&
-        final_states,     // [batch, nheads, headdim, dstate] float32 contiguous
-    const at::Tensor& x,  // [seqlen, nheads, headdim]
-    const at::Tensor&
-        dt,  // [seqlen, nheads] float32 (preprocessed: bias+softplus+clamp)
-    const at::Tensor& A,                 // [nheads] float32
-    const at::Tensor& B,                 // [seqlen, ngroups, dstate]
-    const at::Tensor& C,                 // [seqlen, ngroups, dstate]
-    const c10::optional<at::Tensor>& D,  // [nheads] float32 (optional)
-    const c10::optional<at::Tensor>& z,  // [seqlen, nheads, headdim] (optional)
-    const at::Tensor& cu_seqlens         // [batch+1] int32
-) {
-  const at::ScalarType input_type = x.scalar_type();
-
-  auto ensure_contig = [input_type](const at::Tensor& t) -> at::Tensor {
-    at::Tensor r = (t.scalar_type() != input_type) ? t.to(input_type) : t;
-    return r.is_contiguous() ? r : r.contiguous();
-  };
-  at::Tensor x_in = ensure_contig(x);
-  at::Tensor B_in = ensure_contig(B);
-  at::Tensor C_in = ensure_contig(C);
-  at::Tensor z_in;
-  if (z.has_value() && z.value().defined()) z_in = ensure_contig(z.value());
-
-  // A and D are float32 model parameters, potentially broadcast-expanded.
-  // Strip trailing broadcast dims to get a contiguous (nheads,) array.
-  auto to_per_head_f32 = [](const at::Tensor& t) -> at::Tensor {
-    at::Tensor r = t;
-    while (r.dim() > 1) r = r.select(r.dim() - 1, 0);
-    if (r.scalar_type() != at::kFloat) r = r.to(at::kFloat);
-    return r.is_contiguous() ? r : r.contiguous();
-  };
-  at::Tensor A_f32 = to_per_head_f32(A);
-  at::Tensor D_f32;
-  if (D.has_value() && D.value().defined()) D_f32 = to_per_head_f32(D.value());
-
-  // dt: [seqlen, nheads] float32 — caller has applied bias+softplus+clamp in
-  // Python.
-  at::Tensor dt_c = dt.is_contiguous() ? dt : dt.contiguous();
-  if (dt_c.scalar_type() != at::kFloat) dt_c = dt_c.to(at::kFloat);
-
-  at::Tensor cu_int = cu_seqlens.to(at::kInt).contiguous();
-
-  const int64_t batch = final_states.size(0);
-  const int64_t nheads = final_states.size(1);
-  const int64_t headdim = final_states.size(2);
-  const int64_t dstate = final_states.size(3);
-  const int64_t ngroups = B_in.size(1);
-
-  TORCH_CHECK(final_states.is_contiguous(),
-              "mamba_chunk_scan_fwd_cpu: final_states must be contiguous");
-  TORCH_CHECK(out.is_contiguous(),
-              "mamba_chunk_scan_fwd_cpu: out must be contiguous (writes via "
-              "raw data_ptr)");
-
-  VLLM_DISPATCH_FLOATING_TYPES(input_type, "mamba_chunk_scan_fwd_cpu", [&] {
-    mamba_cpu::mamba_chunk_scan_fwd_kernel<scalar_t>(
-        final_states.data_ptr<float>(), x_in.data_ptr<scalar_t>(),
-        dt_c.data_ptr<float>(), A_f32.data_ptr<float>(),
-        B_in.data_ptr<scalar_t>(), C_in.data_ptr<scalar_t>(),
-        D_f32.defined() ? D_f32.data_ptr<float>() : nullptr,
-        z_in.defined() ? z_in.data_ptr<scalar_t>() : nullptr,
-        out.data_ptr<scalar_t>(), cu_int.data_ptr<int32_t>(), batch, nheads,
-        ngroups, headdim, dstate);
-  });
-}
@@ -1,382 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-//
-// Fused CPU vector kernels for Mamba decode-step hotspots:
-//   - causal_conv1d_update  (depthwise 1-D conv state roll + compute)
-//   - selective_state_update (SSM recurrence, single-step)
-
-#pragma once
-
-#include "cpu_types.hpp"
-#include <cmath>
-#include <cstring>
-#include <cstdint>
-#include <algorithm>
-
-namespace mamba_cpu {
-
-// ---------------------------------------------------------------------------
-// causal_conv1d_update — templated for native BF16/FP32
-//
-// state_ptr may point to a NON-CONTIGUOUS paged KV cache tensor.
-// Explicit strides are passed so the kernel writes directly into the
-// correct memory locations without making a contiguous copy of the full
-// paged tensor (which was the source of the 34-41% direct_copy_kernel).
-//
-//   stride_s_slot  = state.stride(0)  — between cache slots
-//   stride_s_dim   = state.stride(1)  — between conv_dim channels
-//   stride_s_state = state.stride(2)  — between state elements
-//
-// When stride_s_state == 1 (contiguous), the memmove fast path is used.
-// ---------------------------------------------------------------------------
-template <typename scalar_t>
-inline void causal_conv1d_update_kernel(
-    const scalar_t* __restrict__ x_ptr, scalar_t* __restrict__ state_ptr,
-    int64_t stride_s_slot, int64_t stride_s_dim, int64_t stride_s_state,
-    const scalar_t* __restrict__ weight_ptr, const float* __restrict__ bias_ptr,
-    scalar_t* __restrict__ out_ptr, const int32_t* __restrict__ cache_idxs,
-    int32_t pad_slot_id, int64_t batch, int64_t dim, int64_t seqlen,
-    int64_t width, int64_t state_len, bool do_silu) {
-#pragma omp parallel for
-  for (int64_t b = 0; b < batch; ++b) {
-    int64_t cache_idx = (cache_idxs != nullptr) ? cache_idxs[b] : b;
-    if (cache_idx == pad_slot_id) continue;
-
-    for (int64_t t = 0; t < seqlen; ++t) {
-      const scalar_t* x_b = x_ptr + (b * dim * seqlen + t);
-      scalar_t* out_b = out_ptr + (b * dim * seqlen + t);
-      // Base of this slot in the (possibly non-contiguous) paged state
-      scalar_t* s_base = state_ptr + cache_idx * stride_s_slot;
-
-      for (int64_t d = 0; d < dim; ++d) {
-        float x_val = static_cast<float>(x_b[d * seqlen]);
-        scalar_t* sd = s_base + d * stride_s_dim;  // start of this dim's state
-        const scalar_t* w = weight_ptr + d * width;
-
-        // Accumulate in float32 for precision
-        float acc = (bias_ptr != nullptr) ? bias_ptr[d] : 0.0f;
-        for (int64_t k = 0; k < state_len; ++k) {
-          acc += static_cast<float>(w[k]) *
-                 static_cast<float>(sd[k * stride_s_state]);
-        }
-        acc += static_cast<float>(w[state_len]) * x_val;
-
-        // Shift state left and append new input.
-        // Use memmove when contiguous (stride==1); element loop otherwise.
-        if (stride_s_state == 1) {
-          if (state_len > 1)
-            std::memmove(sd, sd + 1, (state_len - 1) * sizeof(scalar_t));
-          if (state_len > 0) sd[state_len - 1] = static_cast<scalar_t>(x_val);
-        } else {
-          for (int64_t k = 0; k < state_len - 1; ++k)
-            sd[k * stride_s_state] = sd[(k + 1) * stride_s_state];
-          if (state_len > 0)
-            sd[(state_len - 1) * stride_s_state] = static_cast<scalar_t>(x_val);
-        }
-
-        if (do_silu) {
-          float sigmoid = (acc >= 0) ? 1.0f / (1.0f + std::exp(-acc))
-                                     : std::exp(acc) / (1.0f + std::exp(acc));
-          acc *= sigmoid;
-        }
-        out_b[d * seqlen] = static_cast<scalar_t>(acc);
-      }
-    }
-  }
-}
-
-// ---------------------------------------------------------------------------
-// selective_state_update
-//
-// Template parameters:
-//   state_t  - dtype of ssm_state cache (typically BFloat16)
-//   input_t  - dtype of x, B, C (typically BFloat16)
-//   out_t    - dtype of output tensor (typically BFloat16)
-//              Write directly — no float32 intermediate buffer needed.
-//
-// A, D, dt_bias are accepted as const float* (they are always float32
-// model parameters in Mamba2).  This eliminates the per-call float32→BF16
-// conversion and the .contiguous() materialisation of the broadcast-expand.
-//
-// dt is accepted as a (N, nheads) scalar-per-head tensor, not as the
-// (N, nheads, head_dim) expansion, so no .contiguous() copy is needed.
-// ---------------------------------------------------------------------------
-template <typename state_t, typename input_t, typename out_t = float>
-inline void selective_state_update_kernel(
-    state_t* __restrict__ state_ptr, int64_t stride_state_n,
-    int64_t stride_state_h, int64_t stride_state_d,
-    const input_t* __restrict__ x_ptr, int64_t stride_x_n, int64_t stride_x_h,
-    // dt: (N, nheads) — scalar per head, NOT expanded to head_dim
-    const float* __restrict__ dt_ptr, int64_t stride_dt_n,
-    // A: (nheads,) float32 — scalar per head
-    const float* __restrict__ A_ptr, const input_t* __restrict__ B_ptr,
-    const input_t* __restrict__ C_ptr, int64_t stride_BC_n, int64_t stride_BC_g,
-    // D: (nheads,) float32 — scalar per head (nullptr if not used)
-    const float* __restrict__ D_ptr,
-    // z: same shape as x (optional)
-    const input_t* __restrict__ z_ptr,
-    // dt_bias: (nheads,) float32 — scalar per head (nullptr if not used)
-    const float* __restrict__ dt_bias_ptr, out_t* __restrict__ out_ptr,
-    int64_t stride_out_n, int64_t stride_out_h,
-    const int32_t* __restrict__ state_batch_indices,
-    const int32_t* __restrict__ dst_state_batch_indices, int32_t null_block_id,
-    const int32_t* __restrict__ num_accepted_tokens,
-    const int32_t* __restrict__ cu_seqlens, int64_t N, int64_t nheads,
-    int64_t ngroups, int64_t dim, int64_t dstate, bool dt_softplus) {
-  using state_vec_t = vec_op::vec_t<state_t>;
-  using input_vec_t = vec_op::vec_t<input_t>;
-  constexpr int VEC_ELEM_NUM = 8;
-
-  int64_t nheads_per_group = nheads / ngroups;
-
-  for (int64_t seq_idx = 0; seq_idx < N; ++seq_idx) {
-    int64_t bos, seq_len;
-    if (cu_seqlens != nullptr) {
-      bos = cu_seqlens[seq_idx];
-      seq_len = cu_seqlens[seq_idx + 1] - bos;
-    } else {
-      bos = seq_idx;
-      seq_len = 1;
-    }
-
-    int64_t state_read_idx = (state_batch_indices != nullptr)
-                                 ? state_batch_indices[seq_idx]
-                                 : seq_idx;
-    if (state_read_idx == null_block_id) continue;
-
-    int64_t state_write_idx = (num_accepted_tokens == nullptr)
-                                  ? ((dst_state_batch_indices != nullptr)
-                                         ? dst_state_batch_indices[seq_idx]
-                                         : state_read_idx)
-                                  : -1;
-
-    state_t* s = state_ptr + state_read_idx * stride_state_n;
-
-    for (int64_t t = 0; t < seq_len; ++t) {
-      int64_t token_idx = bos + t;
-      const input_t* x_tok = x_ptr + token_idx * stride_x_n;
-      // dt: (N, nheads) — one float per head per token
-      const float* dt_tok = dt_ptr + token_idx * stride_dt_n;
-      const input_t* B_tok = B_ptr + token_idx * stride_BC_n;
-      const input_t* C_tok = C_ptr + token_idx * stride_BC_n;
-      out_t* out_tok = out_ptr + token_idx * stride_out_n;
-
-#pragma omp parallel for
-      for (int64_t h = 0; h < nheads; ++h) {
-        int64_t g = h / nheads_per_group;
-        const input_t* x_h = x_tok + h * stride_x_h;
-        const input_t* B_g = B_tok + g * stride_BC_g;
-        const input_t* C_g = C_tok + g * stride_BC_g;
-        out_t* out_h = out_tok + h * stride_out_h;
-        state_t* s_h = s + h * stride_state_h;
-
-        // Read scalars-per-head (A, dt, dt_bias, D) — no per-dim indexing
-        float dt_val = dt_tok[h];
-        if (dt_bias_ptr != nullptr) dt_val += dt_bias_ptr[h];
-        if (dt_softplus) {
-          dt_val = (dt_val <= 20.0f) ? std::log1p(std::exp(dt_val)) : dt_val;
-        }
-        const float A_val = A_ptr[h];  // scalar: same for all dim, dstate
-        const float D_val = (D_ptr != nullptr) ? D_ptr[h] : 0.0f;
-
-        const input_t* z_h =
-            (z_ptr != nullptr) ? z_ptr + token_idx * stride_x_n + h * stride_x_h
-                               : nullptr;
-
-        vec_op::FP32Vec8 dt_vec(dt_val);
-        // dA = exp(A * dt): A and dt are SCALARS per head, so compute once
-        // and broadcast.  This saves 7 redundant std::exp() calls that
-        // FP32Vec8::exp() would otherwise make on the broadcast vector.
-        const float dA_scalar = std::exp(A_val * dt_val);
-        vec_op::FP32Vec8 dA(dA_scalar);  // broadcast
-
-        for (int64_t d = 0; d < dim; ++d) {
-          float x_val = static_cast<float>(x_h[d]);
-
-          vec_op::FP32Vec8 out_vec(0.0f);
-          state_t* s_hd = s_h + d * stride_state_d;
-          const input_t* B_g_base = B_g;
-          const input_t* C_g_base = C_g;
-
-          vec_op::FP32Vec8 x_vec(x_val);
-          // dBx = B * x * dt — same dA for all dstate (A is scalar)
-          // s_new = s * dA + B * x * dt
-
-          int64_t n = 0;
-          for (; n <= dstate - VEC_ELEM_NUM; n += VEC_ELEM_NUM) {
-            vec_op::FP32Vec8 B_v((input_vec_t(B_g_base + n)));
-            vec_op::FP32Vec8 C_v((input_vec_t(C_g_base + n)));
-            vec_op::FP32Vec8 s_v((state_vec_t(s_hd + n)));
-
-            vec_op::FP32Vec8 dBx = B_v * x_vec * dt_vec;
-            vec_op::FP32Vec8 s_new = s_v * dA + dBx;
-
-            state_vec_t(s_new).save(s_hd + n);
-            out_vec = out_vec + s_new * C_v;
-          }
-
-          float out_val = out_vec.reduce_sum();
-          for (; n < dstate; ++n) {
-            // Reuse dA_scalar computed once per head — no exp() re-call
-            float dBx = static_cast<float>(B_g[n]) * x_val * dt_val;
-            float s_new = static_cast<float>(s_hd[n]) * dA_scalar + dBx;
-            s_hd[n] = static_cast<state_t>(s_new);
-            out_val += s_new * static_cast<float>(C_g[n]);
-          }
-
-          if (D_ptr != nullptr) out_val += x_val * D_val;
-          if (z_h != nullptr) {
-            float z_val = static_cast<float>(z_h[d]);
-            float sigmoid = (z_val >= 0)
-                                ? 1.0f / (1.0f + std::exp(-z_val))
-                                : std::exp(z_val) / (1.0f + std::exp(z_val));
-            out_val *= z_val * sigmoid;
-          }
-          out_h[d] = static_cast<out_t>(out_val);
-        }
-      }
-
-      if (num_accepted_tokens != nullptr &&
-          dst_state_batch_indices != nullptr) {
-        int64_t token_dst_idx = dst_state_batch_indices[seq_idx * seq_len + t];
-        if (token_dst_idx != null_block_id && token_dst_idx != state_read_idx) {
-          state_t* dst_s = state_ptr + token_dst_idx * stride_state_n;
-          std::memmove(dst_s, s, nheads * stride_state_h * sizeof(state_t));
-        }
-      }
-    }
-
-    if (num_accepted_tokens == nullptr && state_write_idx != null_block_id &&
-        state_write_idx != state_read_idx) {
-      state_t* dst_s = state_ptr + state_write_idx * stride_state_n;
-      std::memmove(dst_s, s, nheads * stride_state_h * sizeof(state_t));
-    }
-  }
-}
-
-// ---------------------------------------------------------------------------
-// mamba_chunk_scan_fwd
-//
-// Prefill SSM recurrence for Mamba2 / SSD models.
-//
-// Key difference from selective_state_update_kernel (decode path):
-//   - #pragma omp parallel for collapse(2) is OUTSIDE the time loop.
-//     Each thread owns a (batch, head) slice and runs the entire token
-//     sequence without any per-token OpenMP synchronisation overhead.
-//     For seqlen=256, this eliminates 256 thread-barrier launches per batch.
-//
-// `dt` arrives already processed (float32, after bias + softplus + clamp)
-// to keep this kernel simple. Preprocessing is done in the Python wrapper.
-//
-// `states_ptr` points to the [batch, nheads, headdim, dstate] float32 output
-// tensor, pre-initialised by the caller (zero or from initial_states).
-// Each (b, h) slice is private to exactly one thread via collapse(2), so
-// there are no write conflicts.
-//
-// D is treated as a scalar per head ([nheads] float32).
-// ---------------------------------------------------------------------------
-template <typename input_t>
-inline void mamba_chunk_scan_fwd_kernel(
-    float* __restrict__ states_ptr,     // [batch, nheads, headdim, dstate] f32
-    const input_t* __restrict__ x_ptr,  // [seqlen, nheads, headdim]
-    const float* __restrict__ dt_ptr,   // [seqlen, nheads] f32 (preprocessed)
-    const float* __restrict__ A_ptr,    // [nheads] f32
-    const input_t* __restrict__ B_ptr,  // [seqlen, ngroups, dstate]
-    const input_t* __restrict__ C_ptr,  // [seqlen, ngroups, dstate]
-    const float* __restrict__ D_ptr,    // [nheads] f32 (nullable)
-    const input_t* __restrict__ z_ptr,  // [seqlen, nheads, headdim] (nullable)
-    input_t* __restrict__ out_ptr,      // [seqlen, nheads, headdim]
-    const int32_t* __restrict__ cu_seqlens,  // [batch+1] int32
-    int64_t batch, int64_t nheads, int64_t ngroups, int64_t headdim,
-    int64_t dstate) {
-  using input_vec_t = vec_op::vec_t<input_t>;
-  constexpr int VEC_ELEM_NUM = 8;
-
-  const int64_t nheads_per_group = nheads / ngroups;
-  // states layout: [batch, nheads, headdim, dstate] contiguous (caller
-  // guarantee)
-  const int64_t stride_s_b = nheads * headdim * dstate;
-  const int64_t stride_s_h = headdim * dstate;
-  // stride_s_d = dstate, stride_s_n = 1
-
-#pragma omp parallel for collapse(2) schedule(static)
-  for (int64_t b = 0; b < batch; ++b) {
-    for (int64_t h = 0; h < nheads; ++h) {
-      const int64_t seq_start = cu_seqlens[b];
-      const int64_t seq_end = cu_seqlens[b + 1];
-      const int64_t g = h / nheads_per_group;
-
-      const float A_val = A_ptr[h];
-      const float D_val = (D_ptr != nullptr) ? D_ptr[h] : 0.0f;
-
-      // Working state slice: states[b, h, :, :] — float32, headdim * dstate.
-      // Fits in L1/L2 for typical dims (e.g. 64*128*4 = 32 KB).
-      float* s_bh = states_ptr + b * stride_s_b + h * stride_s_h;
-
-      for (int64_t t = seq_start; t < seq_end; ++t) {
-        const input_t* x_h = x_ptr + t * nheads * headdim + h * headdim;
-        const float* dt_h = dt_ptr + t * nheads + h;
-        const input_t* B_g = B_ptr + t * ngroups * dstate + g * dstate;
-        const input_t* C_g = C_ptr + t * ngroups * dstate + g * dstate;
-        const input_t* z_h = (z_ptr != nullptr)
-                                 ? z_ptr + t * nheads * headdim + h * headdim
-                                 : nullptr;
-        input_t* out_h = out_ptr + t * nheads * headdim + h * headdim;
-
-        const float dt_val = *dt_h;
-        const float dA_val = std::exp(A_val * dt_val);
-        const vec_op::FP32Vec8 dA_vec(dA_val);  // broadcast scalar
-        const vec_op::FP32Vec8 dt_vec(dt_val);
-
-        for (int64_t d = 0; d < headdim; ++d) {
-          const float x_val = static_cast<float>(x_h[d]);
-          float* s_bhd = s_bh + d * dstate;  // [dstate] contiguous float32
-
-          // Vectorised SSM update + readout over dstate:
-          //   s_new = s * dA + x * dt * B
-          //   y    += s_new * C
-          int64_t n = 0;
-          vec_op::FP32Vec8 y_vec(0.0f);
-          const vec_op::FP32Vec8 x_vec(x_val);
-
-          for (; n <= dstate - VEC_ELEM_NUM; n += VEC_ELEM_NUM) {
-            const vec_op::FP32Vec8 B_v((input_vec_t(B_g + n)));
-            const vec_op::FP32Vec8 C_v((input_vec_t(C_g + n)));
-            const vec_op::FP32Vec8 s_v(s_bhd + n);
-
-            const vec_op::FP32Vec8 s_new = s_v * dA_vec + x_vec * dt_vec * B_v;
-            s_new.save(s_bhd + n);
-            y_vec = y_vec + s_new * C_v;
-          }
-
-          float y_val = y_vec.reduce_sum();
-
-          // Scalar tail for remaining dstate elements
-          for (; n < dstate; ++n) {
-            const float B_n = static_cast<float>(B_g[n]);
-            const float C_n = static_cast<float>(C_g[n]);
-            const float s_new = s_bhd[n] * dA_val + x_val * dt_val * B_n;
-            s_bhd[n] = s_new;
-            y_val += s_new * C_n;
-          }
-
-          // D skip connection (scalar per head)
-          if (D_ptr != nullptr) y_val += x_val * D_val;
-
-          // z gating: out = y * z * sigmoid(z)  (SiLU)
-          if (z_h != nullptr) {
-            const float z_val = static_cast<float>(z_h[d]);
-            const float sigmoid =
-                (z_val >= 0.0f) ? 1.0f / (1.0f + std::exp(-z_val))
-                                : std::exp(z_val) / (1.0f + std::exp(z_val));
-            y_val *= z_val * sigmoid;
-          }
-
-          out_h[d] = static_cast<input_t>(y_val);
-        }
-      }
-    }
-  }
-}
-
-}  // namespace mamba_cpu
@@ -1116,164 +1116,6 @@ void fused_sigmoid_gating_delta_rule_update_kernel_impl(
  });
 }

-// Speculative-decode variant: processes a varlen batch where each sequence has
-// ``q_len`` draft tokens, runs the recurrence sequentially over those tokens
-// (inside the kernel, so one dispatch handles the whole draft block), reads the
-// initial state from cache slot ``num_accepted-1`` and stores the state *after*
-// token ``t`` into cache slot ``t`` (multi-slot rollback, matching the GPU
-// kernel). Parallelized over (sequence, v_head); the per-sequence token loop is
-// sequential as required by the recurrence.
-template <typename scalar_t, typename param_t>
-void fused_sigmoid_gating_delta_rule_update_spec_kernel_impl(
-    const scalar_t* __restrict__ q_ptr,  // [T, HK, EK]
-    const scalar_t* __restrict__ k_ptr,  // [T, HK, EK]
-    const scalar_t* __restrict__ v_ptr,  // [T, HV, EV]
-    const param_t* __restrict__ A_log_ptr,
-    const scalar_t* __restrict__ a_ptr,  // [T, HV]
-    const scalar_t* __restrict__ dt_bias_ptr,
-    const scalar_t* __restrict__ b_ptr,  // [T, HV]
-    const int32_t* __restrict__ spec_indices_ptr,  // [N, S]
-    const int32_t* __restrict__ num_accepted_ptr,  // [N]
-    const int32_t* __restrict__ cu_seqlens_ptr,     // [N + 1]
-    float* __restrict__ state_ptr,
-    scalar_t* __restrict__ o_ptr,  // [T, HV, EV]
-    float* __restrict__ qk_scale_buf,  // [2, T, HK]
-    int64_t total_tokens,
-    int64_t batch_size,
-    int64_t spec_stride,
-    int64_t num_heads,
-    int64_t head_dim,
-    int64_t v_num_heads,
-    int64_t v_head_dim,
-    int64_t q_strideT,
-    int64_t q_strideH,
-    int64_t k_strideT,
-    int64_t k_strideH,
-    int64_t v_strideT,
-    int64_t v_strideH,
-    int64_t state_slot_stride,
-    bool use_qk_l2norm_in_kernel,
-    double softplus_threshold) {
-  using bVec = at::vec::Vectorized<scalar_t>;
-  using fVec = at::vec::Vectorized<float>;
-  constexpr int64_t VecSize = bVec::size();
-  constexpr int64_t fVecSize = fVec::size();
-  int64_t group_size = v_num_heads / num_heads;
-  double scale = 1 / std::sqrt((double)head_dim);
-  fVec scale_vec = fVec((float)scale);
-
-  if (use_qk_l2norm_in_kernel) {
-    float eps = 1e-5f;
-    at::parallel_for(0, total_tokens * num_heads, 0, [&](int64_t begin, int64_t end) {
-      for (int64_t i = begin; i < end; ++i) {
-        int64_t ti = i / num_heads;
-        int64_t ni = i % num_heads;
-        const scalar_t* qp = q_ptr + ti * q_strideT + ni * q_strideH;
-        const scalar_t* kp = k_ptr + ti * k_strideT + ni * k_strideH;
-        float sq = 0.f, sk = 0.f;
-        for (int64_t d = 0; d < head_dim; ++d) {
-          float qv = (float)qp[d];
-          sq += qv * qv;
-          float kv = (float)kp[d];
-          sk += kv * kv;
-        }
-        qk_scale_buf[ti * num_heads + ni] = 1.f / std::sqrt(sq + eps);
-        qk_scale_buf[total_tokens * num_heads + ti * num_heads + ni] = 1.f / std::sqrt(sk + eps);
-      }
-    });
-  }
-
-  at::parallel_for(0, batch_size * v_num_heads, 0, [&](int64_t begin, int64_t end) {
-    for (int64_t idx = begin; idx < end; ++idx) {
-      int64_t bi = idx / v_num_heads;
-      int64_t ni = idx % v_num_heads;
-      int64_t kh = ni / group_size;
-      int64_t q_start = cu_seqlens_ptr[bi];
-      int64_t q_len = cu_seqlens_ptr[bi + 1] - q_start;
-      if (q_len <= 0) {
-        continue;
-      }
-      int64_t acc = (int64_t)num_accepted_ptr[bi];
-      // Clamp acc-1 to >=0: when num_accepted is 0 the unclamped index reads
-      // out of bounds and yields an arbitrary prev_slot used to index the SSM
-      // state. Mirrors the GPU guard tl.maximum(num_accepted - 1, 0).
-      int64_t prev_slot =
-          (int64_t)spec_indices_ptr[bi * spec_stride + (acc > 0 ? acc - 1 : 0)];
-      for (int64_t t = 0; t < q_len; ++t) {
-        int64_t cur_slot = (int64_t)spec_indices_ptr[bi * spec_stride + t];
-        int64_t token = q_start + t;
-        const float* src = state_ptr + prev_slot * state_slot_stride + ni * head_dim * v_head_dim;
-        float* dst = state_ptr + cur_slot * state_slot_stride + ni * head_dim * v_head_dim;
-        float g_val = -std::exp((float)A_log_ptr[ni]) *
-            softplus((float)a_ptr[token * v_num_heads + ni] + (float)dt_bias_ptr[ni], softplus_threshold);
-        float g_val_exp = std::exp(g_val);
-        fVec g_val_exp_vec = fVec(g_val_exp);
-        float beta_val = 1.f / (1.f + std::exp(-(float)b_ptr[token * v_num_heads + ni]));
-        fVec beta_vec = fVec(beta_val);
-        int64_t q_offset = token * q_strideT + kh * q_strideH;
-        int64_t k_offset = token * k_strideT + kh * k_strideH;
-        float q_scale = use_qk_l2norm_in_kernel ? qk_scale_buf[token * num_heads + kh] : 1.f;
-        float k_scale =
-            use_qk_l2norm_in_kernel ? qk_scale_buf[total_tokens * num_heads + token * num_heads + kh] : 1.f;
-        int64_t v_offset = token * v_strideT + ni * v_strideH;
-        int64_t o_offset = (token * v_num_heads + ni) * v_head_dim;
-        int64_t dvi = 0;
-        for (; dvi <= v_head_dim - VecSize; dvi += VecSize) {
-          fVec kv_mem_vec0 = fVec(0.f);
-          fVec kv_mem_vec1 = fVec(0.f);
-          for (int di = 0; di < head_dim; ++di) {
-            fVec k_val_vec = fVec((float)k_ptr[k_offset + di] * k_scale);
-            fVec sv0 = fVec::loadu(src + di * v_head_dim + dvi);
-            fVec sv1 = fVec::loadu(src + di * v_head_dim + dvi + fVecSize);
-            kv_mem_vec0 = kv_mem_vec0 + sv0 * g_val_exp_vec * k_val_vec;
-            kv_mem_vec1 = kv_mem_vec1 + sv1 * g_val_exp_vec * k_val_vec;
-          }
-          bVec v_bvec = bVec::loadu(v_ptr + v_offset + dvi);
-          fVec v_vec0, v_vec1;
-          std::tie(v_vec0, v_vec1) = at::vec::convert_to_float(v_bvec);
-          fVec dt_vec0 = (v_vec0 - kv_mem_vec0) * beta_vec;
-          fVec dt_vec1 = (v_vec1 - kv_mem_vec1) * beta_vec;
-          fVec o_vec0 = fVec(0.f);
-          fVec o_vec1 = fVec(0.f);
-          for (int di = 0; di < head_dim; ++di) {
-            fVec q_vec = fVec((float)q_ptr[q_offset + di] * q_scale);
-            fVec k_vec = fVec((float)k_ptr[k_offset + di] * k_scale);
-            fVec sv0 = fVec::loadu(src + di * v_head_dim + dvi);
-            fVec sv1 = fVec::loadu(src + di * v_head_dim + dvi + fVecSize);
-            sv0 = sv0 * g_val_exp_vec + k_vec * dt_vec0;
-            sv1 = sv1 * g_val_exp_vec + k_vec * dt_vec1;
-            o_vec0 = o_vec0 + sv0 * q_vec * scale_vec;
-            o_vec1 = o_vec1 + sv1 * q_vec * scale_vec;
-            sv0.store(dst + di * v_head_dim + dvi);
-            sv1.store(dst + di * v_head_dim + dvi + fVecSize);
-          }
-          bVec o_vec = at::vec::convert_from_float<scalar_t>(o_vec0, o_vec1);
-          o_vec.store(o_ptr + o_offset + dvi);
-        }
-        for (; dvi < v_head_dim; ++dvi) {
-          float kv_mem_val = 0.f;
-          for (int di = 0; di < head_dim; ++di) {
-            float k_val = (float)k_ptr[k_offset + di] * k_scale;
-            kv_mem_val += src[di * v_head_dim + dvi] * g_val_exp * k_val;
-          }
-          float v_val = (float)v_ptr[v_offset + dvi];
-          float dt_val = (v_val - kv_mem_val) * beta_val;
-          float o_val = 0.f;
-          for (int di = 0; di < head_dim; ++di) {
-            float q_val = (float)q_ptr[q_offset + di] * q_scale;
-            float k_val = (float)k_ptr[k_offset + di] * k_scale;
-            float ns = src[di * v_head_dim + dvi] * g_val_exp + k_val * dt_val;
-            dst[di * v_head_dim + dvi] = ns;
-            o_val += ns * q_val * scale;
-          }
-          o_ptr[o_offset + dvi] = (scalar_t)o_val;
-        }
-        prev_slot = cur_slot;
-      }
-    }
-  });
-}
-
 template <typename scalar_t>
 void fused_gdn_gating_kernel_impl(
    float* __restrict__ A_log,
@@ -1658,103 +1500,6 @@ at::Tensor fused_sigmoid_gating_delta_rule_update_cpu(
  return core_attn_out;
 }

-// Speculative-decode update (multi-token, multi-slot rollback).
-// q: [T, HK, EK]  k: [T, HK, EK]  v: [T, HV, EV]
-// a: [T, HV]  b: [T, HV]
-// initial_state_source: [N_slots, HV, EK, EV] FP32 (updated in place)
-// spec_state_indices: [batch, S] INT32 (S = num_spec + 1)
-// num_accepted_tokens: [batch] INT32
-// cu_seqlens: [batch + 1] INT32
-// Returns output: [T, HV, EV]
-at::Tensor fused_sigmoid_gating_delta_rule_update_spec_cpu(
-    const at::Tensor& A_log,
-    const at::Tensor& dt_bias,
-    const at::Tensor& q,
-    const at::Tensor& k,
-    const at::Tensor& v,
-    const at::Tensor& a,
-    const at::Tensor& b,
-    at::Tensor& initial_state_source,
-    const at::Tensor& spec_state_indices,
-    const at::Tensor& num_accepted_tokens,
-    const at::Tensor& cu_seqlens,
-    bool use_qk_l2norm_in_kernel,
-    double softplus_beta = 1.0,
-    double softplus_threshold = 20.0) {
-  CHECK_DIM(3, q);
-  CHECK_DIM(3, v);
-  CHECK_LAST_DIM_CONTIGUOUS_INPUT(q);
-  int64_t total_tokens = q.size(0);
-  int64_t num_heads = q.size(1);
-  int64_t head_dim = q.size(2);
-  int64_t v_num_heads = v.size(1);
-  int64_t v_head_dim = v.size(2);
-  int64_t batch_size = cu_seqlens.size(0) - 1;
-  int64_t spec_stride = spec_state_indices.stride(0);
-  CHECK_INPUT_SHAPE_DTYPE<true>(k, {total_tokens, num_heads, head_dim}, q.scalar_type());
-  CHECK_INPUT_SHAPE_DTYPE<true>(v, {total_tokens, v_num_heads, v_head_dim}, q.scalar_type());
-  CHECK_INPUT_SHAPE_DTYPE<true>(a, {total_tokens, v_num_heads}, q.scalar_type());
-  CHECK_INPUT_SHAPE_DTYPE<true>(b, {total_tokens, v_num_heads}, q.scalar_type());
-  CHECK_INPUT_SHAPE_DTYPE<true>(dt_bias, {v_num_heads}, q.scalar_type());
-  CHECK_INPUT_SHAPE_DTYPE<true>(num_accepted_tokens, {batch_size}, at::kInt);
-  CHECK_INPUT_SHAPE_DTYPE<true>(cu_seqlens, {batch_size + 1}, at::kInt);
-  CHECK_EQ(v_num_heads % num_heads, 0);
-  TORCH_CHECK(A_log.sizes() == at::IntArrayRef({v_num_heads}));
-  CHECK_INPUT_SHAPE_DTYPE<true>(
-      initial_state_source,
-      {initial_state_source.size(0), v_num_heads, head_dim, v_head_dim},
-      at::kFloat);
-  TORCH_CHECK(initial_state_source.size(0) >= batch_size,
-      "initial_state_source capacity too small: size(0)=",
-      initial_state_source.size(0), ", batch_size=", batch_size);
-
-  int64_t q_strideT = q.stride(0);
-  int64_t q_strideH = q.stride(1);
-  int64_t k_strideT = k.stride(0);
-  int64_t k_strideH = k.stride(1);
-  int64_t v_strideT = v.stride(0);
-  int64_t v_strideH = v.stride(1);
-  int64_t state_slot_stride = initial_state_source.stride(0);
-
-  at::Tensor o = at::empty({total_tokens, v_num_heads, v_head_dim}, q.options());
-  at::Tensor qk_scale_buf = at::empty({2, total_tokens, num_heads}, at::kFloat);
-
-  CPU_DISPATCH_REDUCED_FLOATING_TYPES_EXT(
-      q.scalar_type(), A_log.scalar_type(), "fused_sigmoid_gating_delta_rule_update_spec_kernel_impl", [&] {
-        fused_sigmoid_gating_delta_rule_update_spec_kernel_impl<scalar_t, param_t>(
-            q.data_ptr<scalar_t>(),
-            k.data_ptr<scalar_t>(),
-            v.data_ptr<scalar_t>(),
-            A_log.data_ptr<param_t>(),
-            a.data_ptr<scalar_t>(),
-            dt_bias.data_ptr<scalar_t>(),
-            b.data_ptr<scalar_t>(),
-            spec_state_indices.data_ptr<int32_t>(),
-            num_accepted_tokens.data_ptr<int32_t>(),
-            cu_seqlens.data_ptr<int32_t>(),
-            initial_state_source.data_ptr<float>(),
-            o.data_ptr<scalar_t>(),
-            qk_scale_buf.data_ptr<float>(),
-            total_tokens,
-            batch_size,
-            spec_stride,
-            num_heads,
-            head_dim,
-            v_num_heads,
-            v_head_dim,
-            q_strideT,
-            q_strideH,
-            k_strideT,
-            k_strideH,
-            v_strideT,
-            v_strideH,
-            state_slot_stride,
-            use_qk_l2norm_in_kernel,
-            softplus_threshold);
-      });
-  return o;
-}
-
 // A_log: [num_v_heads]
 // a: [batch, num_v_heads]
 // b: [batch, num_v_heads]
@@ -120,14 +120,6 @@ at::Tensor fused_sigmoid_gating_delta_rule_update_cpu(
    bool use_qk_l2norm_in_kernel, double softplus_beta = 1.0,
    double softplus_threshold = 20.0);

-at::Tensor fused_sigmoid_gating_delta_rule_update_spec_cpu(
-    const at::Tensor& A_log, const at::Tensor& dt_bias, const at::Tensor& q,
-    const at::Tensor& k, const at::Tensor& v, const at::Tensor& a,
-    const at::Tensor& b, at::Tensor& initial_state_source,
-    const at::Tensor& spec_state_indices, const at::Tensor& num_accepted_tokens,
-    const at::Tensor& cu_seqlens, bool use_qk_l2norm_in_kernel,
-    double softplus_beta = 1.0, double softplus_threshold = 20.0);
-
 std::tuple<at::Tensor, at::Tensor> fused_gdn_gating_cpu(
    const at::Tensor& A_log, const at::Tensor& a, const at::Tensor& b,
    const at::Tensor& dt_bias);
@@ -213,32 +205,6 @@ void compute_slot_mapping_kernel_impl(const torch::Tensor query_start_loc,
                                      torch::Tensor slot_mapping,
                                      const int64_t block_size);

-at::Tensor causal_conv1d_update_cpu_impl(
-    at::Tensor& x, at::Tensor& conv_state, const at::Tensor& weight,
-    const c10::optional<at::Tensor>& bias,
-    const c10::optional<std::string>& activation,
-    const c10::optional<at::Tensor>& conv_state_indices,
-    const c10::optional<at::Tensor>& query_start_loc, int64_t pad_slot_id);
-
-void selective_state_update_cpu_impl(
-    at::Tensor& state, const at::Tensor& x, const at::Tensor& dt,
-    const at::Tensor& A, const at::Tensor& B, const at::Tensor& C,
-    const c10::optional<at::Tensor>& D, const c10::optional<at::Tensor>& z,
-    const c10::optional<at::Tensor>& dt_bias, bool dt_softplus,
-    const c10::optional<at::Tensor>& state_batch_indices,
-    const c10::optional<at::Tensor>& dst_state_batch_indices,
-    int64_t null_block_id, at::Tensor& out,
-    const c10::optional<at::Tensor>& num_accepted_tokens,
-    const c10::optional<at::Tensor>& cu_seqlens);
-
-void mamba_chunk_scan_fwd_cpu_impl(at::Tensor& out, at::Tensor& final_states,
-                                   const at::Tensor& x, const at::Tensor& dt,
-                                   const at::Tensor& A, const at::Tensor& B,
-                                   const at::Tensor& C,
-                                   const c10::optional<at::Tensor>& D,
-                                   const c10::optional<at::Tensor>& z,
-                                   const at::Tensor& cu_seqlens);
-
 void init_cpu_memory_env(std::vector<int64_t> node_ids);

 namespace cpu_utils {
@@ -542,15 +508,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      "softplus_threshold=20.0) -> Tensor");
  ops.impl("fused_sigmoid_gating_delta_rule_update_cpu", torch::kCPU,
           &fused_sigmoid_gating_delta_rule_update_cpu);
-  ops.def(
-      "fused_sigmoid_gating_delta_rule_update_spec_cpu(Tensor A_log, Tensor "
-      "dt_bias, Tensor q, Tensor k, Tensor v, Tensor a, Tensor b, "
-      "Tensor(a!) initial_state_source, Tensor spec_state_indices, "
-      "Tensor num_accepted_tokens, Tensor cu_seqlens, bool "
-      "use_qk_l2norm_in_kernel, float softplus_beta=1.0, float "
-      "softplus_threshold=20.0) -> Tensor");
-  ops.impl("fused_sigmoid_gating_delta_rule_update_spec_cpu", torch::kCPU,
-           &fused_sigmoid_gating_delta_rule_update_spec_cpu);
  ops.def(
      "fused_gdn_gating_cpu(Tensor A_log, Tensor a, Tensor b, Tensor dt_bias) "
      "-> (Tensor, Tensor)");
@@ -596,8 +553,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
 #endif

  // fused moe
-#if defined(__AVX512F__) || \
-    (defined(__aarch64__) && !defined(__APPLE__) && defined(ARM_BF16_SUPPORT))
+#if defined(__AVX512F__) || (defined(ARM_BF16_SUPPORT))
  ops.def(
      "prepack_moe_weight(Tensor weight, Tensor(a1!) packed_weight, str isa) "
      "-> ()");
@@ -608,7 +564,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      "bool skip_weighted, "
      "str act, str isa) -> ()");
  ops.impl("cpu_fused_moe", torch::kCPU, &cpu_fused_moe);
-#endif
+#endif  // #if defined(__AVX512F__) || (defined(ARM_BF16_SUPPORT))
  ops.def(
      "mla_decode_kvcache("
      "   Tensor! out, Tensor query, Tensor kv_cache,"
@@ -621,30 +577,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      "block_size) -> ()",
      &compute_slot_mapping_kernel_impl);

-  // Mamba CPU kernels
-  ops.def(
-      "causal_conv1d_update_cpu_vec("
-      "Tensor(a0!) x, Tensor(a1!) conv_state, Tensor weight, "
-      "Tensor? bias, str? activation, Tensor? conv_state_indices, "
-      "Tensor? query_start_loc, SymInt pad_slot_id) -> Tensor",
-      &causal_conv1d_update_cpu_impl);
-
-  ops.def(
-      "selective_state_update_cpu("
-      "Tensor(a0!) state, Tensor x, Tensor dt, Tensor A, Tensor B, Tensor C, "
-      "Tensor? D, Tensor? z, Tensor? dt_bias, bool dt_softplus, "
-      "Tensor? state_batch_indices, Tensor? dst_state_batch_indices, "
-      "SymInt null_block_id, Tensor(a13!) out, "
-      "Tensor? num_accepted_tokens, Tensor? cu_seqlens) -> ()",
-      &selective_state_update_cpu_impl);
-
-  ops.def(
-      "mamba_chunk_scan_fwd_cpu("
-      "Tensor(a0!) out, Tensor(a1!) final_states, "
-      "Tensor x, Tensor dt, Tensor A, Tensor B, Tensor C, "
-      "Tensor? D, Tensor? z, Tensor cu_seqlens) -> ()",
-      &mamba_chunk_scan_fwd_cpu_impl);
-
  ops.def("init_cpu_memory_env(SymInt[] node_ids) -> ()", &init_cpu_memory_env);

  // Speculative decoding kernels
@@ -1,326 +0,0 @@
-#pragma once
-
-#include "custom_collective_common.cuh"
-
-namespace vllm {
-
-constexpr int kMnnvlLamportAgThreads = 128;
-constexpr int kMnnvlLamportRsThreads = 256;
-constexpr int kMnnvlLamportConcurrentPollMaxPacks = 8192;
-
-using CopyPack = array_t<uint64_t, 2>;
-
-template <int ngpus>
-__global__ void __launch_bounds__(512, 1)
-    cross_device_all_gather(RankData* _dp, RankSignals sg, Signal* self_sg,
-                            CopyPack* __restrict__ result, int rank,
-                            int size_per_rank) {
-  auto dp = *_dp;
-  int tid = blockIdx.x * blockDim.x + threadIdx.x;
-  int stride = gridDim.x * blockDim.x;
-  barrier_at_start<ngpus>(sg, self_sg, rank);
-#pragma unroll
-  for (int src_rank = 0; src_rank < ngpus; ++src_rank) {
-    auto src = reinterpret_cast<const CopyPack*>(dp.ptrs[src_rank]);
-    auto dst = result + src_rank * size_per_rank;
-    for (int idx = tid; idx < size_per_rank; idx += stride) {
-      dst[idx] = src[idx];
-    }
-  }
-  barrier_at_end<ngpus, true>(sg, self_sg, rank);
-}
-
-template <typename T, int ngpus>
-__global__ void __launch_bounds__(512, 1)
-    cross_device_reduce_scatter(RankData* _dp, RankSignals sg, Signal* self_sg,
-                                T* __restrict__ result, int rank,
-                                int size_per_rank) {
-  using P = typename packed_t<T>::P;
-  using A = typename packed_t<T>::A;
-  auto dp = *_dp;
-  auto offset = rank * size_per_rank;
-  barrier_at_start<ngpus>(sg, self_sg, rank);
-  for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size_per_rank;
-       idx += gridDim.x * blockDim.x) {
-    reinterpret_cast<P*>(result)[idx] =
-        packed_reduce<P, ngpus, A>((const P**)&dp.ptrs[0], offset + idx);
-  }
-  barrier_at_end<ngpus, true>(sg, self_sg, rank);
-}
-
-template <typename P>
-union LamportPack {
-  P packed;
-  uint32_t words[sizeof(P) / sizeof(uint32_t)];
-};
-
-template <typename P>
-DINLINE LamportPack<P> load_lamport_pack(const P* ptr) {
-  static_assert(sizeof(P) == 16);
-  LamportPack<P> value;
-#if !defined(USE_ROCM)
-  asm volatile("ld.volatile.global.v4.u32 {%0, %1, %2, %3}, [%4];"
-               : "=r"(value.words[0]), "=r"(value.words[1]),
-                 "=r"(value.words[2]), "=r"(value.words[3])
-               : "l"(ptr)
-               : "memory");
-#else
-  const volatile uint32_t* src = reinterpret_cast<const volatile uint32_t*>(ptr);
-#pragma unroll
-  for (int i = 0; i < sizeof(P) / sizeof(uint32_t); ++i) {
-    value.words[i] = src[i];
-  }
-#endif
-  return value;
-}
-
-template <typename P>
-DINLINE bool is_lamport_dirty(const LamportPack<P>& value) {
-#pragma unroll
-  for (int i = 0; i < sizeof(P) / sizeof(uint32_t); ++i) {
-    if (value.words[i] == 0x80000000U) return true;
-  }
-  return false;
-}
-
-template <typename P>
-DINLINE P lamport_sentinel() {
-  LamportPack<P> value;
-#pragma unroll
-  for (int i = 0; i < sizeof(P) / sizeof(uint32_t); ++i) {
-    value.words[i] = 0x80000000U;
-  }
-  return value.packed;
-}
-
-template <typename P>
-DINLINE P sanitize_lamport_payload(P packed) {
-  LamportPack<P> value{.packed = packed};
-#pragma unroll
-  for (int i = 0; i < sizeof(P) / sizeof(uint32_t); ++i) {
-    if (value.words[i] == 0x80000000U) value.words[i] = 0;
-  }
-  return value.packed;
-}
-
-template <typename P>
-DINLINE P wait_lamport_payload(const P* ptr) {
-  auto value = load_lamport_pack(ptr);
-  while (is_lamport_dirty(value)) value = load_lamport_pack(ptr);
-  return value.packed;
-}
-
-template <typename P, int ngpus>
-DINLINE void wait_lamport_payloads(const P* base, int rank, int rank_stride,
-                                   P local_value, P (&values)[ngpus]) {
-  bool ready[ngpus];
-#pragma unroll
-  for (int src_rank = 0; src_rank < ngpus; ++src_rank) {
-    ready[src_rank] = src_rank == rank;
-    if (src_rank == rank) values[src_rank] = local_value;
-  }
-
-  int remaining = ngpus - 1;
-  while (remaining != 0) {
-#pragma unroll
-    for (int src_rank = 0; src_rank < ngpus; ++src_rank) {
-      if (!ready[src_rank]) {
-        auto value = load_lamport_pack(base + src_rank * rank_stride);
-        if (!is_lamport_dirty(value)) {
-          values[src_rank] = value.packed;
-          ready[src_rank] = true;
-          --remaining;
-        }
-      }
-    }
-  }
-}
-
-template <typename P, typename A, int ngpus>
-DINLINE P reduce_lamport_payloads(const P* current_local, const P* packed_input,
-                                  int rank, int size_per_rank, int idx) {
-  P source_zero =
-      rank == 0 ? packed_input[idx] : wait_lamport_payload(current_local + idx);
-  A tmp = upcast(source_zero);
-#pragma unroll
-  for (int src_rank = 1; src_rank < ngpus; ++src_rank) {
-    P value = src_rank == rank
-                  ? packed_input[rank * size_per_rank + idx]
-                  : wait_lamport_payload(current_local +
-                                         src_rank * size_per_rank + idx);
-    packed_assign_add(tmp, upcast(value));
-  }
-  return sanitize_lamport_payload(downcast<P>(tmp));
-}
-
-DINLINE void lamport_cta_arrive(uint32_t* counter) {
-#if !defined(USE_ROCM)
-  if (threadIdx.x < 32) {
-    asm volatile("barrier.cta.sync 1, %0;" : : "r"(blockDim.x) : "memory");
-    if (threadIdx.x == 0) {
-  #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000
-      asm volatile("red.async.release.global.gpu.add.u32 [%0], 1;"
-                   :
-                   : "l"(counter)
-                   : "memory");
-  #elif defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
-      asm volatile("red.release.global.gpu.add.u32 [%0], 1;"
-                   :
-                   : "l"(counter)
-                   : "memory");
-  #else
-      atomicAdd(counter, 1);
-  #endif
-    }
-  } else {
-    asm volatile("barrier.cta.arrive 1, %0;" : : "r"(blockDim.x) : "memory");
-  }
-#else
-  __syncthreads();
-  if (threadIdx.x == 0) atomicAdd(counter, 1);
-#endif
-}
-
-template <typename T, int ngpus>
-__global__ void __launch_bounds__(kMnnvlLamportAgThreads, 1)
-    mnnvl_lamport_all_gather(RankData* _dp, const T* __restrict__ input,
-                             T* __restrict__ result,
-                             T* __restrict__ multicast_buffer,
-                             uint32_t* __restrict__ epochs, int rank,
-                             int size_per_rank, int stage_size) {
-  using P = typename packed_t<T>::P;
-#if !defined(USE_ROCM) && CUDA_VERSION >= 12000 && defined(__CUDA_ARCH__) && \
-    (__CUDA_ARCH__ >= 900)
-  cudaGridDependencySynchronize();
-#endif
-  auto dp = *_dp;
-  int tid = blockIdx.x * blockDim.x + threadIdx.x;
-  int stride = gridDim.x * blockDim.x;
-  uint32_t epoch = epochs[0];
-  int current_stage = epoch % 3;
-  int dirty_stage = (epoch + 1) % 3;
-  int dirty_size = epochs[2 + dirty_stage];
-  auto local_buffer = reinterpret_cast<P*>(const_cast<void*>(dp.ptrs[rank]));
-  auto current_local = local_buffer + current_stage * stage_size;
-  auto dirty_local = local_buffer + dirty_stage * stage_size;
-  auto current_multicast =
-      reinterpret_cast<P*>(multicast_buffer) + current_stage * stage_size;
-  auto packed_input = reinterpret_cast<const P*>(input);
-  auto packed_result = reinterpret_cast<P*>(result);
-
-  int total_size = size_per_rank * ngpus;
-  P local_value;
-  if (tid < size_per_rank) {
-    local_value = packed_input[tid];
-    current_multicast[rank * size_per_rank + tid] =
-        sanitize_lamport_payload(local_value);
-  }
-#if !defined(USE_ROCM) && CUDA_VERSION >= 12000 && defined(__CUDA_ARCH__) && \
-    (__CUDA_ARCH__ >= 900)
-  cudaTriggerProgrammaticLaunchCompletion();
-#endif
-
-  lamport_cta_arrive(&epochs[1]);
-
-  for (int idx = tid; idx < dirty_size; idx += stride) {
-    dirty_local[idx] = lamport_sentinel<P>();
-  }
-
-  if (tid < size_per_rank) {
-#pragma unroll
-    for (int src_rank = 0; src_rank < ngpus; ++src_rank) {
-      int output_idx = src_rank * size_per_rank + tid;
-      P value = src_rank == rank
-                    ? local_value
-                    : wait_lamport_payload(current_local + output_idx);
-      packed_result[output_idx] = value;
-    }
-  }
-
-  if (tid == 0) {
-    while (*reinterpret_cast<volatile uint32_t*>(&epochs[1]) < gridDim.x);
-    epochs[2 + current_stage] = total_size;
-    epochs[0] = epoch + 1;
-    epochs[1] = 0;
-  }
-}
-
-template <typename T, int ngpus>
-__global__ void __launch_bounds__(kMnnvlLamportRsThreads, 1)
-    mnnvl_lamport_reduce_scatter_kernel(RankData* _dp,
-                                        const T* __restrict__ input,
-                                        T* __restrict__ result,
-                                        uint32_t* __restrict__ epochs, int rank,
-                                        int size_per_rank, int stage_size) {
-  using P = typename packed_t<T>::P;
-  using A = typename packed_t<T>::A;
-#if !defined(USE_ROCM) && CUDA_VERSION >= 12000 && defined(__CUDA_ARCH__) && \
-    (__CUDA_ARCH__ >= 900)
-  cudaGridDependencySynchronize();
-#endif
-  auto dp = *_dp;
-  int dst_rank = blockIdx.x % ngpus;
-  int tile = blockIdx.x / ngpus;
-  int idx = tile * blockDim.x + threadIdx.x;
-  int tid = blockIdx.x * blockDim.x + threadIdx.x;
-  int stride = gridDim.x * blockDim.x;
-  uint32_t epoch = epochs[0];
-  int current_stage = epoch % 3;
-  int dirty_stage = (epoch + 1) % 3;
-  int dirty_size = epochs[2 + dirty_stage];
-  auto local_buffer = reinterpret_cast<P*>(const_cast<void*>(dp.ptrs[rank]));
-  auto current_local = local_buffer + current_stage * stage_size;
-  auto dirty_local = local_buffer + dirty_stage * stage_size;
-  auto packed_input = reinterpret_cast<const P*>(input);
-
-  if (idx < size_per_rank && dst_rank != rank) {
-    auto dst = reinterpret_cast<P*>(const_cast<void*>(dp.ptrs[dst_rank])) +
-               current_stage * stage_size + rank * size_per_rank;
-    auto src = packed_input + dst_rank * size_per_rank;
-    dst[idx] = sanitize_lamport_payload(src[idx]);
-  }
-#if !defined(USE_ROCM) && CUDA_VERSION >= 12000 && defined(__CUDA_ARCH__) && \
-    (__CUDA_ARCH__ >= 900)
-  cudaTriggerProgrammaticLaunchCompletion();
-#endif
-
-  lamport_cta_arrive(&epochs[1]);
-
-  for (int idx = tid; idx < dirty_size; idx += stride) {
-    dirty_local[idx] = lamport_sentinel<P>();
-  }
-
-  if (idx < size_per_rank && dst_rank == rank) {
-    if constexpr (ngpus == 4) {
-      if (size_per_rank > kMnnvlLamportConcurrentPollMaxPacks) {
-        reinterpret_cast<P*>(result)[idx] =
-            reduce_lamport_payloads<P, A, ngpus>(current_local, packed_input,
-                                                 rank, size_per_rank, idx);
-      } else {
-        P values[ngpus];
-        wait_lamport_payloads<P, ngpus>(
-            current_local + idx, rank, size_per_rank,
-            packed_input[rank * size_per_rank + idx], values);
-        A tmp = upcast(values[0]);
-#pragma unroll
-        for (int src_rank = 1; src_rank < ngpus; ++src_rank) {
-          packed_assign_add(tmp, upcast(values[src_rank]));
-        }
-        reinterpret_cast<P*>(result)[idx] =
-            sanitize_lamport_payload(downcast<P>(tmp));
-      }
-    } else {
-      reinterpret_cast<P*>(result)[idx] = reduce_lamport_payloads<P, A, ngpus>(
-          current_local, packed_input, rank, size_per_rank, idx);
-    }
-  }
-
-  if (tid == 0) {
-    while (*reinterpret_cast<volatile uint32_t*>(&epochs[1]) < gridDim.x);
-    epochs[2 + current_stage] = size_per_rank * ngpus;
-    epochs[0] = epoch + 1;
-    epochs[1] = 0;
-  }
-}
-
-}  // namespace vllm
@@ -1,8 +1,299 @@
 #pragma once

-#include "custom_collective_common.cuh"
+#include <cuda.h>
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+
+#if defined(USE_ROCM)
+typedef __hip_bfloat16 nv_bfloat16;
+#endif
+
+#include <iostream>
+#include <array>
+#include <limits>
+#include <map>
+#include <unordered_map>
+#include <vector>
+#include <cstdlib>
+#include <cstring>

 namespace vllm {
+#define CUDACHECK(cmd)                                              \
+  do {                                                              \
+    cudaError_t e = cmd;                                            \
+    if (e != cudaSuccess) {                                         \
+      printf("Failed: Cuda error %s:%d '%s'\n", __FILE__, __LINE__, \
+             cudaGetErrorString(e));                                \
+      exit(EXIT_FAILURE);                                           \
+    }                                                               \
+  } while (0)
+
+// Maximal number of blocks in allreduce kernel.
+constexpr int kMaxBlocks = 36;
+
+// Default number of blocks in allreduce kernel.
+#ifndef USE_ROCM
+const int defaultBlockLimit = 36;
+CUpointer_attribute rangeStartAddrAttr = CU_POINTER_ATTRIBUTE_RANGE_START_ADDR;
+#else
+const int defaultBlockLimit = 16;
+hipPointer_attribute rangeStartAddrAttr =
+    HIP_POINTER_ATTRIBUTE_RANGE_START_ADDR;
+#endif
+
+// Counter may overflow, but it's fine since unsigned int overflow is
+// well-defined behavior.
+using FlagType = uint32_t;
+
+// Two sets of peer counters are needed for two syncs: starting and ending an
+// operation. The reason is that it's possible for peer GPU block to arrive at
+// the second sync point while the current GPU block haven't passed the first
+// sync point. Thus, peer GPU may write counter+1 while current GPU is busy
+// waiting for counter. We use alternating counter array to avoid this
+// possibility.
+struct Signal {
+  alignas(128) FlagType start[kMaxBlocks][8];
+  alignas(128) FlagType end[kMaxBlocks][8];
+  alignas(128) FlagType _flag[kMaxBlocks];  // incremental flags for each rank
+};
+
+struct __align__(16) RankData {
+  const void* ptrs[8];
+};
+
+struct __align__(16) RankSignals {
+  Signal* signals[8];
+};
+
+// like std::array, but aligned
+template <typename T, int sz>
+struct __align__(alignof(T) * sz) array_t {
+  T data[sz];
+  using type = T;
+  static constexpr int size = sz;
+};
+
+// use packed type to maximize memory efficiency
+// goal: generate ld.128 and st.128 instructions
+template <typename T>
+struct packed_t {
+  // the (P)acked type for load/store
+  using P = array_t<T, 16 / sizeof(T)>;
+  // the (A)ccumulator type for reduction
+  using A = array_t<float, 16 / sizeof(T)>;
+};
+
+#define DINLINE __device__ __forceinline__
+
+// scalar cast functions
+DINLINE float upcast_s(half val) { return __half2float(val); }
+
+template <typename T>
+DINLINE T downcast_s(float val);
+template <>
+DINLINE half downcast_s(float val) {
+  return __float2half(val);
+}
+
+// scalar add functions
+// for some reason when compiling with Pytorch, the + operator for half and
+// bfloat is disabled so we call the intrinsics directly
+DINLINE half& assign_add(half& a, half b) {
+  a = __hadd(a, b);
+  return a;
+}
+DINLINE float& assign_add(float& a, float b) { return a += b; }
+
+#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
+DINLINE float upcast_s(nv_bfloat16 val) { return __bfloat162float(val); }
+template <>
+DINLINE nv_bfloat16 downcast_s(float val) {
+  return __float2bfloat16(val);
+}
+DINLINE nv_bfloat16& assign_add(nv_bfloat16& a, nv_bfloat16 b) {
+  a = __hadd(a, b);
+  return a;
+}
+#endif
+
+template <typename T, int N>
+DINLINE array_t<T, N>& packed_assign_add(array_t<T, N>& a, array_t<T, N> b) {
+#pragma unroll
+  for (int i = 0; i < N; i++) {
+    assign_add(a.data[i], b.data[i]);
+  }
+  return a;
+}
+
+template <typename T, int N>
+DINLINE array_t<float, N> upcast(array_t<T, N> val) {
+  if constexpr (std::is_same<T, float>::value) {
+    return val;
+  } else {
+    array_t<float, N> out;
+#pragma unroll
+    for (int i = 0; i < N; i++) {
+      out.data[i] = upcast_s(val.data[i]);
+    }
+    return out;
+  }
+}
+
+template <typename O>
+DINLINE O downcast(array_t<float, O::size> val) {
+  if constexpr (std::is_same<typename O::type, float>::value) {
+    return val;
+  } else {
+    O out;
+#pragma unroll
+    for (int i = 0; i < O::size; i++) {
+      out.data[i] = downcast_s<typename O::type>(val.data[i]);
+    }
+    return out;
+  }
+}
+
+#if !defined(USE_ROCM)
+
+static DINLINE void st_flag_release(FlagType* flag_addr, FlagType flag) {
+  #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
+  asm volatile("st.release.sys.global.u32 [%1], %0;" ::"r"(flag),
+               "l"(flag_addr));
+  #else
+  asm volatile("membar.sys; st.volatile.global.u32 [%1], %0;" ::"r"(flag),
+               "l"(flag_addr));
+  #endif
+}
+
+static DINLINE FlagType ld_flag_acquire(FlagType* flag_addr) {
+  FlagType flag;
+  #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
+  asm volatile("ld.acquire.sys.global.u32 %0, [%1];"
+               : "=r"(flag)
+               : "l"(flag_addr));
+  #else
+  asm volatile("ld.volatile.global.u32 %0, [%1]; membar.gl;"
+               : "=r"(flag)
+               : "l"(flag_addr));
+  #endif
+  return flag;
+}
+
+static DINLINE void st_flag_volatile(FlagType* flag_addr, FlagType flag) {
+  asm volatile("st.volatile.global.u32 [%1], %0;" ::"r"(flag), "l"(flag_addr));
+}
+
+static DINLINE FlagType ld_flag_volatile(FlagType* flag_addr) {
+  FlagType flag;
+  asm volatile("ld.volatile.global.u32 %0, [%1];"
+               : "=r"(flag)
+               : "l"(flag_addr));
+  return flag;
+}
+
+// This function is meant to be used as the first synchronization in the all
+// reduce kernel. Thus, it doesn't need to make any visibility guarantees for
+// prior memory accesses. Note: volatile writes will not be reordered against
+// other volatile writes.
+template <int ngpus>
+DINLINE void barrier_at_start(const RankSignals& sg, Signal* self_sg,
+                              int rank) {
+  uint32_t flag = self_sg->_flag[blockIdx.x] + 1;
+  if (threadIdx.x < ngpus) {
+    auto peer_counter_ptr = &sg.signals[threadIdx.x]->start[blockIdx.x][rank];
+    auto self_counter_ptr = &self_sg->start[blockIdx.x][threadIdx.x];
+    // Write the expected counter value to peer and wait for correct value
+    // from peer.
+    st_flag_volatile(peer_counter_ptr, flag);
+    while (ld_flag_volatile(self_counter_ptr) != flag);
+  }
+  __syncthreads();
+  // use one thread to update flag
+  if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag;
+}
+
+// This function is meant to be used as the second or the final
+// synchronization barrier in the all reduce kernel. If it's the final
+// synchronization barrier, we don't need to make any visibility guarantees
+// for prior memory accesses.
+template <int ngpus, bool final_sync = false>
+DINLINE void barrier_at_end(const RankSignals& sg, Signal* self_sg, int rank) {
+  __syncthreads();
+  uint32_t flag = self_sg->_flag[blockIdx.x] + 1;
+  if (threadIdx.x < ngpus) {
+    auto peer_counter_ptr = &sg.signals[threadIdx.x]->end[blockIdx.x][rank];
+    auto self_counter_ptr = &self_sg->end[blockIdx.x][threadIdx.x];
+    // Write the expected counter value to peer and wait for correct value from
+    // peer.
+    if constexpr (!final_sync) {
+      st_flag_release(peer_counter_ptr, flag);
+      while (ld_flag_acquire(self_counter_ptr) != flag);
+    } else {
+      st_flag_volatile(peer_counter_ptr, flag);
+      while (ld_flag_volatile(self_counter_ptr) != flag);
+    }
+  }
+  if constexpr (!final_sync) __syncthreads();
+
+  // use one thread to update flag
+  if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag;
+}
+
+#else
+
+template <int ngpus>
+DINLINE void barrier_at_start(const RankSignals& sg, Signal* self_sg,
+                              int rank) {
+  uint32_t flag = self_sg->_flag[blockIdx.x] + 1;
+  if (threadIdx.x < ngpus) {
+    // simultaneously write to the corresponding flag of all ranks.
+    // Latency = 1 p2p write
+    __scoped_atomic_store_n(&sg.signals[threadIdx.x]->start[blockIdx.x][rank],
+                            flag, __ATOMIC_RELAXED, __MEMORY_SCOPE_SYSTEM);
+    // wait until we got true from all ranks
+    while (__scoped_atomic_load_n(&self_sg->start[blockIdx.x][threadIdx.x],
+                                  __ATOMIC_RELAXED,
+                                  __MEMORY_SCOPE_DEVICE) < flag);
+  }
+  __syncthreads();
+  // use one thread to update flag
+  if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag;
+}
+
+template <int ngpus, bool final_sync = false>
+DINLINE void barrier_at_end(const RankSignals& sg, Signal* self_sg, int rank) {
+  __syncthreads();
+  uint32_t flag = self_sg->_flag[blockIdx.x] + 1;
+  if (threadIdx.x < ngpus) {
+    // simultaneously write to the corresponding flag of all ranks.
+    // Latency = 1 p2p write
+    __scoped_atomic_store_n(&sg.signals[threadIdx.x]->end[blockIdx.x][rank],
+                            flag,
+                            final_sync ? __ATOMIC_RELAXED : __ATOMIC_RELEASE,
+                            __MEMORY_SCOPE_SYSTEM);
+    // wait until we got true from all ranks
+    while (
+        __scoped_atomic_load_n(&self_sg->end[blockIdx.x][threadIdx.x],
+                               final_sync ? __ATOMIC_RELAXED : __ATOMIC_ACQUIRE,
+                               __MEMORY_SCOPE_DEVICE) < flag);
+  }
+  if constexpr (!final_sync) __syncthreads();
+  // use one thread to update flag
+  if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag;
+}
+
+#endif
+
+template <typename P, int ngpus, typename A>
+DINLINE P packed_reduce(const P* ptrs[], int idx) {
+  A tmp = upcast(ptrs[0][idx]);
+#pragma unroll
+  for (int i = 1; i < ngpus; i++) {
+    packed_assign_add(tmp, upcast(ptrs[i][idx]));
+  }
+  return downcast<P>(tmp);
+}

 template <typename T, int ngpus>
 __global__ void __launch_bounds__(512, 1)
@@ -325,21 +616,6 @@ class CustomAllreduce {
 #undef KL
  }

-  void allgather(cudaStream_t stream, void* input, void* output, int size_bytes,
-                 int threads = 512, int block_limit = defaultBlockLimit);
-  template <typename T>
-  void mnnvl_lamport_allgather(cudaStream_t stream, T* input, T* output,
-                               void* local_buffer, void* multicast_buffer,
-                               uint32_t* epochs, int size_bytes,
-                               int stage_size_bytes);
-  template <typename T>
-  void reduce_scatter(cudaStream_t stream, T* input, T* output, int size,
-                      int threads = 512, int block_limit = defaultBlockLimit);
-  template <typename T>
-  void mnnvl_lamport_reduce_scatter(cudaStream_t stream, T* input, T* output,
-                                    void* local_buffer, uint32_t* epochs,
-                                    int size, int stage_size_bytes);
-
  ~CustomAllreduce() {
    for (auto [_, ptr] : ipc_handles_) {
      CUDACHECK(cudaIpcCloseMemHandle(ptr));
@@ -349,8 +625,8 @@ class CustomAllreduce {

 /**
 * To inspect PTX/SASS, copy paste this header file to compiler explorer and
- * add a template instantiation:
+ add a template instantiation:
 * template void vllm::CustomAllreduce::allreduce<half>(cudaStream_t, half *,
- *                                                       half *, int, int, int);
- */
-}  // namespace vllm
+ half *, int, int, int);
+*/
+}  // namespace vllm
@@ -1,332 +0,0 @@
-#pragma once
-
-#include <cuda.h>
-#include <cuda_bf16.h>
-#include <cuda_fp16.h>
-#include <cuda_runtime.h>
-
-#if defined(USE_ROCM)
-typedef __hip_bfloat16 nv_bfloat16;
-#endif
-
-#include <iostream>
-#include <array>
-#include <limits>
-#include <map>
-#include <unordered_map>
-#include <vector>
-#include <cstdlib>
-#include <cstring>
-
-namespace vllm {
-constexpr int kMaxCustomCollectiveRanks = 16;
-
-#define CUDACHECK(cmd)                                              \
-  do {                                                              \
-    cudaError_t e = cmd;                                            \
-    if (e != cudaSuccess) {                                         \
-      printf("Failed: Cuda error %s:%d '%s'\n", __FILE__, __LINE__, \
-             cudaGetErrorString(e));                                \
-      exit(EXIT_FAILURE);                                           \
-    }                                                               \
-  } while (0)
-
-// Maximal number of blocks in allreduce kernel.
-constexpr int kMaxBlocks = 36;
-
-// Default number of blocks in allreduce kernel.
-#ifndef USE_ROCM
-inline constexpr int defaultBlockLimit = 36;
-inline CUpointer_attribute rangeStartAddrAttr =
-    CU_POINTER_ATTRIBUTE_RANGE_START_ADDR;
-#else
-inline constexpr int defaultBlockLimit = 16;
-inline hipPointer_attribute rangeStartAddrAttr =
-    HIP_POINTER_ATTRIBUTE_RANGE_START_ADDR;
-#endif
-
-// Counter may overflow, but it's fine since unsigned int overflow is
-// well-defined behavior.
-using FlagType = uint32_t;
-
-// Two sets of peer counters are needed for two syncs: starting and ending an
-// operation. The reason is that it's possible for peer GPU block to arrive at
-// the second sync point while the current GPU block haven't passed the first
-// sync point. Thus, peer GPU may write counter+1 while current GPU is busy
-// waiting for counter. We use alternating counter array to avoid this
-// possibility.
-struct Signal {
-  alignas(128) FlagType start[kMaxBlocks][kMaxCustomCollectiveRanks];
-  alignas(128) FlagType end[kMaxBlocks][kMaxCustomCollectiveRanks];
-  alignas(128) FlagType _flag[kMaxBlocks];  // incremental flags for each rank
-};
-
-struct __align__(16) RankData {
-  const void* ptrs[kMaxCustomCollectiveRanks];
-};
-
-struct __align__(16) RankSignals {
-  Signal* signals[kMaxCustomCollectiveRanks];
-};
-
-// like std::array, but aligned
-template <typename T, int sz>
-struct __align__(alignof(T) * sz) array_t {
-  T data[sz];
-  using type = T;
-  static constexpr int size = sz;
-};
-
-// use packed type to maximize memory efficiency
-// goal: generate ld.128 and st.128 instructions
-template <typename T>
-struct packed_t {
-  // the (P)acked type for load/store
-  using P = array_t<T, 16 / sizeof(T)>;
-  // the (A)ccumulator type for reduction
-  using A = array_t<float, 16 / sizeof(T)>;
-};
-
-#define DINLINE __device__ __forceinline__
-
-// scalar cast functions
-DINLINE float upcast_s(half val) { return __half2float(val); }
-
-template <typename T>
-DINLINE T downcast_s(float val);
-template <>
-DINLINE half downcast_s(float val) {
-  return __float2half(val);
-}
-
-// scalar add functions
-// for some reason when compiling with Pytorch, the + operator for half and
-// bfloat is disabled so we call the intrinsics directly
-DINLINE half& assign_add(half& a, half b) {
-  a = __hadd(a, b);
-  return a;
-}
-DINLINE float& assign_add(float& a, float b) { return a += b; }
-
-#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
-DINLINE float upcast_s(nv_bfloat16 val) { return __bfloat162float(val); }
-template <>
-DINLINE nv_bfloat16 downcast_s(float val) {
-  return __float2bfloat16(val);
-}
-DINLINE nv_bfloat16& assign_add(nv_bfloat16& a, nv_bfloat16 b) {
-  a = __hadd(a, b);
-  return a;
-}
-#endif
-
-template <typename T, int N>
-DINLINE array_t<T, N>& packed_assign_add(array_t<T, N>& a, array_t<T, N> b) {
-#pragma unroll
-  for (int i = 0; i < N; i++) {
-    assign_add(a.data[i], b.data[i]);
-  }
-  return a;
-}
-
-template <typename T, int N>
-DINLINE array_t<float, N> upcast(array_t<T, N> val) {
-  if constexpr (std::is_same<T, float>::value) {
-    return val;
-  } else {
-    array_t<float, N> out;
-#pragma unroll
-    for (int i = 0; i < N; i++) {
-      out.data[i] = upcast_s(val.data[i]);
-    }
-    return out;
-  }
-}
-
-template <typename O>
-DINLINE O downcast(array_t<float, O::size> val) {
-  if constexpr (std::is_same<typename O::type, float>::value) {
-    return val;
-  } else {
-    O out;
-#pragma unroll
-    for (int i = 0; i < O::size; i++) {
-      out.data[i] = downcast_s<typename O::type>(val.data[i]);
-    }
-    return out;
-  }
-}
-
-#if !defined(USE_ROCM)
-
-static DINLINE void st_flag_release(FlagType* flag_addr, FlagType flag) {
-  #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
-  asm volatile("st.release.sys.global.u32 [%1], %0;" ::"r"(flag),
-               "l"(flag_addr));
-  #else
-  asm volatile("membar.sys; st.volatile.global.u32 [%1], %0;" ::"r"(flag),
-               "l"(flag_addr));
-  #endif
-}
-
-static DINLINE FlagType ld_flag_acquire(FlagType* flag_addr) {
-  FlagType flag;
-  #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
-  asm volatile("ld.acquire.sys.global.u32 %0, [%1];"
-               : "=r"(flag)
-               : "l"(flag_addr));
-  #else
-  asm volatile("ld.volatile.global.u32 %0, [%1]; membar.gl;"
-               : "=r"(flag)
-               : "l"(flag_addr));
-  #endif
-  return flag;
-}
-
-static DINLINE void st_flag_volatile(FlagType* flag_addr, FlagType flag) {
-  asm volatile("st.volatile.global.u32 [%1], %0;" ::"r"(flag), "l"(flag_addr));
-}
-
-static DINLINE FlagType ld_flag_volatile(FlagType* flag_addr) {
-  FlagType flag;
-  asm volatile("ld.volatile.global.u32 %0, [%1];"
-               : "=r"(flag)
-               : "l"(flag_addr));
-  return flag;
-}
-
-// This function is meant to be used as the first synchronization in the all
-// reduce kernel. Thus, it doesn't need to make any visibility guarantees for
-// prior memory accesses. Note: volatile writes will not be reordered against
-// other volatile writes.
-template <int ngpus>
-DINLINE void barrier_at_start(const RankSignals& sg, Signal* self_sg,
-                              int rank) {
-  uint32_t flag = self_sg->_flag[blockIdx.x] + 1;
-  if (threadIdx.x < ngpus) {
-    auto peer_counter_ptr = &sg.signals[threadIdx.x]->start[blockIdx.x][rank];
-    auto self_counter_ptr = &self_sg->start[blockIdx.x][threadIdx.x];
-    // Write the expected counter value to peer and wait for correct value
-    // from peer.
-    st_flag_volatile(peer_counter_ptr, flag);
-    while (ld_flag_volatile(self_counter_ptr) != flag);
-  }
-  __syncthreads();
-  // use one thread to update flag
-  if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag;
-}
-
-template <int ngpus>
-DINLINE void barrier_at_start_release(const RankSignals& sg, Signal* self_sg,
-                                      int rank) {
-  __syncthreads();
-  uint32_t flag = self_sg->_flag[blockIdx.x] + 1;
-  if (threadIdx.x < ngpus) {
-    auto peer_counter_ptr = &sg.signals[threadIdx.x]->start[blockIdx.x][rank];
-    auto self_counter_ptr = &self_sg->start[blockIdx.x][threadIdx.x];
-    st_flag_release(peer_counter_ptr, flag);
-    while (ld_flag_acquire(self_counter_ptr) != flag);
-  }
-  __syncthreads();
-  if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag;
-}
-
-// This function is meant to be used as the second or the final
-// synchronization barrier in the all reduce kernel. If it's the final
-// synchronization barrier, we don't need to make any visibility guarantees
-// for prior memory accesses.
-template <int ngpus, bool final_sync = false>
-DINLINE void barrier_at_end(const RankSignals& sg, Signal* self_sg, int rank) {
-  __syncthreads();
-  uint32_t flag = self_sg->_flag[blockIdx.x] + 1;
-  if (threadIdx.x < ngpus) {
-    auto peer_counter_ptr = &sg.signals[threadIdx.x]->end[blockIdx.x][rank];
-    auto self_counter_ptr = &self_sg->end[blockIdx.x][threadIdx.x];
-    // Write the expected counter value to peer and wait for correct value from
-    // peer.
-    if constexpr (!final_sync) {
-      st_flag_release(peer_counter_ptr, flag);
-      while (ld_flag_acquire(self_counter_ptr) != flag);
-    } else {
-      st_flag_volatile(peer_counter_ptr, flag);
-      while (ld_flag_volatile(self_counter_ptr) != flag);
-    }
-  }
-  if constexpr (!final_sync) __syncthreads();
-
-  // use one thread to update flag
-  if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag;
-}
-
-#else
-
-template <int ngpus>
-DINLINE void barrier_at_start(const RankSignals& sg, Signal* self_sg,
-                              int rank) {
-  uint32_t flag = self_sg->_flag[blockIdx.x] + 1;
-  if (threadIdx.x < ngpus) {
-    // simultaneously write to the corresponding flag of all ranks.
-    // Latency = 1 p2p write
-    __scoped_atomic_store_n(&sg.signals[threadIdx.x]->start[blockIdx.x][rank],
-                            flag, __ATOMIC_RELAXED, __MEMORY_SCOPE_SYSTEM);
-    // wait until we got true from all ranks
-    while (__scoped_atomic_load_n(&self_sg->start[blockIdx.x][threadIdx.x],
-                                  __ATOMIC_RELAXED,
-                                  __MEMORY_SCOPE_DEVICE) < flag);
-  }
-  __syncthreads();
-  // use one thread to update flag
-  if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag;
-}
-
-template <int ngpus>
-DINLINE void barrier_at_start_release(const RankSignals& sg, Signal* self_sg,
-                                      int rank) {
-  __syncthreads();
-  uint32_t flag = self_sg->_flag[blockIdx.x] + 1;
-  if (threadIdx.x < ngpus) {
-    __scoped_atomic_store_n(&sg.signals[threadIdx.x]->start[blockIdx.x][rank],
-                            flag, __ATOMIC_RELEASE, __MEMORY_SCOPE_SYSTEM);
-    while (__scoped_atomic_load_n(&self_sg->start[blockIdx.x][threadIdx.x],
-                                  __ATOMIC_ACQUIRE,
-                                  __MEMORY_SCOPE_DEVICE) < flag);
-  }
-  __syncthreads();
-  if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag;
-}
-
-template <int ngpus, bool final_sync = false>
-DINLINE void barrier_at_end(const RankSignals& sg, Signal* self_sg, int rank) {
-  __syncthreads();
-  uint32_t flag = self_sg->_flag[blockIdx.x] + 1;
-  if (threadIdx.x < ngpus) {
-    // simultaneously write to the corresponding flag of all ranks.
-    // Latency = 1 p2p write
-    __scoped_atomic_store_n(&sg.signals[threadIdx.x]->end[blockIdx.x][rank],
-                            flag,
-                            final_sync ? __ATOMIC_RELAXED : __ATOMIC_RELEASE,
-                            __MEMORY_SCOPE_SYSTEM);
-    // wait until we got true from all ranks
-    while (
-        __scoped_atomic_load_n(&self_sg->end[blockIdx.x][threadIdx.x],
-                               final_sync ? __ATOMIC_RELAXED : __ATOMIC_ACQUIRE,
-                               __MEMORY_SCOPE_DEVICE) < flag);
-  }
-  if constexpr (!final_sync) __syncthreads();
-  // use one thread to update flag
-  if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag;
-}
-
-#endif
-
-template <typename P, int ngpus, typename A>
-DINLINE P packed_reduce(const P* ptrs[], int idx) {
-  A tmp = upcast(ptrs[0][idx]);
-#pragma unroll
-  for (int i = 1; i < ngpus; i++) {
-    packed_assign_add(tmp, upcast(ptrs[i][idx]));
-  }
-  return downcast<P>(tmp);
-}
-
-}  // namespace vllm
--- a/Show More
+++ b/Show More