[Build] Update pre-commit pip-compile hook to use cu128 torch backend

Matches the test.txt lockfile change to CUDA 12.8. Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
[Build] Recompile test.txt lockfile with cu128 torch backend
2026-02-18 16:41:53 -05:00 · 2026-02-18 14:38:23 -05:00 · 2026-02-18 14:38:23 -05:00
1261 changed files with 29539 additions and 102594 deletions
@@ -10,7 +10,7 @@ steps:
      docker build
      --build-arg max_jobs=16
      --build-arg REMOTE_VLLM=1
-      --build-arg ARG_PYTORCH_ROCM_ARCH='gfx942;gfx950'
+      --build-arg ARG_PYTORCH_ROCM_ARCH='gfx90a;gfx942'
      --build-arg VLLM_BRANCH=$BUILDKITE_COMMIT
      --tag "rocm/vllm-ci:${BUILDKITE_COMMIT}"
      -f docker/Dockerfile.rocm
@@ -14,7 +14,7 @@ BUILDKITE_COMMIT=$3
 aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY"

 # skip build if image already exists
-if [[ -z $(docker manifest inspect "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-arm64-cpu) ]]; then
+if [[ -z $(docker manifest inspect "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-cpu) ]]; then
  echo "Image not found, proceeding with build..."
 else
  echo "Image found"
@@ -25,9 +25,9 @@ fi
 docker build --file docker/Dockerfile.cpu \
  --build-arg max_jobs=16 \
  --build-arg buildkite_commit="$BUILDKITE_COMMIT" \
-  --tag "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-arm64-cpu \
+  --tag "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-cpu \
  --target vllm-test \
  --progress plain .

 # push
-docker push "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-arm64-cpu
+docker push "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-cpu
@@ -51,56 +51,5 @@
            "max-model-len": 256,
            "async-scheduling": ""
        }
-    },
-    {
-        "test_name": "latency_deepseek_r1",
-        "environment_variables": {
-            "PT_HPU_LAZY_MODE": 1,
-            "PT_HPU_ENABLE_LAZY_COLLECTIVES": 1,
-            "VLLM_CONTIGUOUS_PA": 1,
-            "VLLM_DEFRAG": 1
-        },
-        "parameters": {
-            "model": "deepseek-ai/DeepSeek-R1",
-            "tensor_parallel_size": 8,
-            "load_format": "dummy",
-            "max-model-len": 2048,
-            "dtype": "bfloat16"
-        }
-    },
-    {
-        "test_name": "latency_llama4_maverick_17b128e_instruct_fp8",
-        "environment_variables": {
-            "PT_HPU_LAZY_MODE": 1,
-            "PT_HPU_ENABLE_LAZY_COLLECTIVES": 1,
-            "VLLM_CONTIGUOUS_PA": 1,
-            "VLLM_DEFRAG": 1
-        },
-        "parameters": {
-            "model": "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8",
-            "tensor_parallel_size": 8,
-            "max-model-len": 512,
-            "max-num-seqs": 128,
-            "async-scheduling": "",
-            "gpu-memory-utilization": 0.95,
-            "enable_expert_parallel": ""
-        }
-    },
-    {
-        "test_name": "latency_qwen3_8b",
-        "environment_variables": {
-            "PT_HPU_LAZY_MODE": 1,
-            "PT_HPU_ENABLE_LAZY_COLLECTIVES": 1,
-            "VLLM_CONTIGUOUS_PA": 1,
-            "VLLM_DEFRAG": 1
-        },
-        "parameters": {
-            "model": "Qwen/Qwen3-8B",
-            "tensor_parallel_size": 1,
-            "max-model-len": 2048,
-            "max-num-seqs": 128,
-            "dtype": "bfloat16",
-            "async-scheduling": ""
-        }
    }
 ]
@@ -78,84 +78,5 @@
            "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json",
            "num_prompts": 200
        }
-    },
-    {
-        "test_name": "serving_deepseek_r1",
-        "qps_list": [1, 4, 16, "inf"],
-        "server_environment_variables": {
-            "PT_HPU_LAZY_MODE": 1,
-            "PT_HPU_ENABLE_LAZY_COLLECTIVES": 1,
-            "VLLM_CONTIGUOUS_PA": 1,
-            "VLLM_DEFRAG": 1
-        },
-        "server_parameters": {
-            "model": "deepseek-ai/DeepSeek-R1",
-            "tensor_parallel_size": 8,
-            "swap_space": 16,
-            "disable_log_stats": "",
-            "load_format": "dummy",
-            "max-model-len": 2048,
-            "max-num-seqs": 200,
-            "async-scheduling": "",
-            "dtype": "bfloat16"
-        },
-        "client_parameters": {
-            "model": "deepseek-ai/DeepSeek-R1",
-            "backend": "vllm",
-            "dataset_name": "sharegpt",
-            "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json",
-            "num_prompts": 200
-        }
-    },
-    {
-        "test_name": "serving_llama4_maverick_17b128e_instruct_fp8",
-        "qps_list": [1, 4, 16, "inf"],
-        "server_environment_variables": {
-            "PT_HPU_LAZY_MODE": 1,
-            "PT_HPU_ENABLE_LAZY_COLLECTIVES": 1,
-            "VLLM_CONTIGUOUS_PA": 1,
-            "VLLM_DEFRAG": 1
-        },
-        "server_parameters": {
-            "model": "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8",
-            "tensor_parallel_size": 8,
-            "disable_log_stats": "",
-            "max-model-len": 2048,
-            "max-num-seqs": 128,
-            "async-scheduling": "",
-            "enable_expert_parallel": "",
-            "max-num-batched-tokens": 4096
-        },
-        "client_parameters": {
-            "model": "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8",
-            "backend": "vllm",
-            "dataset_name": "sharegpt",
-            "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json",
-            "num_prompts": 200
-        }
-    },
-    {
-        "test_name": "serving_qwen3_8b",
-        "qps_list": [1, 4, 10, "inf"],
-        "server_environment_variables": {
-            "PT_HPU_LAZY_MODE": 1,
-            "PT_HPU_ENABLE_LAZY_COLLECTIVES": 1,
-            "VLLM_CONTIGUOUS_PA": 1,
-            "VLLM_DEFRAG": 1
-        },
-        "server_parameters": {
-            "model": "Qwen/Qwen-3-8B",
-            "tensor_parallel_size": 1,
-            "dtype": "bfloat16",
-            "disable_log_stats": "",
-            "async-scheduling": ""
-        },
-        "client_parameters": {
-            "model": "Qwen/Qwen-3-8B",
-            "backend": "vllm",
-            "dataset_name": "sharegpt",
-            "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json",
-            "num_prompts": 200
-        }
    }
 ]
@@ -57,67 +57,5 @@
            "max-num-seqs": 512,
            "async-scheduling": ""
        }
-    },
-    {
-        "test_name": "throughput_deepseek_r1",
-        "environment_variables": {
-            "PT_HPU_LAZY_MODE": 1,
-            "PT_HPU_ENABLE_LAZY_COLLECTIVES": 1,
-            "VLLM_CONTIGUOUS_PA": 1,
-            "VLLM_DEFRAG": 1
-        },
-        "parameters": {
-            "model": "deepseek-ai/DeepSeek-R1",
-            "tensor_parallel_size": 8,
-            "load_format": "dummy",
-            "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json",
-            "dataset_name": "sharegpt",
-            "num_prompts": 1000,
-            "backend": "vllm",
-            "max-model-len": 2048,
-            "max-num-seqs": 384,
-            "async-scheduling": ""
-        }
-    },
-    {
-        "test_name": "throughput_llama4_maverick_17b128e_instruct_fp8",
-        "environment_variables": {
-            "PT_HPU_LAZY_MODE": 1,
-            "PT_HPU_ENABLE_LAZY_COLLECTIVES": 1,
-            "VLLM_CONTIGUOUS_PA": 1,
-            "VLLM_DEFRAG": 1
-        },
-        "parameters": {
-            "model": "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8",
-            "tensor_parallel_size": 8,
-            "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json",
-            "dataset_name": "sharegpt",
-            "num_prompts": 1000,
-            "backend": "vllm",
-            "max-model-len": 2048,
-            "max-num-seqs": 512,
-            "async-scheduling": "",
-            "enable_expert_parallel": ""
-        }
-    },
-    {
-        "test_name": "throughput_qwen3_8b",
-        "environment_variables": {
-            "PT_HPU_LAZY_MODE": 1,
-            "PT_HPU_ENABLE_LAZY_COLLECTIVES": 1,
-            "VLLM_CONTIGUOUS_PA": 1,
-            "VLLM_DEFRAG": 1
-        },
-        "parameters": {
-            "model": "Qwen/Qwen-3-8B",
-            "tensor_parallel_size": 1,
-            "load_format": "dummy",
-            "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json",
-            "dataset_name": "sharegpt",
-            "num_prompts": 1000,
-            "max-num-seqs": 512,
-            "backend": "vllm",
-            "async-scheduling": ""
-        }
    }
 ]
@@ -31,7 +31,7 @@ steps:
        commands:
          # #NOTE: torch_cuda_arch_list is derived from upstream PyTorch build files here:
          # https://github.com/pytorch/pytorch/blob/main/.ci/aarch64_linux/aarch64_ci_build.sh#L7
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04  --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg BUILDER_CUDA_VERSION=13.0.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04  --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
          - "mkdir artifacts"
          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
          - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35"
@@ -70,7 +70,7 @@ steps:
        agents:
          queue: cpu_queue_postmerge
        commands:
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg BUILDER_CUDA_VERSION=13.0.1 --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
          - "mkdir artifacts"
          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
          - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35"
@@ -123,7 +123,7 @@ steps:
          queue: cpu_queue_postmerge
        commands:
          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 --target vllm-openai --progress plain -f docker/Dockerfile ."
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg BUILDER_CUDA_VERSION=13.0.1 --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 --target vllm-openai --progress plain -f docker/Dockerfile ."
          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130"
          # re-tag to default image tag and push, just in case arm64 build fails
          - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130"
@@ -137,7 +137,7 @@ steps:
        commands:
          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
          # compute capability 12.0 for RTX-50 series / RTX PRO 6000 Blackwell, 12.1 for DGX Spark
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0 12.1' --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 --target vllm-openai --progress plain -f docker/Dockerfile ."
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg BUILDER_CUDA_VERSION=13.0.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0 12.1' --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 --target vllm-openai --progress plain -f docker/Dockerfile ."
          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130"

      - block: "Build release image for x86_64 CPU"
@@ -68,7 +68,7 @@ aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/triton
 aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/torchvision-*.whl .
 aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/torchaudio-*.whl .
 aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/amdsmi-*.whl .
-aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/amd_aiter-*.whl .
+aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/aiter-*.whl .
 aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/flash-attn-*.whl .
 \`\`\`

@@ -80,7 +80,7 @@ aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/flash-
 - **torchvision**: TorchVision for ROCm PyTorch
 - **torchaudio**: Torchaudio for ROCm PyTorch
 - **amdsmi**: AMD SMI Python bindings
- **amd_aiter**: Aiter for ROCm
+- **aiter**: Aiter for ROCm
 - **flash-attn**: Flash Attention for ROCm

 ### :warning: Notes
@@ -1,213 +0,0 @@
-#!/bin/bash
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-#
-# Check if Ray LLM can generate lock files that are compatible with this
-# version of vllm. Downloads Ray's requirement files and runs a full
-# dependency resolution with the installed vllm's constraints to see if
-# a valid lock file can be produced.
-#
-# See: https://github.com/vllm-project/vllm/issues/33599
-
-set -eo pipefail
-
-RAY_BASE_URL="https://raw.githubusercontent.com/ray-project/ray/master/python"
-
-WORK_DIR=$(mktemp -d)
-trap 'rm -rf "$WORK_DIR"' EXIT
-
-# Fetch all Ray requirement files used in the LLM depset pipeline
-echo ">>> Fetching Ray requirement files"
-RAY_FILES=(
-    "requirements.txt"
-    "requirements/cloud-requirements.txt"
-    "requirements/base-test-requirements.txt"
-    "requirements/llm/llm-requirements.txt"
-    "requirements/llm/llm-test-requirements.txt"
-)
-for FILE in "${RAY_FILES[@]}"; do
-    LOCAL_PATH="${WORK_DIR}/$(basename "$FILE")"
-    echo "    ${FILE}"
-    curl -fsSL -o "$LOCAL_PATH" "${RAY_BASE_URL}/${FILE}"
-done
-
-# Extract installed vllm deps
-echo ">>> Extracting installed vllm dependency constraints"
-python3 - "${WORK_DIR}/vllm-constraints.txt" <<'PYEOF'
-"""Write out the installed vllm's dependencies as pip constraint lines.
-
-Ray uses vllm[audio], so audio-extra deps are included with their extra
-markers stripped. The resolver cannot evaluate extra markers for a
-package that is not itself being resolved from an index, so we activate
-them manually here.
-"""
-import importlib.metadata
-import re
-import sys
-
-out_path = sys.argv[1]
-raw_reqs = importlib.metadata.requires("vllm") or []
-
-# Ray uses vllm[audio] – activate that extra.
-ACTIVE_EXTRAS = {"audio"}
-EXTRA_RE = re.compile(r"""extra\s*==\s*['"]([^'"]+)['"]""")
-
-lines = []
-for r in raw_reqs:
-    if ";" not in r:
-        # Unconditional dep — always include.
-        lines.append(r.strip())
-        continue
-
-    req_part, _, marker_part = r.partition(";")
-    marker_part = marker_part.strip()
-
-    extra_matches = EXTRA_RE.findall(marker_part)
-    if not extra_matches:
-        # Non-extra marker (python_version, etc.) — keep as-is.
-        lines.append(r.strip())
-        continue
-
-    if not ACTIVE_EXTRAS.intersection(extra_matches):
-        continue  # Skip inactive extras (tensorizer, bench, …).
-
-    # Strip the extra== conditions but keep any remaining markers
-    # (e.g. python_version).
-    cleaned = EXTRA_RE.sub("", marker_part)
-    cleaned = re.sub(r"\band\b\s*\band\b", "and", cleaned)
-    cleaned = re.sub(r"^\s*and\s+|\s+and\s*$", "", cleaned).strip()
-
-    if cleaned:
-        lines.append(f"{req_part.strip()} ; {cleaned}")
-    else:
-        lines.append(req_part.strip())
-
-with open(out_path, "w") as f:
-    for line in lines:
-        f.write(line + "\n")
-
-print(f"Wrote {len(lines)} constraints to {out_path}")
-PYEOF
-
-echo ">>> Installed vllm deps (first 20 lines):"
-head -20 "${WORK_DIR}/vllm-constraints.txt"
-
-# Remove Ray's vllm pin — the installed vllm's transitive deps
-# (written above) replace it in the resolution. vllm itself cannot
-# be resolved from PyPI for in-development versions, so we test
-# whether Ray's requirements can coexist with vllm's dependency
-# constraints instead.
-sed -i '/^vllm/d' "${WORK_DIR}/llm-requirements.txt"
-
-# Install uv if needed
-if ! command -v uv &>/dev/null; then
-    echo ">>> Installing uv"
-    pip install uv -q
-fi
-
-# Resolve: given vllm's constraints, can Ray compile a lock file?
-#
-# vllm's dependency constraints are the fixed side — Ray is flexible and
-# can regenerate its lock files. We pass vllm's constraints via -c so
-# the resolver treats them as non-negotiable bounds, then check whether
-# Ray's own requirements can still be satisfied within those bounds.
-echo ""
-echo "============================================================"
-echo ">>> Resolving: Can Ray generate compatible lock files?"
-echo "============================================================"
-
-set +e
-uv pip compile \
-    "${WORK_DIR}/requirements.txt" \
-    "${WORK_DIR}/cloud-requirements.txt" \
-    "${WORK_DIR}/base-test-requirements.txt" \
-    "${WORK_DIR}/llm-requirements.txt" \
-    "${WORK_DIR}/llm-test-requirements.txt" \
-    -c "${WORK_DIR}/vllm-constraints.txt" \
-    --python-version 3.12 \
-    --python-platform x86_64-manylinux_2_31 \
-    --extra-index-url https://download.pytorch.org/whl/cu129 \
-    --index-strategy unsafe-best-match \
-    --unsafe-package setuptools \
-    --unsafe-package ray \
-    --no-header \
-    -o "${WORK_DIR}/resolved.txt" \
-    2>&1
-EXIT_CODE=$?
-set -e
-
-echo ""
-echo "=========================================="
-if [ $EXIT_CODE -eq 0 ]; then
-    echo "SUCCESS: Ray can generate lock files compatible with this vllm."
-    echo ""
-    echo "Key resolved versions:"
-    grep -E '^(protobuf|torch|numpy|transformers)==' \
-        "${WORK_DIR}/resolved.txt" | sort || true
-    echo "=========================================="
-    exit 0
-fi
-
-echo "FAILURE: Ray cannot generate lock files compatible with this vllm."
-echo "This means a fundamental dependency conflict exists that Ray"
-echo "cannot resolve by regenerating its lock files."
-echo "See: https://github.com/vllm-project/vllm/issues/33599"
-echo "=========================================="
-
-# Buildkite annotation
-if [ -f /usr/bin/buildkite-agent ]; then
-    buildkite-agent annotate --style 'warning' --context 'ray-compat' << EOF
-### :warning: Ray Dependency Compatibility Warning
-This PR introduces dependencies that **cannot** be resolved with Ray's requirements.
-Ray would not be able to regenerate its lock files to accommodate this vllm version.
-
-Please check the **Ray Dependency Compatibility Check** step logs for details.
-See [issue #33599](https://github.com/vllm-project/vllm/issues/33599) for context.
-EOF
-fi
-
-# Notify Slack if webhook is configured and PR/branch are valid.
-if [ -n "$RAY_COMPAT_SLACK_WEBHOOK_URL" ]; then
-    PR="${BUILDKITE_PULL_REQUEST:-}"
-    BRANCH="${BUILDKITE_BRANCH:-}"
-
-    # Skip notification if PR is invalid or branch is empty
-    if [[ "$PR" = "false" || -z "$PR" || -z "$BRANCH" ]]; then
-        echo ">>> Skipping Slack notification (invalid PR or empty branch: PR=$PR, branch=$BRANCH)"
-    else
-        echo ">>> Sending Slack notification"
-        # Single quotes are intentional: the f-string expressions are Python, not shell.
-        # shellcheck disable=SC2016
-        PAYLOAD=$(python3 -c '
-import json, os, sys
-pr = os.getenv("BUILDKITE_PULL_REQUEST", "N/A")
-branch = os.getenv("BUILDKITE_BRANCH", "unknown")
-url = os.getenv("BUILDKITE_BUILD_URL", "#")
-data = {
-    "text": ":warning: Ray Dependency Compatibility Check Failed",
-    "blocks": [{
-        "type": "section",
-        "text": {
-            "type": "mrkdwn",
-            "text": (
-                "*:warning: Ray Dependency Compatibility Check Failed*\n"
-                f"PR #{pr} on branch `{branch}` introduces dependencies "
-                f"that cannot be resolved with Ray'\''s requirements.\n"
-                f"<{url}|View Build>"
-            ),
-        },
-    }],
-}
-print(json.dumps(data))
-')
-
-        HTTP_CODE=$(curl -s -o /dev/null -w "%{http_code}" -X POST "$RAY_COMPAT_SLACK_WEBHOOK_URL" \
-            -H 'Content-type: application/json' \
-            -d "$PAYLOAD")
-        echo "    Slack webhook response: $HTTP_CODE"
-    fi
-else
-    echo ">>> Skipping Slack notification (RAY_COMPAT_SLACK_WEBHOOK_URL not set)"
-fi
-
-exit 1
@@ -1,57 +1,25 @@
 #!/bin/bash

-# 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
-# "[node0_cmds] && [node1_cmds]" or via the NUM_NODES environment variable.
-#
-###############################################################################
-# QUOTING / COMMAND PASSING
-#
-# Passing commands as positional arguments ($*) is fragile when the command
-# string itself contains double quotes, e.g.:
-#
-#   bash run-amd-test.sh "export FLAGS="value" && pytest -m "not slow""
-#
-# The outer shell resolves the nested quotes *before* this script runs, so
-# the script receives mangled input it cannot fully recover.
-#
-# Preferred: pass commands via the VLLM_TEST_COMMANDS environment variable:
-#
-#   export VLLM_TEST_COMMANDS='export FLAGS="value" && pytest -m "not slow"'
-#   bash run-amd-test.sh
-#
-# Single-quoted assignment preserves all inner double quotes verbatim.
-# The $* path is kept for backward compatibility but callers should migrate.
-###############################################################################
+# This script runs test inside the corresponding ROCm docker container.
 set -o pipefail

 # Export Python path
 export PYTHONPATH=".."

-###############################################################################
-# Helper Functions
-###############################################################################
+# Print ROCm version
+echo "--- Confirming Clean Initial State"
+while true; do
+        sleep 3
+        if grep -q clean /opt/amdgpu/etc/gpu_state; then
+                echo "GPUs state is \"clean\""
+                break
+        fi
+done

-wait_for_clean_gpus() {
-  local timeout=${1:-300}
-  local start=$SECONDS
-  echo "--- Waiting for clean GPU state (timeout: ${timeout}s)"
-  while true; do
-    if grep -q clean /opt/amdgpu/etc/gpu_state; then
-      echo "GPUs state is \"clean\""
-      return
-    fi
-    if (( SECONDS - start >= timeout )); then
-      echo "Error: GPUs did not reach clean state within ${timeout}s" >&2
-      exit 1
-    fi
-    sleep 3
-  done
-}
+echo "--- ROCm info"
+rocminfo

+# cleanup older docker images
 cleanup_docker() {
  # Get Docker's root directory
  docker_root=$(docker info -f '{{.DockerRootDir}}')
@@ -60,12 +28,15 @@ cleanup_docker() {
    exit 1
  fi
  echo "Docker root directory: $docker_root"
-
+  # Check disk usage of the filesystem where Docker's root directory is located
  disk_usage=$(df "$docker_root" | tail -1 | awk '{print $5}' | sed 's/%//')
+  # Define the threshold
  threshold=70
  if [ "$disk_usage" -gt "$threshold" ]; then
    echo "Disk usage is above $threshold%. Cleaning up Docker images and volumes..."
+    # Remove dangling images (those that are not tagged and not used by any container)
    docker image prune -f
+    # Remove unused volumes / force the system prune for old images as well.
    docker volume prune -f && docker system prune --force --filter "until=72h" --all
    echo "Docker images and volumes cleanup completed."
  else
@@ -74,431 +45,193 @@ cleanup_docker() {
 }

 cleanup_network() {
-  local max_nodes=${NUM_NODES:-2}
-  for node in $(seq 0 $((max_nodes - 1))); do
-    if docker ps -a -q -f name="node${node}" | grep -q .; then
-      docker stop "node${node}" || true
+  for node in $(seq 0 $((NUM_NODES-1))); do
+    if docker pr -a -q -f name="node${node}" | grep -q .; then
+      docker stop "node${node}"
    fi
  done
-  if docker network ls | grep -q docker-net; then
-    docker network rm docker-net || true
+  if docker network ls | grep docker-net; then
+    docker network rm docker-net
  fi
 }

-is_multi_node() {
-  local cmds="$1"
-  # Primary signal: NUM_NODES environment variable set by the pipeline
-  if [[ "${NUM_NODES:-1}" -gt 1 ]]; then
-    return 0
-  fi
-  # Fallback: detect the bracket syntax structurally
-  # Pattern: [...] && [...] (per-node command arrays)
-  if [[ "$cmds" =~ \[.*\].*\&\&.*\[.*\] ]]; then
-    return 0
-  fi
-  return 1
-}
-
-handle_pytest_exit() {
-  local exit_code=$1
-  if [ "$exit_code" -eq 5 ]; then
-    echo "Pytest exit code 5 (no tests collected) - treating as success."
-    exit 0
-  fi
-  exit "$exit_code"
-}
-
-###############################################################################
-# Pytest marker/keyword re-quoting
-#
-# When commands are passed through Buildkite -> shell -> $* -> bash -c,
-# quotes around multi-word pytest -m/-k expressions get stripped:
-#   pytest -v -s -m 'not cpu_test' v1/core
-# becomes:
-#   pytest -v -s -m not cpu_test v1/core
-#
-# pytest then interprets "cpu_test" as a file path, not part of the marker.
-#
-# This function detects unquoted expressions after -m/-k and re-quotes them
-# by collecting tokens until a recognizable boundary is reached:
-#   - test path (contains '/')
-#   - test file (ends with '.py')
-#   - another pytest flag (--xxx or -x single-char flags)
-#   - command separator (&& || ; |)
-#   - environment variable assignment (FOO=bar)
-#
-# Single-word markers (e.g. -m cpu_test, -m hybrid_model) pass through
-# unquoted since they have no spaces and work fine.
-#
-# Already-quoted expressions (containing literal single quotes) are passed
-# through untouched to avoid double-quoting values injected by
-# apply_rocm_test_overrides.
-#
-# NOTE: This ONLY fixes -m/-k flags. It cannot recover arbitrary inner
-# double-quotes stripped by the calling shell (see header comment).
-# Use VLLM_TEST_COMMANDS to avoid the problem entirely.
-###############################################################################
-re_quote_pytest_markers() {
-  local input="$1"
-  local output=""
-  local collecting=false
-  local marker_buf=""
-
-  # Strip backslash-newline continuations, then flatten remaining newlines
-  local flat="${input//$'\\\n'/ }"
-  flat="${flat//$'\n'/ }"
-
-  # Disable globbing to prevent *.py etc. from expanding during read -ra
-  local restore_glob
-  restore_glob="$(shopt -p -o noglob 2>/dev/null || true)"
-  set -o noglob
-  local -a words
-  read -ra words <<< "$flat"
-  eval "$restore_glob"
-
-  for word in "${words[@]}"; do
-    if $collecting; then
-      # If the token we're about to collect already contains a literal
-      # single quote, the expression was already quoted upstream.
-      # Flush and stop collecting.
-      if [[ "$word" == *"'"* ]]; then
-        if [[ -n "$marker_buf" ]]; then
-          # Should not normally happen (partial buf + quote), flush raw
-          output+="${marker_buf} "
-          marker_buf=""
-        fi
-        output+="${word} "
-        collecting=false
-        continue
-      fi
-
-      local is_boundary=false
-      case "$word" in
-        # Line-continuation artifact
-        "\\")
-          is_boundary=true ;;
-        # Command separators
-        "&&"|"||"|";"|"|")
-          is_boundary=true ;;
-        # Long flags (--ignore, --shard-id, etc.)
-        --*)
-          is_boundary=true ;;
-        # Short flags (-v, -s, -x, etc.) but NOT negative marker tokens
-        # like "not" which don't start with "-". Also skip -k/-m which
-        # would start a new marker (handled below).
-        -[a-zA-Z])
-          is_boundary=true ;;
-        # Test path (contains /)
-        */*)
-          is_boundary=true ;;
-        # Test file (ends with .py, possibly with ::method)
-        *.py|*.py::*)
-          is_boundary=true ;;
-        # Environment variable assignment preceding a command (FOO=bar)
-        *=*)
-          # Only treat as boundary if it looks like VAR=value, not
-          # pytest filter expressions like num_gpus=2 inside markers
-          if [[ "$word" =~ ^[A-Z_][A-Z0-9_]*= ]]; then
-            is_boundary=true
-          fi
-          ;;
-      esac
-
-      if $is_boundary; then
-        # Flush the collected marker expression
-        if [[ "$marker_buf" == *" "* || "$marker_buf" == *"("* ]]; then
-          output+="'${marker_buf}' "
-        else
-          output+="${marker_buf} "
-        fi
-        collecting=false
-        marker_buf=""
-        # Check if this boundary word itself starts a new -m/-k
-        if [[ "$word" == "-m" || "$word" == "-k" ]]; then
-          output+="${word} "
-          collecting=true
-        # Drop stray backslash tokens silently
-        elif [[ "$word" == "\\" ]]; then
-          :
-        else
-          output+="${word} "
-        fi
-      else
-        # Accumulate into marker buffer
-        if [[ -n "$marker_buf" ]]; then
-          marker_buf+=" ${word}"
-        else
-          marker_buf="${word}"
-        fi
-      fi
-    elif [[ "$word" == "-m" || "$word" == "-k" ]]; then
-      output+="${word} "
-      collecting=true
-      marker_buf=""
-    else
-      output+="${word} "
-    fi
-  done
-
-  # Flush any trailing marker expression (marker at end of command)
-  if $collecting && [[ -n "$marker_buf" ]]; then
-    if [[ "$marker_buf" == *" "* || "$marker_buf" == *"("* ]]; then
-      output+="'${marker_buf}'"
-    else
-      output+="${marker_buf}"
-    fi
-  fi
-
-  echo "${output% }"
-}
-
-###############################################################################
-# ROCm-specific pytest command rewrites
-#
-# These apply ignore flags and environment overrides for tests that are not
-# yet supported or behave differently on ROCm hardware. Kept as a single
-# function so new exclusions are easy to add in one place.
-###############################################################################
-
-apply_rocm_test_overrides() {
-  local cmds="$1"
-
-  # --- Model registry filter ---
-  if [[ $cmds == *"pytest -v -s models/test_registry.py"* ]]; then
-    cmds=${cmds//"pytest -v -s models/test_registry.py"/"pytest -v -s models/test_registry.py -k 'not BambaForCausalLM and not GritLM and not Mamba2ForCausalLM and not Zamba2ForCausalLM'"}
-  fi
-
-  # --- LoRA: disable custom paged attention ---
-  if [[ $cmds == *"pytest -v -s lora"* ]]; then
-    cmds=${cmds//"pytest -v -s lora"/"VLLM_ROCM_CUSTOM_PAGED_ATTN=0 pytest -v -s lora"}
-  fi
-
-  # --- Kernel ignores ---
-  if [[ $cmds == *" kernels/core"* ]]; then
-    cmds="${cmds} \
-    --ignore=kernels/core/test_fused_quant_layernorm.py \
-    --ignore=kernels/core/test_permute_cols.py"
-  fi
-
-  if [[ $cmds == *" kernels/attention"* ]]; then
-    cmds="${cmds} \
-    --ignore=kernels/attention/test_attention_selector.py \
-    --ignore=kernels/attention/test_encoder_decoder_attn.py \
-    --ignore=kernels/attention/test_flash_attn.py \
-    --ignore=kernels/attention/test_flashinfer.py \
-    --ignore=kernels/attention/test_prefix_prefill.py \
-    --ignore=kernels/attention/test_cascade_flash_attn.py \
-    --ignore=kernels/attention/test_mha_attn.py \
-    --ignore=kernels/attention/test_lightning_attn.py \
-    --ignore=kernels/attention/test_attention.py"
-  fi
-
-  if [[ $cmds == *" kernels/quantization"* ]]; then
-    cmds="${cmds} \
-    --ignore=kernels/quantization/test_int8_quant.py \
-    --ignore=kernels/quantization/test_machete_mm.py \
-    --ignore=kernels/quantization/test_block_fp8.py \
-    --ignore=kernels/quantization/test_block_int8.py \
-    --ignore=kernels/quantization/test_marlin_gemm.py \
-    --ignore=kernels/quantization/test_cutlass_scaled_mm.py \
-    --ignore=kernels/quantization/test_int8_kernel.py"
-  fi
-
-  if [[ $cmds == *" kernels/mamba"* ]]; then
-    cmds="${cmds} \
-    --ignore=kernels/mamba/test_mamba_mixer2.py \
-    --ignore=kernels/mamba/test_causal_conv1d.py \
-    --ignore=kernels/mamba/test_mamba_ssm_ssd.py"
-  fi
-
-  if [[ $cmds == *" kernels/moe"* ]]; then
-    cmds="${cmds} \
-    --ignore=kernels/moe/test_moe.py \
-    --ignore=kernels/moe/test_cutlass_moe.py \
-    --ignore=kernels/moe/test_triton_moe_ptpc_fp8.py"
-  fi
-
-  # --- Entrypoint ignores ---
-  if [[ $cmds == *" entrypoints/openai "* ]]; then
-    cmds=${cmds//" entrypoints/openai "/" entrypoints/openai \
-    --ignore=entrypoints/openai/test_audio.py \
-    --ignore=entrypoints/openai/test_shutdown.py \
-    --ignore=entrypoints/openai/test_completion.py \
-    --ignore=entrypoints/openai/test_models.py \
-    --ignore=entrypoints/openai/test_lora_adapters.py \
-    --ignore=entrypoints/openai/test_return_tokens_as_ids.py \
-    --ignore=entrypoints/openai/test_root_path.py \
-    --ignore=entrypoints/openai/test_tokenization.py \
-    --ignore=entrypoints/openai/test_prompt_validation.py "}
-  fi
-
-  if [[ $cmds == *" entrypoints/llm "* ]]; then
-    cmds=${cmds//" entrypoints/llm "/" entrypoints/llm \
-    --ignore=entrypoints/llm/test_chat.py \
-    --ignore=entrypoints/llm/test_accuracy.py \
-    --ignore=entrypoints/llm/test_init.py \
-    --ignore=entrypoints/llm/test_prompt_validation.py "}
-  fi
-
-  # Clean up escaped newlines from --ignore appends
-  cmds=$(echo "$cmds" | sed 's/ \\ / /g')
-
-  echo "$cmds"
-}
-
-###############################################################################
-# Main
-###############################################################################
-
-# --- GPU initialization ---
-echo "--- Confirming Clean Initial State"
-wait_for_clean_gpus
-
-echo "--- ROCm info"
-rocminfo
-
-# --- Docker housekeeping ---
+# Call the cleanup docker function
 cleanup_docker

 echo "--- Resetting GPUs"
-echo "reset" > /opt/amdgpu/etc/gpu_state
-wait_for_clean_gpus

-# --- Pull test image ---
+echo "reset" > /opt/amdgpu/etc/gpu_state
+
+while true; do
+        sleep 3
+        if grep -q clean /opt/amdgpu/etc/gpu_state; then
+                echo "GPUs state is \"clean\""
+                break
+        fi
+done
+
 echo "--- Pulling container"
 image_name="rocm/vllm-ci:${BUILDKITE_COMMIT}"
 container_name="rocm_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | head -c 10; echo)"
 docker pull "${image_name}"

 remove_docker_container() {
-  docker rm -f "${container_name}" || docker image rm -f "${image_name}" || true
+   docker rm -f "${container_name}" || docker image rm -f "${image_name}" || true
 }
 trap remove_docker_container EXIT

-# --- Prepare commands ---
 echo "--- Running container"

 HF_CACHE="$(realpath ~)/huggingface"
 mkdir -p "${HF_CACHE}"
 HF_MOUNT="/root/.cache/huggingface"

-# ---- Command source selection ----
-# Prefer VLLM_TEST_COMMANDS (preserves all inner quoting intact).
-# Fall back to $* for backward compatibility, but warn that inner
-# double-quotes will have been stripped by the calling shell.
-if [[ -n "${VLLM_TEST_COMMANDS:-}" ]]; then
-  commands="${VLLM_TEST_COMMANDS}"
-  echo "Commands sourced from VLLM_TEST_COMMANDS (quoting preserved)"
-else
-  commands="$*"
-  if [[ -z "$commands" ]]; then
-    echo "Error: No test commands provided." >&2
-    echo "Usage:" >&2
-    echo "  Preferred:  VLLM_TEST_COMMANDS='...' bash $0" >&2
-    echo "  Legacy:     bash $0 \"commands here\"" >&2
-    exit 1
-  fi
-  echo "Commands sourced from positional args (legacy mode)"
-  echo "WARNING: Inner double-quotes in the command string may have been"
-  echo "  stripped by the calling shell. If you see syntax errors, switch to:"
-  echo "  export VLLM_TEST_COMMANDS='your commands here'"
-  echo "  bash $0"
-fi
-
+commands=$@
 echo "Raw commands: $commands"

-# Fix quoting before ROCm overrides (so overrides see correct structure)
-commands=$(re_quote_pytest_markers "$commands")
-echo "After re-quoting: $commands"
+commands=${commands//"pytest -v -s basic_correctness/test_basic_correctness.py"/"pytest -v -s basic_correctness/test_basic_correctness.py"}

-commands=$(apply_rocm_test_overrides "$commands")
+if [[ $commands == *"pytest -v -s models/test_registry.py"* ]]; then
+  commands=${commands//"pytest -v -s models/test_registry.py"/"pytest -v -s models/test_registry.py -k 'not BambaForCausalLM and not GritLM and not Mamba2ForCausalLM and not Zamba2ForCausalLM'"}
+fi
+
+commands=${commands//"pytest -v -s compile/test_basic_correctness.py"/"pytest -v -s compile/test_basic_correctness.py"}
+
+if [[ $commands == *"pytest -v -s lora"* ]]; then
+  commands=${commands//"pytest -v -s lora"/"VLLM_ROCM_CUSTOM_PAGED_ATTN=0 pytest -v -s lora"}
+fi
+
+#ignore certain kernels tests
+if [[ $commands == *" kernels/core"* ]]; then
+  commands="${commands} \
+  --ignore=kernels/core/test_fused_quant_layernorm.py \
+  --ignore=kernels/core/test_permute_cols.py"
+fi
+
+if [[ $commands == *" kernels/attention"* ]]; then
+  commands="${commands} \
+  --ignore=kernels/attention/test_attention_selector.py \
+  --ignore=kernels/attention/test_encoder_decoder_attn.py \
+  --ignore=kernels/attention/test_flash_attn.py \
+  --ignore=kernels/attention/test_flashinfer.py \
+  --ignore=kernels/attention/test_prefix_prefill.py \
+  --ignore=kernels/attention/test_cascade_flash_attn.py \
+  --ignore=kernels/attention/test_mha_attn.py \
+  --ignore=kernels/attention/test_lightning_attn.py \
+  --ignore=kernels/attention/test_attention.py"
+fi
+
+if [[ $commands == *" kernels/quantization"* ]]; then
+  commands="${commands} \
+  --ignore=kernels/quantization/test_int8_quant.py \
+  --ignore=kernels/quantization/test_machete_mm.py \
+  --ignore=kernels/quantization/test_block_fp8.py \
+  --ignore=kernels/quantization/test_block_int8.py \
+  --ignore=kernels/quantization/test_marlin_gemm.py \
+  --ignore=kernels/quantization/test_cutlass_scaled_mm.py \
+  --ignore=kernels/quantization/test_int8_kernel.py"
+fi
+
+if [[ $commands == *" kernels/mamba"* ]]; then
+  commands="${commands} \
+  --ignore=kernels/mamba/test_mamba_mixer2.py \
+  --ignore=kernels/mamba/test_causal_conv1d.py \
+  --ignore=kernels/mamba/test_mamba_ssm_ssd.py"
+fi
+
+if [[ $commands == *" kernels/moe"* ]]; then
+  commands="${commands} \
+  --ignore=kernels/moe/test_moe.py \
+  --ignore=kernels/moe/test_cutlass_moe.py \
+  --ignore=kernels/moe/test_triton_moe_ptpc_fp8.py"
+fi
+
+#ignore certain Entrypoints/openai tests
+if [[ $commands == *" entrypoints/openai "* ]]; then
+  commands=${commands//" entrypoints/openai "/" entrypoints/openai \
+  --ignore=entrypoints/openai/test_audio.py \
+  --ignore=entrypoints/openai/test_shutdown.py \
+  --ignore=entrypoints/openai/test_completion.py \
+  --ignore=entrypoints/openai/test_models.py \
+  --ignore=entrypoints/openai/test_lora_adapters.py \
+  --ignore=entrypoints/openai/test_return_tokens_as_ids.py \
+  --ignore=entrypoints/openai/test_root_path.py \
+  --ignore=entrypoints/openai/test_tokenization.py \
+  --ignore=entrypoints/openai/test_prompt_validation.py "}
+fi
+
+#ignore certain Entrypoints/llm tests
+if [[ $commands == *" entrypoints/llm "* ]]; then
+  commands=${commands//" entrypoints/llm "/" entrypoints/llm \
+  --ignore=entrypoints/llm/test_chat.py \
+  --ignore=entrypoints/llm/test_accuracy.py \
+  --ignore=entrypoints/llm/test_init.py \
+  --ignore=entrypoints/llm/test_prompt_validation.py "}
+fi
+
+commands=$(echo "$commands" | sed 's/ \\ / /g')
 echo "Final commands: $commands"

+# --ignore=entrypoints/openai/test_encoder_decoder.py \
+# --ignore=entrypoints/openai/test_embedding.py \
+# --ignore=entrypoints/openai/test_oot_registration.py
+# --ignore=entrypoints/openai/test_accuracy.py \
+# --ignore=entrypoints/openai/test_models.py <= Fails on MI250 but passes on MI300 as of 2025-03-13
+
+
 MYPYTHONPATH=".."

-# Verify GPU access
+# Test that we're launching on the machine that has
+# proper access to GPUs
 render_gid=$(getent group render | cut -d: -f3)
 if [[ -z "$render_gid" ]]; then
  echo "Error: 'render' group not found. This is required for GPU access." >&2
  exit 1
 fi

-# --- RDMA device passthrough (conditional) ---
-# If the host has RDMA devices, pass them through so tests like
-# test_moriio_connector can access ibverbs. On hosts without RDMA
-# hardware the tests will gracefully skip via _rdma_available().
-RDMA_FLAGS=""
-if [ -d /dev/infiniband ]; then
-  echo "RDMA devices detected on host, enabling passthrough"
-  RDMA_FLAGS="--device /dev/infiniband --cap-add=IPC_LOCK"
-else
-  echo "No RDMA devices found on host, RDMA tests will be skipped"
-fi
+if [[ $commands == *"VLLM_TEST_GROUP_NAME=mi325_4-2-node-tests-4-gpus-in-total"* ]]; then

-# --- Route: multi-node vs single-node ---
-if is_multi_node "$commands"; then
-  echo "--- Multi-node job detected"
  export DCKR_VER=$(docker --version | sed 's/Docker version \(.*\), build .*/\1/')

-  # Parse the bracket syntax:  prefix ; [node0_cmds] && [node1_cmds]
-  #   BASH_REMATCH[1] = prefix (everything before first bracket)
-  #   BASH_REMATCH[2] = comma-separated node0 commands
-  #   BASH_REMATCH[3] = comma-separated node1 commands
-  if [[ "$commands" =~ ^(.*)\[(.*)"] && ["(.*)\]$ ]]; then
-    prefix=$(echo "${BASH_REMATCH[1]}" | sed 's/;//g')
-    echo "PREFIX: ${prefix}"
-
-    export composite_command="(command rocm-smi || true)"
-    saved_IFS=$IFS
-    IFS=','
-    read -ra node0 <<< "${BASH_REMATCH[2]}"
-    read -ra node1 <<< "${BASH_REMATCH[3]}"
-    IFS=$saved_IFS
-
-    if [[ ${#node0[@]} -ne ${#node1[@]} ]]; then
-      echo "Warning: node0 has ${#node0[@]} commands, node1 has ${#node1[@]}. They will be paired by index."
-    fi
-
-    for i in "${!node0[@]}"; do
-      command_node_0=$(echo "${node0[i]}" | sed 's/\"//g')
-      command_node_1=$(echo "${node1[i]}" | sed 's/\"//g')
-
-      step_cmd="./.buildkite/scripts/run-multi-node-test.sh /vllm-workspace/tests 2 2 ${image_name} '${command_node_0}' '${command_node_1}'"
-      echo "COMMANDS: ${step_cmd}"
-      composite_command="${composite_command} && ${step_cmd}"
-    done
-
-    /bin/bash -c "${composite_command}"
-    exit_code=$?
-    cleanup_network
-    handle_pytest_exit "$exit_code"
+  if [[ "$commands" =~ ^(.*)"["(.*)"] && ["(.*)"]"$ ]]; then
+      prefix=$( echo "${BASH_REMATCH[1]}" | sed 's/;//g')
+      echo "PREFIX: ${prefix}"
+      export composite_command="(command rocm-smi || true)"
+      myIFS=$IFS
+      IFS=','
+      read -ra node0 <<< ${BASH_REMATCH[2]}
+      read -ra node1 <<< ${BASH_REMATCH[3]}
+      IFS=$myIFS
+      for i in "${!node0[@]}";do 
+        command_node_0=$(echo ${node0[i]} | sed 's/\"//g')
+        command_node_1=$(echo ${node1[i]} | sed 's/\"//g')
+        
+        export commands="./.buildkite/scripts/run-multi-node-test.sh /vllm-workspace/tests 2 2 ${image_name} '${command_node_0}' '${command_node_1}'"
+        echo "COMMANDS: ${commands}"
+        composite_command=$(echo "${composite_command} && ${commands}")
+      done
+      /bin/bash -c "${composite_command}"
+      cleanup_network
  else
-    echo "Multi-node job detected but failed to parse bracket command syntax."
-    echo "Expected format: prefix ; [node0_cmd1, node0_cmd2] && [node1_cmd1, node1_cmd2]"
-    echo "Got: $commands"
-    cleanup_network
-    exit 111
+      echo "Failed to parse node commands! Exiting."
+      cleanup_network
+      exit 111
  fi
 else
-  echo "--- Single-node job"
  echo "Render devices: $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES"
  docker run \
-    --device /dev/kfd $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES \
-    $RDMA_FLAGS \
-    --network=host \
-    --shm-size=16gb \
-    --group-add "$render_gid" \
-    --rm \
-    -e HF_TOKEN \
-    -e AWS_ACCESS_KEY_ID \
-    -e AWS_SECRET_ACCESS_KEY \
-    -v "${HF_CACHE}:${HF_MOUNT}" \
-    -e "HF_HOME=${HF_MOUNT}" \
-    -e "PYTHONPATH=${MYPYTHONPATH}" \
-    --name "${container_name}" \
-    "${image_name}" \
-    /bin/bash -c "${commands}"
-
-  exit_code=$?
-  handle_pytest_exit "$exit_code"
+          --device /dev/kfd $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES \
+          --network=host \
+          --shm-size=16gb \
+          --group-add "$render_gid" \
+          --rm \
+          -e HF_TOKEN \
+          -e AWS_ACCESS_KEY_ID \
+          -e AWS_SECRET_ACCESS_KEY \
+          -v "${HF_CACHE}:${HF_MOUNT}" \
+          -e "HF_HOME=${HF_MOUNT}" \
+          -e "PYTHONPATH=${MYPYTHONPATH}" \
+          --name "${container_name}" \
+          "${image_name}" \
+          /bin/bash -c "${commands}"
 fi
@@ -1,43 +1,26 @@
 #!/bin/bash
 set -euox pipefail
-export VLLM_CPU_CI_ENV=0

 echo "--- PP+TP"
 vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -pp=2 &
 server_pid=$!
-timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1
+timeout 600 bash -c "until curl localhost:8000/v1/models; do sleep 1; done" || exit 1
 vllm bench serve \
    --backend vllm \
    --dataset-name random \
    --model meta-llama/Llama-3.2-3B-Instruct \
    --num-prompts 20 \
-    --result-dir ./test_results \
-    --result-filename tp_pp.json \
-    --save-result \
    --endpoint /v1/completions
-kill -s SIGTERM $server_pid; wait $server_pid || true
-failed_req=$(jq '.failed' ./test_results/tp_pp.json)
-if [ "$failed_req" -ne 0 ]; then
-  echo "Some requests were failed!"
-  exit 1
-fi
+kill -s SIGTERM $server_pid &

 echo "--- DP+TP"
 vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -dp=2 &
 server_pid=$!
-timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1
+timeout 600 bash -c "until curl localhost:8000/v1/models; do sleep 1; done" || exit 1
 vllm bench serve \
    --backend vllm \
    --dataset-name random \
    --model meta-llama/Llama-3.2-3B-Instruct \
    --num-prompts 20 \
-    --result-dir ./test_results \
-    --result-filename dp_pp.json \
-    --save-result \
    --endpoint /v1/completions
-kill -s SIGTERM $server_pid; wait $server_pid || true
-failed_req=$(jq '.failed' ./test_results/dp_pp.json)
-if [ "$failed_req" -ne 0 ]; then
-  echo "Some requests were failed!"
-  exit 1
-fi
+kill -s SIGTERM $server_pid &
@@ -1,27 +1,9 @@
 #!/bin/bash

-# This script builds the HPU docker image and runs the offline inference inside the container.
+# This script build the CPU docker image and run the offline inference inside the container.
 # It serves a sanity check for compilation and basic model usage.
-#
-# vllm-gaudi compatibility pinning:
-#   The vllm-gaudi plugin is installed on top of the vllm upstream checkout used by this CI job.
-#   When upstream vllm changes its API, the plugin may break before it has been updated.
-#   To handle this, the vllm-gaudi repository maintains a file:
-#     vllm/last-good-commit-for-vllm-gaudi/VLLM_COMMUNITY_COMMIT
-#   The first line of that file controls what version of vllm is used inside the Docker image:
-#     - "latest"        : no checkout override; the current Buildkite CI commit is used as-is.
-#     - "<commit SHA>"  : vllm is checked out to that specific commit before building, pinning
-#                         the test to a known-compatible baseline.
-#   To unpin (resume testing against the live vllm tip), set the file content back to "latest".
 set -exuo pipefail

-# Fetch the vllm community commit reference from vllm-gaudi (first line only).
-VLLM_COMMUNITY_COMMIT=$(curl -s \
-  https://raw.githubusercontent.com/vllm-project/vllm-gaudi/vllm/last-good-commit-for-vllm-gaudi/VLLM_COMMUNITY_COMMIT \
-  | head -1 | tr -d '\n')
-
-echo "Using vllm community commit: ${VLLM_COMMUNITY_COMMIT}"
-
 # Try building the docker image
 image_name="hpu/upstream-vllm-ci:${BUILDKITE_COMMIT}"
 container_name="hpu-upstream-vllm-ci-${BUILDKITE_COMMIT}-container"
@@ -30,13 +12,6 @@ FROM gaudi-base-image:latest

 COPY ./ /workspace/vllm

-# If VLLM_COMMUNITY_COMMIT is a specific commit (not "latest"), check it out to pin vllm
-# to the version known to be compatible with vllm-gaudi. When the value is "latest",
-# the current checkout (the Buildkite CI commit) is used unchanged.
-RUN if [ "${VLLM_COMMUNITY_COMMIT}" != "latest" ]; then \
-      cd /workspace/vllm && git fetch --unshallow 2>/dev/null || true && git checkout ${VLLM_COMMUNITY_COMMIT}; \
-    fi
-
 WORKDIR /workspace/vllm

 ENV no_proxy=localhost,127.0.0.1
@@ -44,7 +44,7 @@ docker run \
    python3 examples/offline_inference/basic/generate.py --model ibm-research/PowerMoE-3b  --block-size 64 --enforce-eager -tp 2
    python3 examples/offline_inference/basic/generate.py --model ibm-research/PowerMoE-3b  --block-size 64 --enforce-eager -tp 2 --enable-expert-parallel
    cd tests
-    pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py --ignore=v1/core/test_scheduler_e2e.py
+    pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py
    pytest -v -s v1/engine
    pytest -v -s v1/sample --ignore=v1/sample/test_logprobs.py --ignore=v1/sample/test_logprobs_e2e.py
    pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py
@@ -67,7 +67,7 @@ start_nodes() {
        # 3. map the huggingface cache directory to the container
        # 3. assign ip addresses to the containers (head node: 192.168.10.10, worker nodes:
        #    starting from 192.168.10.11)
-        docker run -d $GPU_DEVICES --shm-size=10.24gb -e HF_TOKEN \
+        docker run -d "$GPU_DEVICES" --shm-size=10.24gb -e HF_TOKEN \
            -v ~/.cache/huggingface:/root/.cache/huggingface --name "node$node" \
            --network docker-net --ip 192.168.10.$((10 + $node)) --rm "$DOCKER_IMAGE" \
            /bin/bash -c "tail -f /dev/null"
@@ -0,0 +1,64 @@
+#!/bin/bash
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Setup script for Prime-RL integration tests
+# This script prepares the environment for running Prime-RL tests with nightly vLLM
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)"
+PRIME_RL_REPO="https://github.com/PrimeIntellect-ai/prime-rl.git"
+PRIME_RL_DIR="${REPO_ROOT}/prime-rl"
+
+if command -v rocm-smi &> /dev/null || command -v rocminfo &> /dev/null; then
+    echo "AMD GPU detected. Prime-RL currently only supports NVIDIA. Skipping..."
+    exit 0
+fi
+
+echo "Setting up Prime-RL integration test environment..."
+
+# Clean up any existing Prime-RL directory
+if [ -d "${PRIME_RL_DIR}" ]; then
+    echo "Removing existing Prime-RL directory..."
+    rm -rf "${PRIME_RL_DIR}"
+fi
+
+# Install UV if not available
+if ! command -v uv &> /dev/null; then
+    echo "Installing UV package manager..."
+    curl -LsSf https://astral.sh/uv/install.sh | sh
+    source "$HOME"/.local/bin/env
+fi
+
+# Clone Prime-RL repository at specific branch for reproducible tests
+PRIME_RL_BRANCH="integ-vllm-main"
+echo "Cloning Prime-RL repository at branch: ${PRIME_RL_BRANCH}..."
+git clone --branch "${PRIME_RL_BRANCH}" --single-branch "${PRIME_RL_REPO}" "${PRIME_RL_DIR}"
+cd "${PRIME_RL_DIR}"
+
+echo "Setting up UV project environment..."
+export UV_PROJECT_ENVIRONMENT=/usr/local
+ln -s /usr/bin/python3 /usr/local/bin/python
+
+# Remove vllm pin from pyproject.toml
+echo "Removing vllm pin from pyproject.toml..."
+sed -i '/vllm==/d' pyproject.toml
+
+# Sync Prime-RL dependencies
+echo "Installing Prime-RL dependencies..."
+uv sync --inexact && uv sync --inexact --all-extras
+
+# Verify installation
+echo "Verifying installations..."
+uv run python -c "import vllm; print(f'vLLM version: {vllm.__version__}')"
+uv run python -c "import prime_rl; print('Prime-RL imported successfully')"
+
+echo "Prime-RL integration test environment setup complete!"
+
+echo "Running Prime-RL integration tests..."
+export WANDB_MODE=offline # this makes this test not require a WANDB_API_KEY
+uv run pytest -vs tests/integration/test_rl.py -m gpu
+
+echo "Prime-RL integration tests completed!"
@@ -1,57 +0,0 @@
-#!/usr/bin/env bash
-set -euxo pipefail
-
-# Nightly e2e test for prefetch offloading with a MoE model.
-# Runs DeepSeek-V2-Lite with prefetch offloading of MoE expert weights
-# and validates GSM8K accuracy matches baseline (no offloading).
-#
-# args: [THRESHOLD] [NUM_QUESTIONS] [START_PORT]
-THRESHOLD=${1:-0.25}
-NUM_Q=${2:-1319}
-PORT=${3:-8030}
-OUT_DIR=${OUT_DIR:-/tmp/vllm-scheduled}
-mkdir -p "${OUT_DIR}"
-
-wait_for_server() {
-  local port=$1
-  timeout 600 bash -c '
-    until curl -sf "http://127.0.0.1:'"$port"'/health" > /dev/null; do
-      sleep 1
-    done'
-}
-
-MODEL="deepseek-ai/DeepSeek-V2-Lite"
-
-cleanup() {
-  if [[ -n "${SERVER_PID:-}" ]] && kill -0 "${SERVER_PID}" 2>/dev/null; then
-    kill "${SERVER_PID}" 2>/dev/null || true
-    for _ in {1..20}; do
-      kill -0 "${SERVER_PID}" 2>/dev/null || break
-      sleep 0.5
-    done
-    kill -9 "${SERVER_PID}" 2>/dev/null || true
-  fi
-}
-trap cleanup EXIT
-
-vllm serve "$MODEL" \
-  --max-model-len 2048 \
-  --offload-group-size 8 \
-  --offload-num-in-group 2 \
-  --offload-prefetch-step 1 \
-  --offload-params w13_weight w2_weight \
-  --port "$PORT" &
-SERVER_PID=$!
-wait_for_server "$PORT"
-
-TAG=$(echo "$MODEL" | tr '/: \\n' '_____')
-OUT="${OUT_DIR}/${TAG}_prefetch_offload.json"
-python3 tests/evals/gsm8k/gsm8k_eval.py --host http://127.0.0.1 --port "$PORT" --num-questions "${NUM_Q}" --save-results "${OUT}"
-python3 - <<PY
-import json; acc=json.load(open('${OUT}'))['accuracy']
-print(f"${MODEL} prefetch_offload: accuracy {acc:.3f}")
-assert acc >= ${THRESHOLD}, f"${MODEL} prefetch_offload accuracy {acc}"
-PY
-
-cleanup
-SERVER_PID=
@@ -72,7 +72,7 @@ obj_json="objects.json"
 aws s3api list-objects-v2 --bucket "$BUCKET" --prefix "$SUBPATH/" --delimiter / --output json > "$obj_json"
 mkdir -p "$INDICES_OUTPUT_DIR"

-# call script to generate indices for all existing wheels
+# call script to generate indicies for all existing wheels
 # this indices have relative paths that could work as long as it is next to the wheel directory in s3
 # i.e., the wheels are always in s3://vllm-wheels/<commit>/
 # and indices can be placed in /<commit>/, or /nightly/, or /<version>/
@@ -54,13 +54,10 @@ mkdir -p $DIST_DIR
 # include only wheels for the release version, ignore all files with "dev" or "rc" in the name (without excluding 'aarch64')
 aws s3 cp --recursive --exclude "*" --include "vllm-${PURE_VERSION}*.whl" --exclude "*dev*" --exclude "*rc[0-9]*" "$S3_COMMIT_PREFIX" $DIST_DIR
 echo "Wheels copied to local directory"
-# generate source distribution using setup.py
-python setup.py sdist --dist-dir=$DIST_DIR
+# generate source tarball
+git archive --format=tar.gz --output="$DIST_DIR/vllm-${PURE_VERSION}.tar.gz" "$BUILDKITE_COMMIT"
 ls -la $DIST_DIR

-SDIST_FILE=$(find $DIST_DIR -name "vllm*.tar.gz")
-echo "Found sdist: $SDIST_FILE"
-
 # upload wheels to PyPI (only default variant, i.e. files without '+' in the name)
 PYPI_WHEEL_FILES=$(find $DIST_DIR -name "vllm-${PURE_VERSION}*.whl" -not -name "*+*")
 if [[ -z "$PYPI_WHEEL_FILES" ]]; then
@@ -68,6 +65,6 @@ if [[ -z "$PYPI_WHEEL_FILES" ]]; then
  exit 1
 fi

-python3 -m twine check "$PYPI_WHEEL_FILES" "$SDIST_FILE"
-python3 -m twine upload --non-interactive --verbose "$PYPI_WHEEL_FILES" "$SDIST_FILE"
-echo "Wheels and source distribution uploaded to PyPI"
+python3 -m twine check "$PYPI_WHEEL_FILES"
+python3 -m twine upload --non-interactive --verbose "$PYPI_WHEEL_FILES"
+echo "Wheels uploaded to PyPI"
@@ -36,16 +36,6 @@ steps:
  - export VLLM_TEST_CLEAN_GPU_MEMORY=1
  - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py

- label: AsyncTP Correctness Tests (B200)
-  timeout_in_minutes: 50
-  working_dir: "/vllm-workspace/"
-  device: b200
-  optional: true
-  num_devices: 2
-  commands:
-  - export VLLM_TEST_CLEAN_GPU_MEMORY=1
-  - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py
-
 - label: Distributed Compile Unit Tests (2xH100)
  timeout_in_minutes: 20
  working_dir: "/vllm-workspace/"
@@ -67,7 +67,6 @@ steps:
  - tests/v1/distributed
  - tests/v1/engine/test_engine_core_client.py
  - tests/distributed/test_symm_mem_allreduce.py
-  - tests/distributed/test_multiproc_executor.py
  commands:
  # https://github.com/NVIDIA/nccl/issues/1838
  - export NCCL_CUMEM_HOST_ENABLE=0
@@ -96,8 +95,6 @@ steps:
  - pytest -v -s distributed/test_pynccl.py
  - pytest -v -s distributed/test_events.py
  - pytest -v -s distributed/test_symm_mem_allreduce.py
-  # test multi-node TP with multiproc executor (simulated on single node)
-  - pytest -v -s distributed/test_multiproc_executor.py::test_multiproc_executor_multi_node
  # TODO: create a dedicated test section for multi-GPU example tests
  # when we have multiple distributed example tests
  # OLD rlhf examples
@@ -106,8 +103,8 @@ steps:
  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 RAY_DEDUP_LOGS=0 python3 rlhf_colocate.py
  # NEW rlhf examples
  - cd new_weight_syncing
-  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_nccl.py
-  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_ipc.py
+  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf.py
+  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_async_new_apis.py

 - label: Distributed Tests (8 GPUs)(H100)
  timeout_in_minutes: 10
@@ -149,7 +146,6 @@ steps:
  num_devices: 2
  commands:
    - pytest -v -s tests/distributed/test_context_parallel.py
-    # - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/offline_inference/new_weight_syncing/rlhf_async_new_apis.py --- failing, need to re-enable
    - VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
    - pytest -v -s tests/v1/distributed/test_dbo.py

@@ -169,7 +165,6 @@ steps:
  num_devices: 2
  num_nodes: 2
  no_plugin: true
-  optional: true # TODO: revert once infra issue solved
  source_file_dependencies:
  - vllm/distributed/
  - vllm/engine/
@@ -213,20 +208,7 @@ steps:
    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
    - CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh

- label: NixlConnector PD + Spec Decode acceptance (2 GPUs)
-  timeout_in_minutes: 30
-  device: a100
-  working_dir: "/vllm-workspace/tests"
-  num_devices: 2
-  source_file_dependencies:
-    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
-    - vllm/v1/worker/kv_connector_model_runner_mixin.py
-    - tests/v1/kv_connector/nixl_integration/
-  commands:
-    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
-    - bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
-
- label: Pipeline + Context Parallelism (4 GPUs)
+- label: Pipeline + Context Parallelism (4 GPUs))
  timeout_in_minutes: 60
  working_dir: "/vllm-workspace/tests"
  num_devices: 4
@@ -29,11 +29,15 @@ steps:
  commands:
  - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020 2 1

- label: DeepSeek V2-Lite Prefetch Offload Accuracy (H100)
-  timeout_in_minutes: 60
-  device: h100
+- label: Prime-RL Integration (2 GPUs)
+  timeout_in_minutes: 30
  optional: true
-  num_devices: 1
+  soft_fail: true
+  num_devices: 2
  working_dir: "/vllm-workspace"
+  source_file_dependencies:
+  - vllm/
+  - .buildkite/scripts/run-prime-rl-test.sh
  commands:
-  - bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_prefetch_offload.sh 0.25 200 8030
+    - nvidia-smi
+    - bash .buildkite/scripts/run-prime-rl-test.sh
@@ -14,7 +14,7 @@ steps:
  commands:
  - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py

- label: V1 e2e + engine (1 GPU)
+- label: V1 e2e + engine
  timeout_in_minutes: 45
  source_file_dependencies:
    - vllm/
@@ -28,43 +28,3 @@ steps:
    - pytest -v -s v1/engine/test_preprocess_error_handling.py
    # Run the rest of v1/engine tests
    - pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py
-  mirror:
-    amd:
-      device: mi325_1
-      depends_on:
-      - image-build-amd
-      commands:
-      - pytest -v -s v1/e2e
-      - pytest -v -s v1/engine
-
- label: V1 e2e (2 GPUs)
-  timeout_in_minutes: 60 # TODO: Fix timeout after we have more confidence in the test stability
-  optional: true
-  num_devices: 2
-  source_file_dependencies:
-    - vllm/
-    - tests/v1/e2e
-  commands:
-    # Only run tests that need exactly 2 GPUs
-    - pytest -v -s v1/e2e/test_spec_decode.py -k "tensor_parallelism"
-  mirror:
-    amd:
-      device: mi325_2
-      depends_on:
-      - image-build-amd
-
- label: V1 e2e (4 GPUs)
-  timeout_in_minutes: 60 # TODO: Fix timeout after we have more confidence in the test stability
-  optional: true
-  num_devices: 4
-  source_file_dependencies:
-    - vllm/
-    - tests/v1/e2e
-  commands:
-    # Only run tests that need 4 GPUs
-    - pytest -v -s v1/e2e/test_spec_decode.py -k "eagle_correctness_heavy"
-  mirror:
-    amd:
-      device: mi325_4
-      depends_on:
-      - image-build-amd
@@ -41,11 +41,6 @@ steps:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/test_oot_registration.py --ignore=entrypoints/openai/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/  --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
  - pytest -v -s entrypoints/test_chat_utils.py
-  mirror:
-    amd:
-      device: mi325_1
-      depends_on:
-      - image-build-amd

 - label: Entrypoints Integration (API Server 2)
  timeout_in_minutes: 130
@@ -60,11 +55,6 @@ steps:
  - pytest -v -s entrypoints/instrumentator
  - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc
  - pytest -v -s tool_use
-  mirror:
-    amd:
-      device: mi325_1
-      depends_on:
-      - image-build-amd

 - label: Entrypoints Integration (Pooling)
  timeout_in_minutes: 50
@@ -75,11 +65,6 @@ steps:
  commands:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
  - pytest -v -s entrypoints/pooling
-  mirror:
-    amd:
-      device: mi325_1
-      depends_on:
-      - image-build-amd

 - label: Entrypoints Integration (Responses API)
  timeout_in_minutes: 50
@@ -97,11 +82,6 @@ steps:
    - tests/v1
  commands:
    - pytest -v -s v1/entrypoints
-  mirror:
-    amd:
-      device: mi325_1
-      depends_on:
-      - image-build-amd

 - label: OpenAI API Correctness
  timeout_in_minutes: 30
@@ -20,19 +20,4 @@ steps:
  - tests/distributed/test_eplb_execute.py
  commands:
  - pytest -v -s distributed/test_eplb_execute.py
-  - pytest -v -s distributed/test_eplb_spec_decode.py
-
- label: Elastic EP Scaling Test
-  timeout_in_minutes: 20
-  device: b200
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  num_devices: 4
-  source_file_dependencies:
-  - vllm/distributed/
-  - vllm/engine/
-  - vllm/executor/
-  - vllm/compilation/
-  - tests/distributed/
-  commands:
-  - pytest -v -s distributed/test_elastic_ep.py
+  - pytest -v -s distributed/test_eplb_spec_decode.py
@@ -44,8 +44,7 @@ steps:
  - vllm/envs.py
  - vllm/config
  commands:
-    - pytest -v -s kernels/moe --ignore=kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
-    - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
+    - pytest -v -s kernels/moe --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
  parallelism: 2

 - label: Kernels Mamba Test
@@ -71,7 +70,7 @@ steps:
  - tests/kernels/moe/test_batched_deepgemm.py
  - tests/kernels/attention/test_deepgemm_attention.py
  commands:
-    - pytest -v -s kernels/quantization/test_block_fp8.py
+    - pytest -v -s kernels/quantization/test_block_fp8.py -k deep_gemm
    - pytest -v -s kernels/moe/test_deepgemm.py
    - pytest -v -s kernels/moe/test_batched_deepgemm.py
    - pytest -v -s kernels/attention/test_deepgemm_attention.py
@@ -116,7 +115,6 @@ steps:
    - pytest -v -s tests/kernels/moe/test_nvfp4_moe.py
    - pytest -v -s tests/kernels/moe/test_ocp_mx_moe.py
    - pytest -v -s tests/kernels/moe/test_flashinfer.py
-    - pytest -v -s tests/kernels/moe/test_flashinfer_moe.py
    - pytest -v -s tests/kernels/moe/test_cutedsl_moe.py
    # e2e
    - pytest -v -s tests/models/quantization/test_nvfp4.py
@@ -156,7 +154,9 @@ steps:
  commands:
    - pytest -v -s kernels/moe/test_deepep_deepgemm_moe.py
    - pytest -v -s kernels/moe/test_deepep_moe.py
-
+    - pytest -v -s kernels/moe/test_pplx_cutlass_moe.py
+    # - pytest -v -s kernels/moe/test_pplx_moe.py - failing on main
+  
 - label: Kernels Fp4 MoE Test (B200)
  timeout_in_minutes: 60
  device: b200
@@ -11,17 +11,17 @@ steps:
  commands:
  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.txt

-# - label: LM Eval Large Models (4 GPUs)(A100)
-#   device: a100
-#   optional: true
-#   num_devices: 4
-#   working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
-#   source_file_dependencies:
-#   - csrc/
-#   - vllm/model_executor/layers/quantization
-#   commands:
-#   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-#   - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large.txt --tp-size=4
+- label: LM Eval Large Models (4 GPUs)(A100)
+  device: a100
+  optional: true
+  num_devices: 4
+  working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
+  source_file_dependencies:
+  - csrc/
+  - vllm/model_executor/layers/quantization
+  commands:
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large.txt --tp-size=4

 - label: LM Eval Large Models (4 GPUs)(H100)
  device: h100
@@ -73,29 +73,3 @@ steps:
  num_devices: 2
  commands:
    - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt
-
- label: GPQA Eval (GPT-OSS) (H100)
-  timeout_in_minutes: 120
-  device: h100
-  optional: true
-  num_devices: 2
-  source_file_dependencies:
-  - csrc/
-  - vllm/model_executor/layers/quantization
-  - tests/evals/gpt_oss/
-  commands:
-    - uv pip install --system 'gpt-oss[eval]==0.0.5'
-    - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-h100.txt
-
- label: GPQA Eval (GPT-OSS) (B200)
-  timeout_in_minutes: 120
-  device: b200
-  optional: true
-  num_devices: 2
-  source_file_dependencies:
-  - csrc/
-  - vllm/model_executor/layers/quantization
-  - tests/evals/gpt_oss/
-  commands:
-    - uv pip install --system 'gpt-oss[eval]==0.0.5'
-    - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-b200.txt
@@ -9,7 +9,6 @@ steps:
    - tests/v1
  commands:
    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
    # split the test to avoid interference
    - pytest -v -s -m 'not cpu_test' v1/core
    - pytest -v -s v1/executor
@@ -17,7 +16,6 @@ steps:
    - pytest -v -s v1/sample
    - pytest -v -s v1/logits_processors
    - pytest -v -s v1/worker
-    # TODO: create another `optional` test group for slow tests
    - pytest -v -s -m 'not slow_test' v1/spec_decode
    - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
    - pytest -v -s -m 'not cpu_test' v1/metrics
@@ -27,11 +25,6 @@ steps:
    # Integration test for streaming correctness (requires special branch).
    - pip install -U git+https://github.com/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api
    - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
-  mirror:
-    amd:
-      device: mi325_1
-      depends_on:
-      - image-build-amd

 - label: V1 Others (CPU)
  depends_on:
@@ -87,11 +80,6 @@ steps:
    - python3 offline_inference/spec_decode.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
    # https://github.com/vllm-project/vllm/pull/26682 uses slightly more memory in PyTorch 2.9+ causing this test to OOM in 1xL4 GPU
    - python3 offline_inference/spec_decode.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:
-      device: mi325_1
-      depends_on:
-      - image-build-amd

 - label: Metrics, Tracing (2 GPUs)
  timeout_in_minutes: 20
@@ -120,11 +108,9 @@ steps:
  timeout_in_minutes: 50
  source_file_dependencies:
  - vllm/
-  - tests/detokenizer
  - tests/multimodal
  - tests/utils_
  commands:
-  - pytest -v -s detokenizer
  - pytest -v -s -m 'not cpu_test' multimodal
  - pytest -v -s utils_

@@ -159,6 +145,20 @@ steps:
  - pytest -v -s transformers_utils
  - pytest -v -s config

+- label: GPT-OSS Eval (B200)
+  timeout_in_minutes: 60
+  working_dir: "/vllm-workspace/"
+  device: b200
+  optional: true
+  source_file_dependencies:
+  - tests/evals/gpt_oss
+  - vllm/model_executor/models/gpt_oss.py
+  - vllm/model_executor/layers/quantization/mxfp4.py
+  - vllm/v1/attention/backends/flashinfer.py
+  commands:
+    - uv pip install --system 'gpt-oss[eval]==0.0.5'
+    - pytest -s -v tests/evals/gpt_oss/test_gpqa_correctness.py --model openai/gpt-oss-20b --metric 0.58
+
 - label: Batch Invariance (H100)
  timeout_in_minutes: 25
  device: h100
@@ -55,15 +55,6 @@ steps:
    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
    - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.2'
    - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'
-  mirror:
-    amd:
-      device: mi325_1
-      depends_on:
-      - image-build-amd
-      commands:
-      - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr'
-      - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.2'
-      - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'

 - label: Language Models Test (PPL)
  timeout_in_minutes: 110
@@ -82,11 +73,6 @@ steps:
  - tests/models/language/pooling
  commands:
    - pytest -v -s models/language/pooling -m 'not core_model'
-  mirror:
-    amd:
-      device: mi325_1
-      depends_on:
-      - image-build-amd

 - label: Language Models Test (MTEB)
  timeout_in_minutes: 110
@@ -12,11 +12,6 @@ steps:
    - pip freeze | grep -E 'torch'
    - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/processing
    - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model  # Otherwise, mp_method="spawn" doesn't work
-  mirror:
-    amd:
-      device: mi325_1
-      depends_on:
-      - image-build-amd

 - label: Multi-Modal Processor Test (CPU)
  depends_on: 
@@ -25,7 +20,6 @@ steps:
  source_file_dependencies:
  - vllm/
  - tests/models/multimodal
-  - tests/models/registry.py
  device: cpu
  commands:
    - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
@@ -36,7 +30,6 @@ steps:
  source_file_dependencies:
  - vllm/
  - tests/models/multimodal
-  - tests/models/registry.py
  commands:
    - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
    - pytest -v -s models/multimodal/processing/test_tensor_schema.py
@@ -59,11 +52,6 @@ steps:
  commands:
    - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
    - pytest -v -s models/multimodal -m 'not core_model' --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/processing
-  mirror:
-    amd:
-      device: mi325_1
-      depends_on:
-      - image-build-amd

 - label: Multi-Modal Models (Extended) 2
  optional: true
@@ -82,3 +70,12 @@ steps:
  commands:
    - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
    - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=1) and not core_model'
+
+# This test is used only in PR development phase to test individual models and should never run on main
+- label: Custom Models
+  optional: true
+  commands:
+    - echo 'Testing custom models...'
+    # PR authors can temporarily add commands below to test individual models
+    # e.g. pytest -v -s models/encoder_decoder/vision_language/test_mllama.py
+    # *To avoid merge conflicts, remember to REMOVE (not just comment out) them before merging the PR*
@@ -15,17 +15,10 @@ steps:
  - pytest -v -s plugins_tests/test_platform_plugins.py
  - pip uninstall vllm_add_dummy_platform -y
  # end platform plugin tests
-  # begin io_processor plugins test
-  # test generic io_processor plugins functions
-  - pytest -v -s ./plugins_tests/test_io_processor_plugins.py
-  # test Terratorch io_processor plugins
+  # begin io_processor plugins test, all the code in between uses the prithvi_io_processor plugin
  - pip install -e ./plugins/prithvi_io_processor_plugin
-  - pytest -v -s plugins_tests/test_terratorch_io_processor_plugins.py
+  - pytest -v -s plugins_tests/test_io_processor_plugins.py
  - pip uninstall prithvi_io_processor_plugin -y
-  # test bge_m3_sparse io_processor plugin
-  - pip install -e ./plugins/bge_m3_sparse_plugin
-  - pytest -v -s plugins_tests/test_bge_m3_sparse_io_processor_plugins.py
-  - pip uninstall bge_m3_sparse_plugin -y
  # end io_processor plugins test
  # begin stat_logger plugins test
  - pip install -e ./plugins/vllm_add_dummy_stat_logger
@@ -39,8 +32,3 @@ steps:
  - pytest -v -s entrypoints/openai/test_oot_registration.py # it needs a clean process
  - pytest -v -s models/test_oot_registration.py # it needs a clean process
  - pytest -v -s plugins/lora_resolvers # unit tests for in-tree lora resolver plugins
-  mirror:
-    amd:
-      device: mi325_2
-      depends_on:
-      - image-build-amd
@@ -1,16 +0,0 @@
-group: Ray Compatibility
-depends_on:
-  - image-build
-steps:
- label: Ray Dependency Compatibility Check
-  # Informational only — does not block the pipeline.
-  # If this fails, it means the PR introduces a dependency that
-  # conflicts with Ray's dependency constraints.
-  # See https://github.com/vllm-project/vllm/issues/33599
-  soft_fail: true
-  timeout_in_minutes: 10
-  source_file_dependencies:
-  - requirements/
-  - setup.py
-  commands:
-  - bash /vllm-workspace/.buildkite/scripts/check-ray-compatibility.sh
@@ -18,4 +18,4 @@ steps:
      depends_on:
      - image-build-amd
      commands:
-      - pytest -v -s samplers
+      - pytest -v -s -m 'not skip_v1' samplers
@@ -13,13 +13,13 @@ steps:
  commands:
    - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models.txt

-# - label: Weight Loading Multiple GPU - Large Models # optional
-#   working_dir: "/vllm-workspace/tests"
-#   num_devices: 2
-#   device: a100
-#   optional: true
-#   source_file_dependencies:
-#   - vllm/
-#   - tests/weight_loading
-#   commands:
-#     - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-large.txt
+- label: Weight Loading Multiple GPU - Large Models # optional
+  working_dir: "/vllm-workspace/tests"
+  num_devices: 2
+  device: a100
+  optional: true
+  source_file_dependencies:
+  - vllm/
+  - tests/weight_loading
+  commands:
+    - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-large.txt
@@ -0,0 +1,24 @@
+# doc: https://github.com/pytorch/test-infra/blob/main/tools/stronghold/docs/bc_linter_config.md
+version: 1
+paths:
+# We temporarily disable globally, and will only enable with `annotations.include`
+# include:
+#   - "vllm/v1/attetion/*.py"
+#   - "vllm/v1/core/*.py"
+exclude:
+  - "**/*.py"
+
+scan:
+  functions: true        # check free functions and methods
+  classes: true          # check classes/dataclasses
+  public_only: true      # ignore names starting with "_" at any level
+
+annotations:
+  include:               # decorators that force‑include a symbol
+    - name: "bc_linter_include"  # matched by simple name or dotted suffix
+      propagate_to_members: false # for classes, include methods/inner classes
+  exclude:               # decorators that force‑exclude a symbol
+    - name: "bc_linter_skip"     # matched by simple name or dotted suffix
+      propagate_to_members: true  # for classes, exclude methods/inner classes
+
+excluded_violations: []  # e.g. ["ParameterRenamed", "FieldTypeChanged"]
@@ -2,17 +2,17 @@
 # for more info about CODEOWNERS file

 # This lists cover the "core" components of vLLM that require careful review
-/vllm/compilation @zou3519 @youkaichao @ProExpertProg @BoyuanFeng
+/vllm/compilation @zou3519 @youkaichao @ProExpertProg
 /vllm/distributed/kv_transfer @NickLucche @ApostaC @orozery
 /vllm/lora @jeejeelee
-/vllm/model_executor/layers/attention @LucasWilkinson @MatthewBonanni
+/vllm/model_executor/layers/attention @LucasWilkinson
 /vllm/model_executor/layers/fused_moe @mgoin @pavanimajety
 /vllm/model_executor/layers/quantization @mgoin @robertgshaw2-redhat @tlrmchlsmth @yewentao256 @pavanimajety
 /vllm/model_executor/layers/mamba @tdoublep
 /vllm/model_executor/model_loader @22quinn
 /vllm/model_executor/layers/batch_invariant.py @yewentao256 
 /vllm/multimodal @DarkLight1337 @ywang96 @NickLucche @tjtanaa
-/vllm/vllm_flash_attn @LucasWilkinson @MatthewBonanni
+/vllm/vllm_flash_attn @LucasWilkinson
 CMakeLists.txt @tlrmchlsmth @LucasWilkinson

 # Any change to the VllmConfig changes can have a large user-facing impact,
@@ -43,25 +43,22 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
 /vllm/tool_parsers @aarnphm @chaunceyjiang

 # vLLM V1
-/vllm/v1/attention @LucasWilkinson @MatthewBonanni
+/vllm/v1/attention @LucasWilkinson
 /vllm/v1/attention/backend.py @WoosukKwon @zhuohan123 @youkaichao @alexm-redhat @njhill
 /vllm/v1/attention/backends/mla @pavanimajety
 /vllm/v1/attention/backends/flashinfer.py @mgoin @pavanimajety
 /vllm/v1/attention/backends/triton_attn.py @tdoublep
 /vllm/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC @orozery
 /vllm/v1/sample @22quinn @houseroad @njhill
-/vllm/v1/spec_decode @benchislett @luccafong @MatthewBonanni
+/vllm/v1/spec_decode @benchislett @luccafong
 /vllm/v1/structured_output @mgoin @russellb @aarnphm @benchislett
 /vllm/v1/kv_cache_interface.py @heheda12345
 /vllm/v1/kv_offload @ApostaC @orozery
-/vllm/v1/engine @njhill
-/vllm/v1/executor @njhill
-/vllm/v1/worker @njhill
-/vllm/v1/worker/kv_connector_model_runner_mixin.py @orozery @NickLucche
+/vllm/v1/worker/gpu/kv_connector.py @orozery
+/vllm/v1/worker/kv_connector_model_runner_mixin.py @orozery

 # Model runner V2
-/vllm/v1/worker/gpu @WoosukKwon @njhill
-/vllm/v1/worker/gpu/kv_connector.py @orozery
+/vllm/v1/worker/gpu @WoosukKwon

 # Test ownership
 /.buildkite/lm-eval-harness @mgoin 
@@ -3,7 +3,6 @@ pull_request_rules:
  description: Automatically apply documentation label
  conditions:
    - label != stale
-    - -closed
    - or:
      - files~=^[^/]+\.md$
      - files~=^docs/
@@ -260,7 +259,8 @@ pull_request_rules:
      - files=benchmarks/run_structured_output_benchmark.sh
      - files=docs/features/structured_outputs.md
      - files=examples/offline_inference/structured_outputs.py
-      - files=examples/online_serving/structured_outputs/structured_outputs.py
+      - files=examples/online_serving/openai_chat_completion_structured_outputs.py
+      - files=examples/online_serving/openai_chat_completion_structured_outputs_with_reasoning.py
      - files~=^tests/v1/structured_output/
      - files=tests/v1/entrypoints/llm/test_struct_output_generate.py
      - files~=^vllm/v1/structured_output/
@@ -0,0 +1,29 @@
+name: BC Lint
+
+on:
+  pull_request:
+    types:
+      - opened
+      - synchronize
+      - reopened
+      - labeled
+      - unlabeled
+
+jobs:
+  bc_lint:
+    if: github.repository_owner == 'vllm-project'
+    runs-on: ubuntu-latest
+    steps:
+      - name: Run BC Lint Action
+        uses: pytorch/test-infra/.github/actions/bc-lint@main
+        with:
+          repo: ${{ github.event.pull_request.head.repo.full_name }}
+          base_sha: ${{ github.event.pull_request.base.sha }}
+          head_sha: ${{ github.event.pull_request.head.sha }}
+          suppression: ${{ contains(github.event.pull_request.labels.*.name, 'suppress-bc-linter') }}
+          docs_link: 'https://github.com/pytorch/test-infra/wiki/BC-Linter'
+          config_dir: .github
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
+  cancel-in-progress: true
@@ -6,9 +6,6 @@ on:
      - main
  workflow_dispatch:  # Manual trigger

-permissions:
-  contents: read
-
 jobs:
  macos-m1-smoke-test:
    runs-on: macos-latest
@@ -3,8 +3,6 @@

 # vllm-flash-attn built from source
 vllm/vllm_flash_attn/*
-!vllm/vllm_flash_attn/__init__.py
-!vllm/vllm_flash_attn/flash_attn_interface.py

 # OpenAI triton kernels copied from source
 vllm/third_party/triton_kernels/*
@@ -13,7 +13,7 @@ repos:
    args: [--output-format, github, --fix]
  - id: ruff-format
 - repo: https://github.com/crate-ci/typos
-  rev: v1.43.5
+  rev: v1.38.1
  hooks:
  - id: typos
    args: [--force-exclude]
@@ -38,7 +38,7 @@ repos:
  rev: 0.9.1
  hooks:
    - id: pip-compile
-      args: [requirements/test.in, -o, requirements/test.txt, --index-strategy, unsafe-best-match, --torch-backend, cu129, --python-platform, x86_64-manylinux_2_28, --python-version, "3.12"]
+      args: [requirements/test.in, -o, requirements/test.txt, --index-strategy, unsafe-best-match, --torch-backend, cu128, --python-platform, x86_64-manylinux_2_28, --python-version, "3.12"]
      files: ^requirements/test\.(in|txt)$
 - repo: local
  hooks:
@@ -55,7 +55,7 @@ repos:
      language: python
      types_or: [python, pyi]
      require_serial: true
-      additional_dependencies: ["mypy[faster-cache]==1.15.0", regex, types-cachetools, types-setuptools, types-PyYAML, types-requests, types-torch, pydantic]
+      additional_dependencies: [mypy==1.11.1, regex, types-cachetools, types-setuptools, types-PyYAML, types-requests, types-torch, pydantic]
  - id: mypy-3.10 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward
    name: Run mypy for Python 3.10
    entry: python tools/pre_commit/mypy.py 1 "3.10"
@@ -127,13 +127,6 @@ repos:
    language: python
    types: [python]
    additional_dependencies: [regex]
-  # prevent use torch.cuda APIs
-  - id: check-torch-cuda-call
-    name: "Prevent new 'torch.cuda' APIs call"
-    entry: python tools/pre_commit/check_torch_cuda.py
-    language: python
-    types: [python]
-    additional_dependencies: [regex]
  - id: validate-config
    name: Validate configuration has default values and that each field has a docstring
    entry: python tools/pre_commit/validate_config.py
@@ -9,7 +9,6 @@ build:
    python: "3.12"
  jobs:
    post_checkout:
-      - bash docs/maybe_skip_pr_build.sh
      - git fetch origin main --unshallow --no-tags --filter=blob:none || true
    pre_create_environment:
      - pip install uv
@@ -725,7 +725,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
  # CUTLASS MoE kernels

  # The MoE kernel cutlass_moe_mm requires CUDA 12.3 or later (and ONLY works
-  # on Hopper). get_cutlass_(batched_)moe_mm_data should only be compiled
+  # on Hopper). get_cutlass_(pplx_)moe_mm_data should only be compiled
  # if it's possible to compile MoE kernels that use its output.
  cuda_archs_loose_intersection(SCALED_MM_ARCHS "9.0a" "${CUDA_ARCHS}")
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3 AND SCALED_MM_ARCHS)
@@ -771,51 +771,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
    endif()
  endif()

-  # Expert-specialization MXFP8 blockscaled grouped kernels (SM100+).
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(ES_MXFP8_GROUPED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
-  else()
-    cuda_archs_loose_intersection(ES_MXFP8_GROUPED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
-  endif()
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND ES_MXFP8_GROUPED_MM_ARCHS)
-    set(SRCS
-      "csrc/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm.cu"
-      "csrc/moe/mxfp8_moe/mxfp8_experts_quant.cu")
-    set_gencode_flags_for_srcs(
-      SRCS "${SRCS}"
-      CUDA_ARCHS "${ES_MXFP8_GROUPED_MM_ARCHS}")
-    list(APPEND VLLM_EXT_SRC "${SRCS}")
-    list(APPEND VLLM_GPU_FLAGS "-DENABLE_ES_MXFP8_GROUPED_MM_SM100=1")
-    message(STATUS "Building ES MXFP8 grouped kernels for archs: ${ES_MXFP8_GROUPED_MM_ARCHS}")
-  else()
-    if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8
-        AND ES_MXFP8_GROUPED_MM_ARCHS)
-      message(STATUS "Not building ES MXFP8 grouped kernels as CUDA Compiler version is "
-                     "not >= 12.8.")
-    else()
-      message(STATUS "Not building ES MXFP8 grouped kernels as no compatible archs found "
-                     "in CUDA target architectures.")
-    endif()
-  endif()
-
-  # DeepSeek V3 fused A GEMM kernel (requires SM 9.0+, Hopper and later)
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0f;11.0f" "${CUDA_ARCHS}")
-  else()
-    cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
-  endif()
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND DSV3_FUSED_A_GEMM_ARCHS)
-    set(DSV3_FUSED_A_GEMM_SRC "csrc/dsv3_fused_a_gemm.cu")
-    set_gencode_flags_for_srcs(
-      SRCS "${DSV3_FUSED_A_GEMM_SRC}"
-      CUDA_ARCHS "${DSV3_FUSED_A_GEMM_ARCHS}")
-    list(APPEND VLLM_EXT_SRC ${DSV3_FUSED_A_GEMM_SRC})
-    message(STATUS "Building dsv3_fused_a_gemm for archs: ${DSV3_FUSED_A_GEMM_ARCHS}")
-  else()
-    message(STATUS "Not building dsv3_fused_a_gemm as no compatible archs found "
-                   "in CUDA target architectures.")
-  endif()
-
  # 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;11.0f;12.0f" "${CUDA_ARCHS}")
@@ -998,8 +953,7 @@ set(VLLM_MOE_EXT_SRC
 if(VLLM_GPU_LANG STREQUAL "CUDA")
  list(APPEND VLLM_MOE_EXT_SRC
    "csrc/moe/moe_wna16.cu"
-    "csrc/moe/grouped_topk_kernels.cu"
-    "csrc/moe/router_gemm.cu")
+    "csrc/moe/grouped_topk_kernels.cu")
 endif()

 if(VLLM_GPU_LANG STREQUAL "CUDA")
@@ -1128,27 +1082,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
    message(STATUS "Not building Marlin MOE kernels as no compatible archs found"
                   " in CUDA target architectures")
  endif()
-
-  # DeepSeek V3 router GEMM kernel - requires SM90+
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(DSV3_ROUTER_GEMM_ARCHS "9.0a;10.0f;11.0f" "${CUDA_ARCHS}")
-  else()
-    cuda_archs_loose_intersection(DSV3_ROUTER_GEMM_ARCHS "9.0a;10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
-  endif()
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND DSV3_ROUTER_GEMM_ARCHS)
-    set(DSV3_ROUTER_GEMM_SRC
-      "csrc/moe/dsv3_router_gemm_entry.cu"
-      "csrc/moe/dsv3_router_gemm_float_out.cu"
-      "csrc/moe/dsv3_router_gemm_bf16_out.cu")
-    set_gencode_flags_for_srcs(
-      SRCS "${DSV3_ROUTER_GEMM_SRC}"
-      CUDA_ARCHS "${DSV3_ROUTER_GEMM_ARCHS}")
-    list(APPEND VLLM_MOE_EXT_SRC "${DSV3_ROUTER_GEMM_SRC}")
-    message(STATUS "Building DSV3 router GEMM kernel for archs: ${DSV3_ROUTER_GEMM_ARCHS}")
-  else()
-    message(STATUS "Not building DSV3 router GEMM kernel as no compatible archs found"
-                   " (requires SM90+ and CUDA >= 12.0)")
-  endif()
 endif()

 message(STATUS "Enabling moe extension.")
@@ -15,6 +15,7 @@ from .common import (
    BenchmarkConfig,
    BenchmarkResult,
    MockLayer,
+    MockModelConfig,
    ResultsFormatter,
    get_attention_scale,
    is_mla_backend,
@@ -35,6 +36,7 @@ __all__ = [
    "ResultsFormatter",
    # Mock objects
    "MockLayer",
+    "MockModelConfig",
    # Utilities
    "setup_mla_dims",
    "get_attention_scale",
@@ -10,6 +10,7 @@ from dataclasses import asdict, dataclass
 from pathlib import Path
 from typing import Any

+import numpy as np
 import torch
 from batch_spec import get_batch_type, parse_batch_spec
 from rich.console import Console
@@ -30,7 +31,7 @@ def batch_spec_sort_key(spec: str) -> tuple[int, int, int]:
        max_kv_len = max(r.kv_len for r in requests) if requests else 0
        return (batch_size, max_q_len, max_kv_len)
    except Exception:
-        # Fallback for unparsable specs
+        # Fallback for unparseable specs
        return (0, 0, 0)


@@ -61,7 +62,10 @@ class MockHfConfig:
 # Import AttentionLayerBase at module level to avoid circular dependencies
 try:
    from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
+
+    _HAS_ATTENTION_LAYER_BASE = True
 except ImportError:
+    _HAS_ATTENTION_LAYER_BASE = False
    AttentionLayerBase = object  # Fallback


@@ -163,6 +167,95 @@ class MockLayer(AttentionLayerBase):
        return self._kv_cache_spec


+class MockModelConfig:
+    """Mock model configuration."""
+
+    def __init__(
+        self,
+        num_q_heads: int,
+        num_kv_heads: int,
+        head_dim: int,
+        dtype: torch.dtype = torch.float16,
+        max_model_len: int = 32768,
+    ):
+        self._n_q = num_q_heads
+        self._n_kv = num_kv_heads
+        self._d = head_dim
+        self.dtype = dtype
+        self.max_model_len = max_model_len
+
+    def get_num_attention_heads(self, _=None) -> int:
+        return self._n_q
+
+    def get_num_kv_heads(self, _=None) -> int:
+        return self._n_kv
+
+    def get_head_size(self) -> int:
+        return self._d
+
+    def get_num_layers(self) -> int:
+        """Mock method for layer count queries."""
+        return 1
+
+    def get_sliding_window_for_layer(self, _layer_idx: int):
+        """Mock method for sliding window queries."""
+        return None
+
+    def get_logits_soft_cap_for_layer(self, _layer_idx: int):
+        """Mock method for logits soft cap queries."""
+        return None
+
+    def get_sm_scale_for_layer(self, _layer_idx: int) -> float:
+        """Mock method for SM scale queries."""
+        return 1.0 / (self.get_head_size() ** 0.5)
+
+
+class MockParallelConfig:
+    """Mock parallel configuration."""
+
+    pass
+
+
+class MockCompilationConfig:
+    """Mock compilation configuration."""
+
+    def __init__(self):
+        self.full_cuda_graph = False
+        self.static_forward_context = {}
+
+
+class MockVLLMConfig:
+    """Mock VLLM configuration."""
+
+    def __init__(self):
+        self.compilation_config = MockCompilationConfig()
+
+
+class MockRunner:
+    """Mock GPU runner for metadata builders."""
+
+    def __init__(
+        self,
+        seq_lens: np.ndarray,
+        query_start_locs: np.ndarray,
+        device: torch.device,
+        num_q_heads: int,
+        num_kv_heads: int,
+        head_dim: int,
+        dtype: torch.dtype,
+    ):
+        self.model_config = MockModelConfig(num_q_heads, num_kv_heads, head_dim, dtype)
+        self.parallel_config = MockParallelConfig()
+        self.vllm_config = MockVLLMConfig()
+        self.seq_lens_np = seq_lens
+        self.query_start_loc_np = query_start_locs
+        self.device = device
+        self.attention_chunk_size = None
+        self.num_query_heads = num_q_heads
+        self.num_kv_heads = num_kv_heads
+        self.dtype = dtype
+
+
@dataclass
 class ParameterSweep:
    """Configuration for sweeping a backend parameter."""
@@ -701,7 +701,7 @@ def _run_single_benchmark(
    # Warmup
    for _ in range(config.warmup_iters):
        forward_fn()
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    # Benchmark
    times = []
@@ -714,7 +714,7 @@ def _run_single_benchmark(
            forward_fn()
        end.record()

-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()
        elapsed_ms = start.elapsed_time(end)
        times.append(elapsed_ms / 1000.0 / config.num_layers)

@@ -391,7 +391,7 @@ def _run_single_benchmark(
                attn_metadata,
                output=out,
            )
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    # Benchmark
    times = []
@@ -412,7 +412,7 @@ def _run_single_benchmark(
            )
        end.record()

-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()
        elapsed_ms = start.elapsed_time(end)
        times.append(elapsed_ms / 1000.0 / config.num_layers)  # seconds per layer

@@ -85,6 +85,7 @@ start_server() {
    # Each argument and its value are separate elements.
    local common_args_array=(
        "$MODEL"
+        "--disable-log-requests"
        "--port" "8004"
        "--host" "$HOSTNAME"
        "--gpu-memory-utilization" "$gpu_memory_utilization"
@@ -649,3 +649,9 @@ ASYNC_REQUEST_FUNCS = {
    "sglang": async_request_openai_completions,
    "llama.cpp": async_request_openai_completions,
 }
+
+OPENAI_COMPATIBLE_BACKENDS = [
+    k
+    for k, v in ASYNC_REQUEST_FUNCS.items()
+    if v in (async_request_openai_completions, async_request_openai_chat_completions)
+]
@@ -94,7 +94,7 @@ def create_logits(

 def measure_memory() -> tuple[int, int]:
    """Return (allocated, reserved) memory in bytes."""
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()
    return torch.cuda.memory_allocated(), torch.cuda.max_memory_allocated()


@@ -102,7 +102,7 @@ def reset_memory_stats():
    """Reset peak memory statistics."""
    reset_buffer_cache()
    torch.cuda.reset_peak_memory_stats()
-    torch.accelerator.empty_cache()
+    torch.cuda.empty_cache()
    gc.collect()


@@ -123,7 +123,7 @@ def benchmark_function(
    for _ in range(warmup_iters):
        logits_copy = logits.clone()
        func(logits_copy, k, p)
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    # Reset memory stats before benchmark
    reset_memory_stats()
@@ -140,7 +140,7 @@ def benchmark_function(
        func(logits_copy, k, p)
        end_events[i].record()

-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    # Calculate timing
    times = [
@@ -1,7 +1,78 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import argparse
+import json
+import math
+import os
 import time
 from types import TracebackType
+from typing import Any
+
+
+def convert_to_pytorch_benchmark_format(
+    args: argparse.Namespace, metrics: dict[str, list], extra_info: dict[str, Any]
+) -> list:
+    """
+    Save the benchmark results in the format used by PyTorch OSS benchmark with
+    on metric per record
+    https://github.com/pytorch/pytorch/wiki/How-to-integrate-with-PyTorch-OSS-benchmark-database
+    """
+    records = []
+    if not os.environ.get("SAVE_TO_PYTORCH_BENCHMARK_FORMAT", False):
+        return records
+
+    for name, benchmark_values in metrics.items():
+        record = {
+            "benchmark": {
+                "name": "vLLM benchmark",
+                "extra_info": {
+                    "args": vars(args),
+                },
+            },
+            "model": {
+                "name": args.model,
+            },
+            "metric": {
+                "name": name,
+                "benchmark_values": benchmark_values,
+                "extra_info": extra_info,
+            },
+        }
+
+        tp = record["benchmark"]["extra_info"]["args"].get("tensor_parallel_size")
+        # Save tensor_parallel_size parameter if it's part of the metadata
+        if not tp and "tensor_parallel_size" in extra_info:
+            record["benchmark"]["extra_info"]["args"]["tensor_parallel_size"] = (
+                extra_info["tensor_parallel_size"]
+            )
+
+        records.append(record)
+
+    return records
+
+
+class InfEncoder(json.JSONEncoder):
+    def clear_inf(self, o: Any):
+        if isinstance(o, dict):
+            return {k: self.clear_inf(v) for k, v in o.items()}
+        elif isinstance(o, list):
+            return [self.clear_inf(v) for v in o]
+        elif isinstance(o, float) and math.isinf(o):
+            return "inf"
+        return o
+
+    def iterencode(self, o: Any, *args, **kwargs) -> Any:
+        return super().iterencode(self.clear_inf(o), *args, **kwargs)
+
+
+def write_to_json(filename: str, records: list) -> None:
+    with open(filename, "w") as f:
+        json.dump(
+            records,
+            f,
+            cls=InfEncoder,
+            default=lambda o: f"<{type(o).__name__} object is not JSON serializable>",
+        )


 # Collect time and generate time metrics
@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

 # Cutlass bench utils
+from collections.abc import Iterable

 import torch

@@ -85,3 +86,15 @@ def make_rand_sparse_tensors(

    # Compressed B, Metadata, Original A, B
    return b_compressed, e, a, b
+
+
+def make_n_rand_sparse_tensors(
+    num_tensors: int, dtype: torch.dtype, m: int, n: int, k: int
+) -> tuple[Iterable[torch.Tensor], Iterable[torch.Tensor]]:
+    ABs = []
+    for _ in range(num_tensors):
+        b_comp, e, a, b = make_rand_sparse_tensors(dtype, m, n, k)
+        if b_comp is not None:
+            ABs.append(make_rand_sparse_tensors(dtype, m, n, k))
+    BComps, Es, As, Bs = zip(*ABs)
+    return list(BComps), list(Es), list(As), list(Bs)
@@ -0,0 +1,45 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import asyncio
+import time
+
+
+class RateLimiter:
+    """Token bucket rate limiter implementation"""
+
+    def __init__(self, rate_limit):
+        self.rate_limit = rate_limit  # Requests per second
+        self.num_available_tokens = rate_limit  # Available tokens
+        self.last_refill = time.monotonic()  # Last token refill time
+        self.lock = asyncio.Lock()  # Synchronization lock
+
+    async def acquire(self):
+        """Acquire a token from the rate limiter"""
+        while True:
+            async with self.lock:
+                current_time = time.monotonic()
+                elapsed = current_time - self.last_refill
+
+                # Refill num_available_tokens if more than 1 second has passed
+                if elapsed > 1.0:
+                    self.num_available_tokens = self.rate_limit
+                    self.last_refill = current_time
+
+                # Check if num_available_tokens are available
+                if self.num_available_tokens > 0:
+                    self.num_available_tokens -= 1
+                    return True
+
+                # Calculate wait time if no num_available_tokens available
+                wait_time = 1.0 - elapsed
+            await asyncio.sleep(wait_time)
+
+    async def __aenter__(self):
+        """Enter async context manager - acquire token"""
+        await self.acquire()
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        """Exit async context manager - no cleanup needed"""
+        pass
@@ -0,0 +1,39 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import asyncio
+from collections import deque
+
+
+class RequestQueue:
+    """Request queue manager with concurrency control"""
+
+    def __init__(self, max_concurrent, max_queue_size):
+        # Maximum concurrent requests
+        self.max_concurrent = max_concurrent
+        self.max_queue_size = max_queue_size  # Maximum queue size
+        # Concurrency control
+        self.semaphore = asyncio.Semaphore(max_concurrent)
+        self.queue = deque()  # Request queue
+        self.queue_size = 0  # Current queue size
+        self.lock = asyncio.Lock()  # Sync queue Lock
+
+    async def enqueue(self, task):
+        """Add a request task to the queue"""
+        async with self.lock:
+            if self.queue_size >= self.max_queue_size:
+                return False
+
+            self.queue.append(task)
+            self.queue_size += 1
+            return True
+
+    async def process(self):
+        """Process queued requests using semaphore for concurrency control"""
+        while True:
+            if self.queue:
+                async with self.semaphore, self.lock:
+                    task = self.queue.popleft()
+                    self.queue_size -= 1
+                    await task
+            await asyncio.sleep(0.01)  # Yield control to event loop
@@ -13,7 +13,6 @@ from torch.utils.benchmark import Measurement as TMeasurement
 from tqdm import tqdm

 import vllm._custom_ops as ops
-from vllm.benchmarks.lib.utils import default_vllm_config
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
    per_token_group_quant_fp8,
@@ -292,7 +291,6 @@ def print_timers(timers: Iterable[TMeasurement]):
    compare.print()


-@default_vllm_config()
 def main():
    torch.set_default_device("cuda")
    bench_params = get_bench_params()
@@ -8,7 +8,6 @@ os.environ["VLLM_USE_DEEP_GEMM"] = "0"

 import torch

-from vllm.benchmarks.lib.utils import default_vllm_config
 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
    W8A8BlockFp8LinearOp,
 )
@@ -41,7 +40,6 @@ DEEPSEEK_V3_SHAPES = [
 ]


-@default_vllm_config()
 def build_w8a8_block_fp8_runner(M, N, K, block_size, device, use_cutlass):
    """Build runner function for w8a8 block fp8 matmul."""
    factor_for_scale = 1e-2
@@ -7,7 +7,6 @@ from unittest.mock import patch
 import pandas as pd
 import torch

-from vllm.benchmarks.lib.utils import default_vllm_config
 from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
 from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
 from vllm.triton_utils import triton
@@ -85,7 +84,6 @@ def calculate_diff(
 configs = []


-@default_vllm_config()
 def benchmark_quantization(
    batch_size,
    hidden_size,
@@ -168,7 +168,7 @@ def bench_impl(
    # warmup
    for kwargs in kwargs_list:
        impl_type.get_impl()(**kwargs)
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    # Merge into a single kwargs and qualify arguments as ArgPool
    kwargs = {k: ArgPool([]) for k in kwargs_list[0]}
@@ -202,7 +202,7 @@ def test_correctness(T: int, N: int):
    # reference output
    ref_out_q, ref_out_s = output_from_impl(ImplType.REFERENCE)

-    # test output
+    # test ouptut
    out_q, out_s = output_from_impl(
        ImplType.SILU_MUL_PER_TOKEN_GROUP_QUANT_FP8_COLMAJOR
    )
@@ -7,7 +7,6 @@ import itertools
 import torch

 import vllm.model_executor.layers.activation  # noqa F401
-from vllm.benchmarks.lib.utils import default_vllm_config
 from vllm.model_executor.custom_op import op_registry
 from vllm.triton_utils import triton
 from vllm.utils.argparse_utils import FlexibleArgumentParser
@@ -19,7 +18,6 @@ intermediate_size = [3072, 9728, 12288]
 configs = list(itertools.product(batch_size_range, seq_len_range, intermediate_size))


-@default_vllm_config()
 def benchmark_activation(
    batch_size: int,
    seq_len: int,
@@ -12,12 +12,12 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk
 from tests.kernels.moe.utils import make_dummy_moe_config
 from vllm import _custom_ops as ops
 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
-from vllm.model_executor.layers.fused_moe.all2all_utils import (
-    maybe_make_prepare_finalize,
-)
 from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config
 from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp8
 from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts, fused_topk
+from vllm.model_executor.layers.fused_moe.prepare_finalize import (
+    MoEPrepareAndFinalizeNoEP,
+)
 from vllm.platforms import current_platform
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 from vllm.v1.worker.workspace import init_workspace_manager
@@ -137,21 +137,15 @@ def bench_run(
        per_out_ch_quant=per_out_ch,
    )

-    moe_config = make_dummy_moe_config(
-        num_experts=num_experts,
-        hidden_dim=k,
-        intermediate_size_per_partition=n,
-        in_dtype=a.dtype,
-    )
-    fn = mk.FusedMoEKernel(
-        maybe_make_prepare_finalize(
-            moe=moe_config,
-            quant_config=quant_config,
-            allow_new_interface=True,
-            use_monolithic=False,
-        ),
+    fn = mk.FusedMoEModularKernel(
+        MoEPrepareAndFinalizeNoEP(),
        CutlassExpertsFp8(
-            moe_config=moe_config,
+            moe_config=make_dummy_moe_config(
+                num_experts=num_experts,
+                hidden_dim=k,
+                intermediate_size_per_partition=n,
+                in_dtype=a.dtype,
+            ),
            quant_config=quant_config,
        ),
    )
@@ -171,7 +165,7 @@ def bench_run(
                activation=MoEActivation.SILU,
                global_num_experts=num_experts,
            )
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    # Create CUDA graphs for Triton (match benchmark_moe.py pattern exactly)
    triton_stream = torch.cuda.Stream()
@@ -187,14 +181,14 @@ def bench_run(
                topk_ids,
                quant_config=quant_config,
            )
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    def bench_cuda_graph(graph, num_warmup=5, num_iters=100):
        """Benchmark CUDA graph using events like benchmark_moe.py"""
        # Warmup
        for _ in range(num_warmup):
            graph.replay()
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()

        # Timing
        start_event = torch.Event(enable_timing=True)
@@ -202,7 +196,7 @@ def bench_run(

        latencies = []
        for _ in range(num_iters):
-            torch.accelerator.synchronize()
+            torch.cuda.synchronize()
            start_event.record()
            graph.replay()
            end_event.record()
@@ -15,9 +15,6 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk
 from tests.kernels.moe.utils import make_dummy_moe_config
 from vllm import _custom_ops as ops
 from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
-from vllm.model_executor.layers.fused_moe.all2all_utils import (
-    maybe_make_prepare_finalize,
-)
 from vllm.model_executor.layers.fused_moe.config import (
    fp8_w8a8_moe_quant_config,
    nvfp4_moe_quant_config,
@@ -26,6 +23,9 @@ from vllm.model_executor.layers.fused_moe.cutlass_moe import (
    CutlassExpertsFp4,
 )
 from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts, fused_topk
+from vllm.model_executor.layers.fused_moe.prepare_finalize import (
+    MoEPrepareAndFinalizeNoEP,
+)
 from vllm.scalar_type import scalar_types
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 from vllm.v1.worker.workspace import init_workspace_manager
@@ -196,21 +196,10 @@ def bench_run(
            g2_alphas=w2_gs,
        )

-        moe_config = make_dummy_moe_config(
-            num_experts=num_experts,
-            hidden_dim=k,
-            intermediate_size_per_partition=n,
-            in_dtype=a.dtype,
-        )
-        kernel = mk.FusedMoEKernel(
-            maybe_make_prepare_finalize(
-                moe=moe_config,
-                quant_config=quant_config,
-                allow_new_interface=True,
-                use_monolithic=False,
-            ),
+        kernel = mk.FusedMoEModularKernel(
+            MoEPrepareAndFinalizeNoEP(),
            CutlassExpertsFp4(
-                moe_config=moe_config,
+                make_dummy_moe_config(),
                quant_config=quant_config,
            ),
        )
@@ -251,17 +240,11 @@ def bench_run(
            g1_alphas=w1_gs,
            g2_alphas=w2_gs,
        )
-        moe_config = make_dummy_moe_config()

-        kernel = mk.FusedMoEKernel(
-            maybe_make_prepare_finalize(
-                moe=moe_config,
-                quant_config=quant_config,
-                allow_new_interface=True,
-                use_monolithic=False,
-            ),
+        kernel = mk.FusedMoEModularKernel(
+            MoEPrepareAndFinalizeNoEP(),
            CutlassExpertsFp4(
-                moe_config=moe_config,
+                make_dummy_moe_config(),
                quant_config=quant_config,
            ),
        )
@@ -307,7 +290,7 @@ def bench_run(
    def replay_graph(graph, num_repeats):
        for _ in range(num_repeats):
            graph.replay()
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()

    cutlass_stream = torch.cuda.Stream()
    cutlass_graph = torch.cuda.CUDAGraph()
@@ -330,7 +313,7 @@ def bench_run(
            e=num_experts,
            device=device,
        )
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    triton_stream = torch.cuda.Stream()
    triton_graph = torch.cuda.CUDAGraph()
@@ -345,7 +328,7 @@ def bench_run(
            w2_fp8scale,
            a_fp8_scale,
        )
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    min_run_time = 5
    num_warmup = 5
@@ -30,9 +30,6 @@ import torch.distributed as dist
 from torch.distributed import ProcessGroup

 from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce
-from vllm.distributed.device_communicators.flashinfer_all_reduce import (
-    FlashInferAllReduce,
-)
 from vllm.distributed.device_communicators.pynccl import (
    PyNcclCommunicator,
    register_nccl_symmetric_ops,
@@ -47,7 +44,7 @@ from vllm.utils.argparse_utils import FlexibleArgumentParser
 logger = init_logger(__name__)

 # Default sequence lengths to benchmark
-DEFAULT_SEQUENCE_LENGTHS = [16, 64, 128, 512, 1024, 2048, 4096, 8192]
+DEFAULT_SEQUENCE_LENGTHS = [128, 512, 1024, 2048, 4096, 8192]

 # Fixed hidden size and dtype for all benchmarks
 HIDDEN_SIZE = 8192
@@ -84,7 +81,6 @@ class CommunicatorBenchmark:
        self.symm_mem_comm = None
        self.symm_mem_comm_multimem = None
        self.symm_mem_comm_two_shot = None
-        self.fi_ar_comm = None

        self._init_communicators()

@@ -165,22 +161,6 @@ class CommunicatorBenchmark:
            )
            self.symm_mem_comm_two_shot = None

-        try:
-            self.fi_ar_comm = FlashInferAllReduce(
-                group=self.cpu_group,
-                device=self.device,
-            )
-            if not self.fi_ar_comm.disabled:
-                logger.info("Rank %s: FlashInferAllReduce initialized", self.rank)
-            else:
-                logger.info("Rank %s: FlashInferAllReduce disabled", self.rank)
-                self.fi_ar_comm = None
-        except Exception as e:
-            logger.warning(
-                "Rank %s: Failed to initialize FlashInferAllReduce: %s", self.rank, e
-            )
-            self.fi_ar_comm = None
-
    def benchmark_allreduce(
        self, sequence_length: int, num_warmup: int, num_trials: int
    ) -> dict[str, float]:
@@ -200,8 +180,7 @@ class CommunicatorBenchmark:
                    lambda t, c=comm: c.custom_all_reduce(t),
                    lambda t, c=comm: c.should_custom_ar(t),
                    comm.capture(),
-                    {"VLLM_CUSTOM_ALLREDUCE_ALGO": "1stage"},
-                    None,  # no destroy function
+                    "1stage",  # env variable value
                )
            )
            # CustomAllreduce two-shot
@@ -211,8 +190,7 @@ class CommunicatorBenchmark:
                    lambda t, c=comm: c.custom_all_reduce(t),
                    lambda t, c=comm: c.should_custom_ar(t),
                    comm.capture(),
-                    {"VLLM_CUSTOM_ALLREDUCE_ALGO": "2stage"},
-                    None,  # no destroy function
+                    "2stage",  # env variable value
                )
            )

@@ -224,8 +202,7 @@ class CommunicatorBenchmark:
                    lambda t, c=comm: c.all_reduce(t),
                    lambda t: True,  # Always available if initialized
                    nullcontext(),
-                    {},  # no env variable needed
-                    None,  # no destroy function
+                    None,  # no env variable needed
                )
            )
            communicators.append(
@@ -234,8 +211,7 @@ class CommunicatorBenchmark:
                    lambda t: torch.ops.vllm.all_reduce_symmetric_with_copy(t),
                    lambda t: True,  # Always available if initialized
                    nullcontext(),
-                    {},  # no env variable needed
-                    None,  # no destroy function
+                    None,  # no env variable needed
                )
            )

@@ -247,8 +223,7 @@ class CommunicatorBenchmark:
                    lambda t, c=comm: c.all_reduce(t),
                    lambda t, c=comm: c.should_use_symm_mem(t),
                    nullcontext(),
-                    {},  # no env variable needed
-                    None,  # no destroy function
+                    None,  # no env variable needed
                )
            )

@@ -260,67 +235,29 @@ class CommunicatorBenchmark:
                    lambda t, c=comm: c.all_reduce(t),
                    lambda t, c=comm: c.should_use_symm_mem(t),
                    nullcontext(),
-                    {},  # no env variable needed
-                    None,  # no destroy function needed
-                )
-            )
-
-        if self.fi_ar_comm is not None:
-            comm = self.fi_ar_comm
-            communicators.append(
-                (
-                    "flashinfer_trtllm",
-                    lambda t, c=comm: c.all_reduce(t),
-                    lambda t, c=comm: c.should_use_fi_ar(t),
-                    nullcontext(),
-                    {"VLLM_FLASHINFER_ALLREDUCE_BACKEND": "trtllm"},
-                    lambda c=comm: c.destroy(),
-                )
-            )
-            communicators.append(
-                (
-                    "flashinfer_mnnvl",
-                    lambda t, c=comm: c.all_reduce(t),
-                    lambda t, c=comm: c.should_use_fi_ar(t),
-                    nullcontext(),
-                    {"VLLM_FLASHINFER_ALLREDUCE_BACKEND": "mnnvl"},
-                    lambda c=comm: c.destroy(),
+                    None,  # no env variable needed
                )
            )

        # Benchmark each communicator
-        for (
-            name,
-            allreduce_fn,
-            should_use_fn,
-            context,
-            env_dict,
-            destroy_fn,
-        ) in communicators:
-            # Save original values and apply new environment variables
-            saved_env = {key: os.environ.get(key) for key in env_dict}
-            for key, value in env_dict.items():
-                os.environ[key] = value
-            try:
-                latency = self.benchmark_allreduce_single(
-                    sequence_length,
-                    allreduce_fn,
-                    should_use_fn,
-                    context,
-                    num_warmup,
-                    num_trials,
-                )
-                if latency is not None:
-                    results[name] = latency
-            finally:
-                if destroy_fn is not None:
-                    destroy_fn()
-                # Restore environment variables to their original state
-                for key, original_value in saved_env.items():
-                    if original_value is None:
-                        os.environ.pop(key, None)
-                    else:
-                        os.environ[key] = original_value
+        for name, allreduce_fn, should_use_fn, context, env_var in communicators:
+            # Set environment variable if needed
+            if env_var is not None:
+                os.environ["VLLM_CUSTOM_ALLREDUCE_ALGO"] = env_var
+            else:
+                # Clear the environment variable to avoid interference
+                os.environ.pop("VLLM_CUSTOM_ALLREDUCE_ALGO", None)
+
+            latency = self.benchmark_allreduce_single(
+                sequence_length,
+                allreduce_fn,
+                should_use_fn,
+                context,
+                num_warmup,
+                num_trials,
+            )
+            if latency is not None:
+                results[name] = latency

        return results

@@ -342,7 +279,7 @@ class CommunicatorBenchmark:
            if not should_use_fn(tensor):
                return None

-            torch.accelerator.synchronize()
+            torch.cuda.synchronize()
            stream = torch.cuda.Stream()
            with torch.cuda.stream(stream):
                graph_input = tensor.clone()
@@ -360,17 +297,17 @@ class CommunicatorBenchmark:
                        for _ in range(CUDA_GRAPH_CAPTURE_CYCLES):
                            allreduce_fn(graph_input)

-            torch.accelerator.synchronize()
+            torch.cuda.synchronize()
            for _ in range(num_warmup):
                graph.replay()
-            torch.accelerator.synchronize()
+            torch.cuda.synchronize()

-            torch.accelerator.synchronize()
+            torch.cuda.synchronize()
            start_time = time.perf_counter()

            for _ in range(num_trials):
                graph.replay()
-            torch.accelerator.synchronize()
+            torch.cuda.synchronize()

            end_time = time.perf_counter()

@@ -5,11 +5,8 @@
 Benchmark for FlashInfer fused collective operations vs standard operations.

 This benchmark compares:
-1. FlashInfer's allreduce_fusion with trtllm backend
-   (fused allreduce + rmsnorm + optional FP8/FP4 quant)
-2. FlashInfer's allreduce_fusion with mnnvl backend
-   (fused allreduce + rmsnorm only, no quantization support)
-3. Standard tensor_model_parallel_all_reduce + separate rmsnorm/quant operations
+1. FlashInfer's allreduce_fusion (fused allreduce + rmsnorm + optional quant)
+2. Standard tensor_model_parallel_all_reduce + separate rmsnorm/quant operations

 Usage with torchrun:
    torchrun --nproc_per_node=2 benchmark_fused_collective.py
@@ -51,12 +48,8 @@ SCALED_FP4_QUANT_OP = torch.ops._C.scaled_fp4_quant
 logger = init_logger(__name__)

 # Try to import FlashInfer
-TorchDistBackend = None
 try:
    import flashinfer.comm as flashinfer_comm  # type: ignore
-    from flashinfer.comm.mnnvl import (  # type: ignore
-        TorchDistBackend,
-    )

    if not (
        hasattr(flashinfer_comm, "allreduce_fusion")
@@ -81,15 +74,11 @@ _FI_MAX_SIZES = {
    8: 64 * MiB,  # 64MB
 }

-# Global workspace tensors for FlashInfer (keyed by backend name)
-_FI_WORKSPACES: dict = {}
-
-# Backends to benchmark
-FLASHINFER_BACKENDS = ["trtllm", "mnnvl"]
+# Global workspace tensor for FlashInfer
+_FI_WORKSPACE = None


 def setup_flashinfer_workspace(
-    backend: str,
    world_size: int,
    rank: int,
    hidden_dim: int,
@@ -97,54 +86,41 @@ def setup_flashinfer_workspace(
    dtype: torch.dtype,
 ):
    """Setup FlashInfer workspace for fused allreduce operations."""
-    global FI_WORKSPACES
+    global _FI_WORKSPACE

    if flashinfer_comm is None:
-        return None
+        return None, None

    if world_size not in _FI_MAX_SIZES:
        logger.warning("FlashInfer not supported for world size %s", world_size)
-        return None
+        return None, None

    try:
-        kwargs = {}
-        if TorchDistBackend is not None:
-            kwargs["comm_backend"] = TorchDistBackend(group=dist.group.WORLD)
-
        workspace = flashinfer_comm.create_allreduce_fusion_workspace(
-            backend=backend,
+            backend="trtllm",
            world_size=world_size,
            rank=rank,
            max_token_num=max_token_num,
            hidden_dim=hidden_dim,
            dtype=dtype,
-            **kwargs,
        )

-        _FI_WORKSPACES[backend] = workspace
+        _FI_WORKSPACE = workspace
        return workspace
    except Exception as e:
-        logger.error(
-            "Failed to setup FlashInfer workspace (backend=%s): %s", backend, e
-        )
+        logger.error("Failed to setup FlashInfer workspace: %s", e)
        return None


-def cleanup_flashinfer_workspaces():
-    """Cleanup all FlashInfer workspaces."""
-    if flashinfer_comm is None:
+def cleanup_flashinfer_workspace(workspace):
+    """Cleanup FlashInfer workspace."""
+    if flashinfer_comm is None or workspace is None:
        return

-    for backend, workspace in _FI_WORKSPACES.items():
-        try:
-            workspace.destroy()
-        except Exception as e:
-            logger.error(
-                "Failed to cleanup FlashInfer workspace (backend=%s): %s",
-                backend,
-                e,
-            )
-    _FI_WORKSPACES.clear()
+    try:
+        workspace.destroy()
+    except Exception as e:
+        logger.error("Failed to cleanup FlashInfer workspace: %s", e)


 class FlashInferFusedAllReduceParams:
@@ -158,7 +134,7 @@ class FlashInferFusedAllReduceParams:
        self.fp32_acc = True
        self.max_token_num = max_token_num

-    def get_flashinfer_fused_allreduce_kwargs(self):
+    def get_trtllm_fused_allreduce_kwargs(self):
        return {
            "launch_with_pdl": self.launch_with_pdl,
            "fp32_acc": self.fp32_acc,
@@ -171,12 +147,11 @@ def flashinfer_fused_allreduce_rmsnorm(
    rms_gamma: torch.Tensor,
    rms_eps: float,
    allreduce_params: "FlashInferFusedAllReduceParams",
-    workspace: object,
    use_oneshot: bool,
    norm_out: torch.Tensor | None = None,
 ):
    """FlashInfer fused allreduce + rmsnorm operation."""
-    if flashinfer_comm is None or workspace is None:
+    if flashinfer_comm is None or _FI_WORKSPACE is None:
        raise RuntimeError("FlashInfer not available or workspace not initialized")

    if norm_out is None:
@@ -185,13 +160,9 @@ def flashinfer_fused_allreduce_rmsnorm(
    else:
        residual_out = input_tensor

-    layout_code = None
-    if workspace.backend == "trtllm":
-        layout_code = flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4
-
    flashinfer_comm.allreduce_fusion(
        input=input_tensor,
-        workspace=workspace,
+        workspace=_FI_WORKSPACE,
        pattern=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNorm,
        residual_in=residual,
        residual_out=residual_out,
@@ -200,10 +171,10 @@ def flashinfer_fused_allreduce_rmsnorm(
        rms_eps=rms_eps,
        quant_out=None,
        scale_out=None,
-        layout_code=layout_code,
+        layout_code=flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4,
        scale_factor=None,
        use_oneshot=use_oneshot,
-        **allreduce_params.get_flashinfer_fused_allreduce_kwargs(),
+        **allreduce_params.get_trtllm_fused_allreduce_kwargs(),
    )


@@ -214,16 +185,12 @@ def flashinfer_fused_allreduce_rmsnorm_fp8_quant(
    rms_eps: float,
    scale_factor: torch.Tensor,
    allreduce_params: FlashInferFusedAllReduceParams,
-    workspace: object,
    use_oneshot: bool = True,
    norm_out: torch.Tensor | None = None,
    quant_out: torch.Tensor | None = None,
 ):
-    """FlashInfer fused allreduce + rmsnorm + FP8 quantization.
-
-    Note: Only supported by the trtllm backend.
-    """
-    if flashinfer_comm is None or workspace is None:
+    """FlashInfer fused allreduce + rmsnorm + FP8 quantization."""
+    if flashinfer_comm is None or _FI_WORKSPACE is None:
        raise RuntimeError("FlashInfer not available or workspace not initialized")

    if norm_out is None:
@@ -234,7 +201,7 @@ def flashinfer_fused_allreduce_rmsnorm_fp8_quant(

    flashinfer_comm.allreduce_fusion(
        input=input_tensor,
-        workspace=workspace,
+        workspace=_FI_WORKSPACE,
        pattern=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP8Quant,
        residual_in=residual,
        residual_out=residual_out,
@@ -246,7 +213,7 @@ def flashinfer_fused_allreduce_rmsnorm_fp8_quant(
        layout_code=flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4,
        scale_factor=scale_factor,
        use_oneshot=use_oneshot,
-        **allreduce_params.get_flashinfer_fused_allreduce_kwargs(),
+        **allreduce_params.get_trtllm_fused_allreduce_kwargs(),
    )


@@ -257,17 +224,13 @@ def flashinfer_fused_allreduce_rmsnorm_fp4_quant(
    rms_eps: float,
    input_global_scale: torch.Tensor,
    allreduce_params: FlashInferFusedAllReduceParams,
-    workspace: object,
    quant_out: torch.Tensor,
    use_oneshot: bool,
    output_scale: torch.Tensor,
    norm_out: torch.Tensor | None = None,
 ):
-    """FlashInfer fused allreduce + rmsnorm + FP4 quantization.
-
-    Note: Only supported by the trtllm backend.
-    """
-    if flashinfer_comm is None or workspace is None:
+    """FlashInfer fused allreduce + rmsnorm + FP4 quantization."""
+    if flashinfer_comm is None or _FI_WORKSPACE is None:
        raise RuntimeError("FlashInfer not available or workspace not initialized")

    if norm_out is None:
@@ -278,7 +241,7 @@ def flashinfer_fused_allreduce_rmsnorm_fp4_quant(

    flashinfer_comm.allreduce_fusion(
        input=input_tensor,
-        workspace=workspace,
+        workspace=_FI_WORKSPACE,
        pattern=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP4Quant,
        residual_in=residual,
        residual_out=residual_out,
@@ -290,7 +253,7 @@ def flashinfer_fused_allreduce_rmsnorm_fp4_quant(
        layout_code=flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4,
        scale_factor=input_global_scale,
        use_oneshot=use_oneshot,
-        **allreduce_params.get_flashinfer_fused_allreduce_kwargs(),
+        **allreduce_params.get_trtllm_fused_allreduce_kwargs(),
    )


@@ -385,7 +348,7 @@ def benchmark_operation(
    # Warmup before graph capture
    for _ in range(warmup):
        operation_func(*args, **kwargs)
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    # Create CUDA graph
    graph = torch.cuda.CUDAGraph()
@@ -398,19 +361,19 @@ def benchmark_operation(
            operation_func(*args, **kwargs)

    # Graph warmup
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()
    for _ in range(warmup):
        graph.replay()

    # Benchmark with CUDA graph
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()
    start_time = time.perf_counter()

    for _ in range(trials // num_op_per_cudagraph):
        # operation_func(*args, **kwargs)
        graph.replay()

-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()
    end_time = time.perf_counter()

    avg_time_ms = ((end_time - start_time) / trials) * 1000
@@ -423,16 +386,13 @@ def run_benchmarks(
    dtype: torch.dtype,
    use_residual: bool,
    allreduce_params: FlashInferFusedAllReduceParams | None,
-    workspaces: dict,
    quant_modes: set[str],
    no_oneshot: bool,
 ):
    """Run all benchmarks for given configuration.

    Args:
-        allreduce_params: Shared parameters for FlashInfer fused allreduce.
-        workspaces: Dict mapping backend name ("trtllm", "mnnvl") to workspace.
-        quant_modes: Set of quantization modes: "none", "fp8", "fp4".
+        quant_mode: "none", "fp8_only", "fp4_only", or "all"
    """
    (
        input_tensor,
@@ -448,18 +408,18 @@ def run_benchmarks(

    rms_eps = 1e-6
    results = {}
+    vllm_fused_allreduce = VllmFusedAllreduce(hidden_dim, dtype)
    use_oneshot_options = [False] if no_oneshot else [True, False]

+    # Create RMSNorm and QuantFP8 layers once for native benchmarks
+
    if "none" in quant_modes:
        # Standard AllReduce + RMSNorm
-        # Re-create VllmFusedAllreduce per config so CustomOp binds the
-        # correct forward method (native vs custom kernel).
        for custom_op in ["-rms_norm", "+rms_norm"]:
            with set_current_vllm_config(
                VllmConfig(compilation_config=CompilationConfig(custom_ops=[custom_op]))
            ):
                try:
-                    vllm_fused_allreduce = VllmFusedAllreduce(hidden_dim, dtype)
                    suffix = (
                        "_custom_rms_norm" if "+" in custom_op else "_native_rms_norm"
                    )
@@ -478,7 +438,6 @@ def run_benchmarks(
            VllmConfig(compilation_config=CompilationConfig(custom_ops=["-rms_norm"]))
        ):
            try:
-                vllm_fused_allreduce = VllmFusedAllreduce(hidden_dim, dtype)
                standard_allreduce_rmsnorm_native_compiled = torch.compile(
                    vllm_fused_allreduce.allreduce_rmsnorm,
                    fullgraph=True,
@@ -494,11 +453,10 @@ def run_benchmarks(
                logger.error("Standard AllReduce+RMSNorm Native Compiled failed: %s", e)
                results["standard_allreduce_rmsnorm_native_compiled"] = float("inf")

-        # FlashInfer Fused AllReduce + RMSNorm (all backends)
-        for backend, workspace in workspaces.items():
+        # FlashInfer Fused AllReduce + RMSNorm Oneshot/Twoshot
+        if flashinfer_comm is not None and allreduce_params is not None:
            for use_oneshot in use_oneshot_options:
                suffix = "_oneshot" if use_oneshot else "_twoshot"
-                key = f"flashinfer_{backend}_fused_allreduce_rmsnorm{suffix}"
                try:
                    time_ms = benchmark_operation(
                        flashinfer_fused_allreduce_rmsnorm,
@@ -508,17 +466,14 @@ def run_benchmarks(
                        rms_gamma=rms_gamma,
                        rms_eps=rms_eps,
                        allreduce_params=allreduce_params,
-                        workspace=workspace,
                        use_oneshot=use_oneshot,
                    )
-                    results[key] = time_ms
+                    results[f"flashinfer_fused_allreduce_rmsnorm{suffix}"] = time_ms
                except Exception as e:
-                    logger.error(
-                        "FlashInfer (%s) Fused AllReduce+RMSNorm failed: %s",
-                        backend,
-                        e,
+                    logger.error("FlashInfer Fused AllReduce+RMSNorm failed: %s", e)
+                    results[f"flashinfer_fused_allreduce_rmsnorm{suffix}"] = float(
+                        "inf"
                    )
-                    results[key] = float("inf")

    if "fp8" in quant_modes:
        # Standard AllReduce + RMSNorm + FP8 Quant
@@ -527,7 +482,7 @@ def run_benchmarks(
                "_custom_rms_norm" if "+" in rms_norm_custom_op else "_native_rms_norm"
            )
            for quant_fp8_custom_op in ["-quant_fp8", "+quant_fp8"]:
-                op_suffix = suffix + (
+                suffix += (
                    "_custom_quant_fp8"
                    if "+" in quant_fp8_custom_op
                    else "_native_quant_fp8"
@@ -540,17 +495,16 @@ def run_benchmarks(
                    )
                ):
                    try:
-                        vllm_fused_allreduce = VllmFusedAllreduce(hidden_dim, dtype)
                        time_ms = benchmark_operation(
                            vllm_fused_allreduce.allreduce_rmsnorm_fp8_quant,
                            input_tensor,
                            residual=residual,
                            scale_factor=scale_fp8,
                        )
-                        results[f"standard_allreduce{op_suffix}"] = time_ms
+                        results[f"standard_allreduce{suffix}"] = time_ms
                    except Exception as e:
                        logger.error("Standard AllReduce+RMSNorm+FP8 failed: %s", e)
-                        results[f"standard_allreduce{op_suffix}"] = float("inf")
+                        results[f"standard_allreduce{suffix}"] = float("inf")

        # Standard AllReduce + RMSNorm + FP8 Quant Native Compiled
        with set_current_vllm_config(
@@ -561,7 +515,6 @@ def run_benchmarks(
            )
        ):
            try:
-                vllm_fused_allreduce = VllmFusedAllreduce(hidden_dim, dtype)
                standard_allreduce_rmsnorm_fp8_quant_native_compiled = torch.compile(
                    vllm_fused_allreduce.allreduce_rmsnorm_fp8_quant,
                    fullgraph=True,
@@ -584,12 +537,10 @@ def run_benchmarks(
                    "inf"
                )

-        # FlashInfer Fused AllReduce + RMSNorm + FP8 Quant (trtllm only)
-        if "trtllm" in workspaces:
-            trtllm_ws = workspaces["trtllm"]
+        # FlashInfer Fused AllReduce + RMSNorm + FP8 Quant Oneshot
+        if flashinfer_comm is not None and allreduce_params is not None:
            for use_oneshot in use_oneshot_options:
                suffix = "_oneshot" if use_oneshot else "_twoshot"
-                key = f"flashinfer_trtllm_fused_allreduce_rmsnorm_fp8_quant{suffix}"
                try:
                    time_ms = benchmark_operation(
                        flashinfer_fused_allreduce_rmsnorm_fp8_quant,
@@ -601,16 +552,19 @@ def run_benchmarks(
                        scale_factor=scale_fp8,
                        quant_out=quant_out_fp8,
                        allreduce_params=allreduce_params,
-                        workspace=trtllm_ws,
                        use_oneshot=use_oneshot,
                    )
-                    results[key] = time_ms
+                    results[f"flashinfer_fused_allreduce_rmsnorm_fp8_quant{suffix}"] = (
+                        time_ms
+                    )
                except Exception as e:
                    logger.error(
-                        "FlashInfer (trtllm) Fused AllReduce+RMSNorm+FP8 failed: %s",
+                        "FlashInfer Fused AllReduce+RMSNorm+FP8 Oneshot failed: %s",
                        e,
                    )
-                    results[key] = float("inf")
+                    results[f"flashinfer_fused_allreduce_rmsnorm_fp8_quant{suffix}"] = (
+                        float("inf")
+                    )

    if "fp4" in quant_modes and current_platform.has_device_capability(100):
        # Standard AllReduce + RMSNorm + FP4 Quant
@@ -626,7 +580,6 @@ def run_benchmarks(
                )
            ):
                try:
-                    vllm_fused_allreduce = VllmFusedAllreduce(hidden_dim, dtype)
                    time_ms = benchmark_operation(
                        vllm_fused_allreduce.allreduce_rmsnorm_fp4_quant,
                        input_tensor,
@@ -645,7 +598,6 @@ def run_benchmarks(
            VllmConfig(compilation_config=CompilationConfig(custom_ops=["-rms_norm"]))
        ):
            try:
-                vllm_fused_allreduce = VllmFusedAllreduce(hidden_dim, dtype)
                standard_allreduce_rmsnorm_fp4_quant_native_compiled = torch.compile(
                    vllm_fused_allreduce.allreduce_rmsnorm_fp4_quant,
                    fullgraph=True,
@@ -670,12 +622,10 @@ def run_benchmarks(
                    "inf"
                )

-        # FlashInfer Fused AllReduce + RMSNorm + FP4 Quant (trtllm only)
-        if "trtllm" in workspaces:
-            trtllm_ws = workspaces["trtllm"]
+        # FlashInfer Fused AllReduce + RMSNorm + FP4 Quant Oneshot
+        if flashinfer_comm is not None and allreduce_params is not None:
            for use_oneshot in use_oneshot_options:
                suffix = "_oneshot" if use_oneshot else "_twoshot"
-                key = f"flashinfer_trtllm_fused_allreduce_rmsnorm_fp4_quant{suffix}"
                try:
                    time_ms = benchmark_operation(
                        flashinfer_fused_allreduce_rmsnorm_fp4_quant,
@@ -686,18 +636,49 @@ def run_benchmarks(
                        rms_eps=rms_eps,
                        input_global_scale=scale_fp4,
                        allreduce_params=allreduce_params,
-                        workspace=trtllm_ws,
                        quant_out=fp4_quant_out,
                        output_scale=fp4_output_scale,
                        use_oneshot=use_oneshot,
                    )
-                    results[key] = time_ms
+                    results[f"flashinfer_fused_allreduce_rmsnorm_fp4_quant{suffix}"] = (
+                        time_ms
+                    )
                except Exception as e:
                    logger.error(
-                        "FlashInfer (trtllm) Fused AllReduce+RMSNorm+FP4 failed: %s",
+                        "FlashInfer Fused AllReduce+RMSNorm+FP4 Oneshot failed: %s",
                        e,
                    )
-                    results[key] = float("inf")
+                    results[f"flashinfer_fused_allreduce_rmsnorm_fp4_quant{suffix}"] = (
+                        float("inf")
+                    )
+
+        # FlashInfer Fused AllReduce + RMSNorm + FP4 Quant Two-shot
+        if flashinfer_comm is not None and allreduce_params is not None:
+            try:
+                time_ms = benchmark_operation(
+                    flashinfer_fused_allreduce_rmsnorm_fp4_quant,
+                    input_tensor,
+                    residual=residual,
+                    norm_out=norm_out,
+                    rms_gamma=rms_gamma,
+                    rms_eps=rms_eps,
+                    input_global_scale=scale_fp4,
+                    allreduce_params=allreduce_params,
+                    quant_out=fp4_quant_out,
+                    output_scale=fp4_output_scale,
+                    use_oneshot=False,
+                )
+                results["flashinfer_fused_allreduce_rmsnorm_fp4_quant_twoshot"] = (
+                    time_ms
+                )
+            except Exception as e:
+                logger.error(
+                    "FlashInfer Fused AllReduce+RMSNorm+FP4 Two-shot failed: %s",
+                    e,
+                )
+                results["flashinfer_fused_allreduce_rmsnorm_fp4_quant_twoshot"] = float(
+                    "inf"
+                )

    return results

@@ -1035,7 +1016,8 @@ def main():

    configs = list(itertools.product(args.num_tokens, dtypes, residual_options))

-    # Setup FlashInfer workspaces for all backends
+    # Setup FlashInfer workspace if available
+    workspace = None
    allreduce_params = None

    if flashinfer_comm is not None:
@@ -1050,17 +1032,15 @@ def main():
            args.hidden_dim * max_element_size
        )

-        for backend in FLASHINFER_BACKENDS:
-            setup_flashinfer_workspace(
-                backend=backend,
-                world_size=world_size,
-                rank=rank,
-                hidden_dim=args.hidden_dim,
-                max_token_num=max_num_token,
-                dtype=workspace_dtype,
-            )
+        workspace = setup_flashinfer_workspace(
+            world_size,
+            rank,
+            args.hidden_dim,
+            max_num_token,
+            dtype=workspace_dtype,
+        )

-        if _FI_WORKSPACES:
+        if workspace is not None:
            allreduce_params = FlashInferFusedAllReduceParams(
                max_token_num=max_num_token,
            )
@@ -1086,7 +1066,6 @@ def main():
                dtype,
                use_residual,
                allreduce_params,
-                workspaces=_FI_WORKSPACES,
                quant_modes=quant_modes,
                no_oneshot=args.no_oneshot,
            )
@@ -1125,13 +1104,11 @@ def main():

    finally:
        # Cleanup
-        cleanup_flashinfer_workspaces()
+        if workspace is not None:
+            cleanup_flashinfer_workspace(workspace)

        dist.barrier()


 if __name__ == "__main__":
-    from vllm.config import VllmConfig, set_current_vllm_config
-
-    with set_current_vllm_config(VllmConfig()):
-        main()
+    main()
@@ -9,15 +9,15 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk
 from tests.kernels.moe.utils import make_dummy_moe_config
 from vllm import _custom_ops as ops
 from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
-from vllm.model_executor.layers.fused_moe.all2all_utils import (
-    maybe_make_prepare_finalize,
-)
 from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config
 from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp8
 from vllm.model_executor.layers.fused_moe.fused_moe import (
    fused_experts,
    fused_topk,
 )
+from vllm.model_executor.layers.fused_moe.prepare_finalize import (
+    MoEPrepareAndFinalizeNoEP,
+)
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 from vllm.v1.worker.workspace import init_workspace_manager

@@ -131,22 +131,16 @@ def bench_run(
            w2_scale=w2_scale,
            per_act_token_quant=per_act_token,
        )
-        moe_config = make_dummy_moe_config(
-            num_experts=w2.shape[0],
-            hidden_dim=w2.shape[1],
-            intermediate_size_per_partition=w2.shape[2],
-            in_dtype=a.dtype,
-        )

-        fn = mk.FusedMoEKernel(
-            maybe_make_prepare_finalize(
-                moe=moe_config,
-                quant_config=quant_config,
-                allow_new_interface=True,
-                use_monolithic=False,
-            ),
+        fn = mk.FusedMoEModularKernel(
+            MoEPrepareAndFinalizeNoEP(),
            CutlassExpertsFp8(
-                moe_config=moe_config,
+                moe_config=make_dummy_moe_config(
+                    num_experts=w2.shape[0],
+                    hidden_dim=w2.shape[1],
+                    intermediate_size_per_partition=w2.shape[2],
+                    in_dtype=a.dtype,
+                ),
                quant_config=quant_config,
            ),
        )
@@ -169,22 +163,16 @@ def bench_run(
            w2_scale=w2_scale,
            per_act_token_quant=per_act_token,
        )
-        moe_config = make_dummy_moe_config(
-            num_experts=w2.shape[0],
-            hidden_dim=w2.shape[1],
-            intermediate_size_per_partition=w2.shape[2],
-            in_dtype=a.dtype,
-        )

-        fn = mk.FusedMoEKernel(
-            maybe_make_prepare_finalize(
-                moe=moe_config,
-                quant_config=quant_config,
-                allow_new_interface=True,
-                use_monolithic=False,
-            ),
+        fn = mk.FusedMoEModularKernel(
+            MoEPrepareAndFinalizeNoEP(),
            CutlassExpertsFp8(
-                moe_config=moe_config,
+                moe_config=make_dummy_moe_config(
+                    num_experts=w2.shape[0],
+                    hidden_dim=w2.shape[1],
+                    intermediate_size_per_partition=w2.shape[2],
+                    in_dtype=a.dtype,
+                ),
                quant_config=quant_config,
            ),
        )
@@ -224,7 +212,7 @@ def bench_run(
    def replay_graph(graph, num_repeats):
        for _ in range(num_repeats):
            graph.replay()
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()

    cutlass_stream = torch.cuda.Stream()
    cutlass_graph = torch.cuda.CUDAGraph()
@@ -239,7 +227,7 @@ def bench_run(
            topk_weights,
            topk_ids,
        )
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    triton_stream = torch.cuda.Stream()
    triton_graph = torch.cuda.CUDAGraph()
@@ -254,7 +242,7 @@ def bench_run(
            w2_scale,
            a_scale,
        )
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    min_run_time = 5
    num_warmup = 5
@@ -5,14 +5,12 @@ import time

 import torch

-from vllm.benchmarks.lib.utils import default_vllm_config
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, set_random_seed


@torch.inference_mode()
-@default_vllm_config()
 def main(
    num_tokens: int,
    hidden_size: int,
@@ -34,14 +32,14 @@ def main(
    residual = torch.randn_like(x) * scale if add_residual else None

    def run_cuda_benchmark(num_iters: int, profile: bool = False) -> float:
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()
        if profile:
            torch.cuda.cudart().cudaProfilerStart()
        start_time = time.perf_counter()

        for _ in range(num_iters):
            layer(x, residual)
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()

        end_time = time.perf_counter()
        if profile:
@@ -1035,7 +1035,7 @@ def bench_optype(
    # Run bench function so that _LORA_A_PTR_DICT and _LORA_B_PTR_DICT are set up
    for kwargs in kwargs_list:
        op_type.bench_fn()(**kwargs)
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    # Merge into a single kwargs and qualify arguments as ArgPool
    kwargs = {k: ArgPool([]) for k in kwargs_list[0]}
@@ -47,13 +47,13 @@ def benchmark_method(
    # Warmup
    for _ in range(num_warmup):
        _ = method(k_nope, k_pe)
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    # Benchmark
    start = time.perf_counter()
    for _ in range(num_iters):
        _ = method(k_nope, k_pe)
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()
    end = time.perf_counter()

    return (end - start) / num_iters * 1000  # Convert to ms
@@ -17,9 +17,6 @@ from ray.experimental.tqdm_ray import tqdm

 from vllm.model_executor.layers.fused_moe import fused_topk
 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
-from vllm.model_executor.layers.fused_moe.all2all_utils import (
-    maybe_make_prepare_finalize,
-)
 from vllm.model_executor.layers.fused_moe.config import (
    FusedMoEConfig,
    FusedMoEParallelConfig,
@@ -54,7 +51,7 @@ def clear_triton_cache():

    # Clear CUDA memory cache
    if torch.cuda.is_available():
-        torch.accelerator.empty_cache()
+        torch.cuda.empty_cache()

    # Try to clear Triton's runtime cache
    try:
@@ -245,33 +242,24 @@ def benchmark_config(

        deep_gemm_experts = None
        if use_deep_gemm:
-            moe_config = (
-                FusedMoEConfig(
-                    num_experts=num_experts,
-                    experts_per_token=topk,
-                    hidden_dim=hidden_size,
-                    intermediate_size_per_partition=shard_intermediate_size,
-                    num_local_experts=num_experts,
-                    num_logical_experts=num_experts,
-                    activation=MoEActivation.SILU,
-                    moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
-                    in_dtype=init_dtype,
-                    routing_method=RoutingMethodType.TopK,
-                    device="cuda",
-                ),
-            )
-            deep_gemm_experts = mk.FusedMoEKernel(
-                prepare_finalize=maybe_make_prepare_finalize(
-                    moe=moe_config,
-                    quant_config=quant_config,
-                    allow_new_interface=True,
-                    use_monolithic=False,
-                ),
+            deep_gemm_experts = mk.FusedMoEModularKernel(
+                prepare_finalize=MoEPrepareAndFinalizeNoEP(),
                fused_experts=TritonOrDeepGemmExperts(
-                    moe_config=moe_config,
+                    moe_config=FusedMoEConfig(
+                        num_experts=num_experts,
+                        experts_per_token=topk,
+                        hidden_dim=hidden_size,
+                        intermediate_size_per_partition=shard_intermediate_size,
+                        num_local_experts=num_experts,
+                        num_logical_experts=num_experts,
+                        activation=MoEActivation.SILU,
+                        moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
+                        in_dtype=init_dtype,
+                        routing_method=RoutingMethodType.TopK,
+                        device="cuda",
+                    ),
                    quant_config=quant_config,
                ),
-                inplace=not disable_inplace(),
            )

        with override_config(config):
@@ -281,16 +269,8 @@ def benchmark_config(

            inplace = not disable_inplace()
            if use_deep_gemm:
-                return deep_gemm_experts.apply(
-                    x,
-                    w1,
-                    w2,
-                    topk_weights,
-                    topk_ids,
-                    activation=MoEActivation.SILU,
-                    global_num_experts=num_experts,
-                    apply_router_weight_on_input=False,
-                    expert_map=False,
+                return deep_gemm_experts(
+                    x, w1, w2, topk_weights, topk_ids, inplace=inplace
                )
            return fused_experts(
                x,
@@ -304,19 +284,19 @@ def benchmark_config(

    # JIT compilation & warmup
    run()
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    # Capture 10 invocations with CUDA graph
    graph = torch.cuda.CUDAGraph()
    with torch.cuda.graph(graph):
        for _ in range(10):
            run()
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    # Warmup
    for _ in range(5):
        graph.replay()
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    start_event = torch.Event(enable_timing=True)
    end_event = torch.Event(enable_timing=True)
@@ -324,7 +304,7 @@ def benchmark_config(
    latencies: list[float] = []
    for i in range(num_iters):
        prepare(i)
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()

        start_event.record()
        graph.replay()
@@ -1,278 +0,0 @@
-#!/usr/bin/env python3
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Benchmark comparing old vs new default fused MoE configs.
-
-Runs the triton fused_moe kernel with three configurations for each scenario:
-  1. Tuned config (from JSON file, if available) — the target to match
-  2. Old default (the hardcoded defaults before this change)
-  3. New default (the improved defaults)
-
-Usage:
-    python benchmarks/kernels/benchmark_moe_defaults.py
-
-Produces a table showing kernel time (us) and speedup of new vs old defaults.
-"""
-
-import torch
-
-from vllm.model_executor.layers.fused_moe import fused_topk, override_config
-from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
-from vllm.model_executor.layers.fused_moe.fused_moe import (
-    fused_experts,
-    get_default_config,
-    get_moe_configs,
-)
-from vllm.platforms import current_platform
-from vllm.triton_utils import triton
-from vllm.utils.torch_utils import set_random_seed
-
-FP8_DTYPE = current_platform.fp8_dtype()
-
-
-def old_default_config(M, E, N, K, topk, dtype=None, block_shape=None):
-    """The original defaults before https://github.com/vllm-project/vllm/pull/34846,
-    for comparison."""
-    if dtype == "fp8_w8a8" and block_shape is not None:
-        return {
-            "BLOCK_SIZE_M": 64,
-            "BLOCK_SIZE_N": block_shape[0],
-            "BLOCK_SIZE_K": block_shape[1],
-            "GROUP_SIZE_M": 32,
-            "SPLIT_K": 1,
-            "num_warps": 4,
-            "num_stages": 3 if not current_platform.is_rocm() else 2,
-        }
-    elif M <= E:
-        return {
-            "BLOCK_SIZE_M": 16,
-            "BLOCK_SIZE_N": 32,
-            "BLOCK_SIZE_K": 64,
-            "GROUP_SIZE_M": 1,
-            "SPLIT_K": 1,
-        }
-    else:
-        return {
-            "BLOCK_SIZE_M": 64,
-            "BLOCK_SIZE_N": 64,
-            "BLOCK_SIZE_K": 32,
-            "GROUP_SIZE_M": 8,
-            "SPLIT_K": 1,
-        }
-
-
-def benchmark_config(
-    config,
-    M,
-    E,
-    N,
-    K,
-    topk,
-    dtype,
-    use_fp8=False,
-    block_shape=None,
-    num_iters=100,
-):
-    """Time a single kernel config. Returns kernel time in microseconds."""
-    init_dtype = torch.float16 if use_fp8 else dtype
-
-    a = torch.randn(M, K, device="cuda", dtype=init_dtype) / 10
-    w1 = torch.randn(E, 2 * N, K, device="cuda", dtype=init_dtype) / 10
-    w2 = torch.randn(E, K, N, device="cuda", dtype=init_dtype) / 10
-
-    w1_scale = None
-    w2_scale = None
-    a1_scale = None
-    a2_scale = None
-    if use_fp8:
-        if block_shape is not None:
-            bsn, bsk = block_shape
-            n_tiles_w1 = triton.cdiv(2 * N, bsn)
-            k_tiles_w1 = triton.cdiv(K, bsk)
-            n_tiles_w2 = triton.cdiv(K, bsn)
-            k_tiles_w2 = triton.cdiv(N, bsk)
-            w1_scale = torch.rand(
-                E, n_tiles_w1, k_tiles_w1, device="cuda", dtype=torch.float32
-            )
-            w2_scale = torch.rand(
-                E, n_tiles_w2, k_tiles_w2, device="cuda", dtype=torch.float32
-            )
-        else:
-            w1_scale = torch.rand(E, device="cuda", dtype=torch.float32)
-            w2_scale = torch.rand(E, device="cuda", dtype=torch.float32)
-        a1_scale = torch.rand(1, device="cuda", dtype=torch.float32)
-        a2_scale = torch.rand(1, device="cuda", dtype=torch.float32)
-        # Only weights are stored in fp8; activations stay in bf16/fp16
-        # and get dynamically quantized inside the kernel.
-        w1 = w1.to(FP8_DTYPE)
-        w2 = w2.to(FP8_DTYPE)
-
-    quant_config = FusedMoEQuantConfig.make(
-        quant_dtype=torch.float8_e4m3fn if use_fp8 else None,
-        w1_scale=w1_scale,
-        w2_scale=w2_scale,
-        a1_scale=a1_scale,
-        a2_scale=a2_scale,
-        block_shape=block_shape,
-    )
-
-    gating = torch.randn(M, E, device="cuda", dtype=torch.float32)
-
-    # Warmup
-    for _ in range(20):
-        with override_config(config):
-            topk_weights, topk_ids, _ = fused_topk(a, gating, topk, renormalize=True)
-            fused_experts(
-                a,
-                w1,
-                w2,
-                topk_weights,
-                topk_ids,
-                quant_config=quant_config,
-            )
-    torch.accelerator.synchronize()
-
-    # Benchmark
-    start = torch.cuda.Event(enable_timing=True)
-    end = torch.cuda.Event(enable_timing=True)
-    start.record()
-    for _ in range(num_iters):
-        with override_config(config):
-            topk_weights, topk_ids, _ = fused_topk(a, gating, topk, renormalize=True)
-            fused_experts(
-                a,
-                w1,
-                w2,
-                topk_weights,
-                topk_ids,
-                quant_config=quant_config,
-            )
-    end.record()
-    torch.accelerator.synchronize()
-    return start.elapsed_time(end) / num_iters * 1000  # ms -> us
-
-
-# Model configurations: (name, E, N, K, topk, dtype_str, use_fp8, block_shape)
-# N = moe_intermediate_size // tp_size (the value used in config file lookup)
-MODELS = [
-    # --- Few experts ---
-    ("Mixtral bf16", 8, 7168, 4096, 2, None, False, None),
-    ("Mixtral fp8", 8, 7168, 4096, 2, "fp8_w8a8", True, None),
-    # --- Many experts: real model shapes at tp=1 ---
-    # Qwen2-MoE-57B: E=60, topk=4, N=1408, K=2048
-    ("Qwen2-MoE bf16", 60, 1408, 2048, 4, None, False, None),
-    # DeepSeek-V2: E=64, topk=6, N=1407, K=4096
-    # (use 1408 to avoid odd alignment; real model is 1407)
-    ("DeepSeek-V2 bf16", 64, 1408, 4096, 6, None, False, None),
-    # OLMoE-7B: E=64, topk=8, N=2048, K=2048
-    ("OLMoE bf16", 64, 2048, 2048, 8, None, False, None),
-    # GLM-4-100B-A10B: E=128, topk=8, N=1408, K=4096
-    ("GLM-4-MoE bf16", 128, 1408, 4096, 8, None, False, None),
-    # Qwen3-30B-A3B: E=128, topk=8, N=768, K=2048
-    ("Qwen3-MoE bf16", 128, 768, 2048, 8, None, False, None),
-    # DeepSeek-V3 / MiMo-V2-Flash: E=256, topk=8, N=2048, K=7168
-    ("DeepSeek-V3 bf16", 256, 2048, 7168, 8, None, False, None),
-    # Qwen3.5-70B-A22B (Qwen3-Next): E=512, topk=10, N=512, K=2048
-    ("Qwen3-Next bf16", 512, 512, 2048, 10, None, False, None),
-    # E=128 N=1856 bf16
-    ("E128 N1856 bf16", 128, 1856, 4096, 8, None, False, None),
-    # E=256 N=512 bf16 (DS-V3 tp=4)
-    ("DS-V3 tp4 bf16", 256, 512, 7168, 8, None, False, None),
-    # E=512 N=512 bf16 (Qwen3-Next tp=1)
-    ("Qwen3-Next bf16", 512, 512, 2048, 10, None, False, None),
-    # E=512 N=256 bf16 (Qwen3-Next tp=2)
-    ("Qwen3-Next tp2", 512, 256, 2048, 10, None, False, None),
-    # --- FP8 block quant (many experts) ---
-    # DS-V3 tp=4: E=256, N=512, fp8 block
-    ("DS-V3 tp4 fp8blk", 256, 512, 7168, 8, "fp8_w8a8", True, [128, 128]),
-    # DS-V3 tp=8: E=256, N=256, fp8 block
-    ("DS-V3 tp8 fp8blk", 256, 256, 7168, 8, "fp8_w8a8", True, [128, 128]),
-    # Qwen3-Next tp=2 fp8 block
-    ("Qwen3-Next tp2 fp8blk", 512, 256, 2048, 10, "fp8_w8a8", True, [128, 128]),
-]
-
-BATCH_SIZES = [1, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096]
-
-
-def main():
-    set_random_seed(0)
-    torch.set_default_device("cuda")
-    dtype = torch.bfloat16
-
-    for name, E, N, K, topk, dtype_str, use_fp8, block_shape in MODELS:
-        print(f"\n{'=' * 90}")
-        print(f"  {name}  (E={E}, N={N}, K={K}, topk={topk})")
-        print(f"{'=' * 90}")
-
-        # Try to load tuned config
-        block_n = block_shape[0] if block_shape else None
-        block_k = block_shape[1] if block_shape else None
-        tuned = get_moe_configs(E, N, dtype_str, block_n, block_k)
-        has_tuned = tuned is not None
-        print(f"  Tuned config available: {has_tuned}")
-
-        hdr = (
-            f"{'Batch':>6} | {'Tuned (us)':>11} | {'Old (us)':>11} | "
-            f"{'New (us)':>11} | {'New/Old':>8} | {'New/Tuned':>10}"
-        )
-        print(f"  {hdr}")
-        print(f"  {'-' * len(hdr)}")
-
-        for M in BATCH_SIZES:
-            old_cfg = old_default_config(M, E, N, K, topk, dtype_str, block_shape)
-            new_cfg = get_default_config(M, E, N, K, topk, dtype_str, block_shape)
-
-            if has_tuned:
-                tuned_cfg = tuned[min(tuned.keys(), key=lambda x: abs(x - M))]
-                t_tuned = benchmark_config(
-                    tuned_cfg,
-                    M,
-                    E,
-                    N,
-                    K,
-                    topk,
-                    dtype,
-                    use_fp8=use_fp8,
-                    block_shape=block_shape,
-                )
-            else:
-                t_tuned = None
-
-            t_old = benchmark_config(
-                old_cfg,
-                M,
-                E,
-                N,
-                K,
-                topk,
-                dtype,
-                use_fp8=use_fp8,
-                block_shape=block_shape,
-            )
-            t_new = benchmark_config(
-                new_cfg,
-                M,
-                E,
-                N,
-                K,
-                topk,
-                dtype,
-                use_fp8=use_fp8,
-                block_shape=block_shape,
-            )
-
-            ratio_new_old = t_new / t_old
-            tuned_str = f"{t_tuned:11.2f}" if t_tuned else f"{'N/A':>11}"
-            ratio_tuned = f"{t_new / t_tuned:10.2f}x" if t_tuned else f"{'N/A':>10}"
-            # flag regressions where new default is >5% slower than old
-            marker = " <--" if ratio_new_old > 1.05 else ""
-
-            print(
-                f"  {M:>6} | {tuned_str} | {t_old:11.2f} | {t_new:11.2f} "
-                f"| {ratio_new_old:7.2f}x | {ratio_tuned}{marker}"
-            )
-
-
-if __name__ == "__main__":
-    main()
@@ -69,19 +69,19 @@ def benchmark_permute(

    # JIT compilation & warmup
    run()
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    # Capture 10 invocations with CUDA graph
    graph = torch.cuda.CUDAGraph()
    with torch.cuda.graph(graph):
        for _ in range(10):
            run()
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    # Warmup
    for _ in range(5):
        graph.replay()
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    start_event = torch.Event(enable_timing=True)
    end_event = torch.Event(enable_timing=True)
@@ -89,7 +89,7 @@ def benchmark_permute(
    latencies: list[float] = []
    for i in range(num_iters):
        prepare(i)
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()

        start_event.record()
        graph.replay()
@@ -159,26 +159,26 @@ def benchmark_unpermute(
    # JIT compilation & warmup
    input = prepare()
    run(input)
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    # Capture 10 invocations with CUDA graph
    graph = torch.cuda.CUDAGraph()
    with torch.cuda.graph(graph):
        for _ in range(10):
            run(input)
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    # Warmup
    for _ in range(5):
        graph.replay()
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    start_event = torch.Event(enable_timing=True)
    end_event = torch.Event(enable_timing=True)

    latencies: list[float] = []
    for i in range(num_iters):
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()
        start_event.record()
        graph.replay()
        end_event.record()
@@ -36,7 +36,6 @@ from typing import Any
 import numpy as np
 import torch

-from vllm.benchmarks.lib.utils import default_vllm_config
 from vllm.model_executor.layers.rotary_embedding import get_rope
 from vllm.transformers_utils.config import get_config
 from vllm.utils.argparse_utils import FlexibleArgumentParser
@@ -79,7 +78,6 @@ def calculate_stats(times: list[float]) -> dict[str, float]:
    }


-@default_vllm_config()
 def benchmark_mrope(
    model_name: str,
    num_tokens: int,
@@ -135,14 +133,14 @@ def benchmark_mrope(
            key.clone(),
        )

-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    # Time reference implementation
    torch_times = []
    for _ in range(benchmark_iter):
        query_clone = query.clone()
        key_clone = key.clone()
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()
        start_time = time.time()

        mrope_helper_class.forward_native(
@@ -151,7 +149,7 @@ def benchmark_mrope(
            key_clone,
        )

-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()
        torch_times.append(time.time() - start_time)

    # Time triton kernel implementation
@@ -159,14 +157,14 @@ def benchmark_mrope(
    for _ in range(benchmark_iter):
        query_clone = query.clone()
        key_clone = key.clone()
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()
        start_time = time.time()
        mrope_helper_class.forward_cuda(
            positions,
            query_clone,
            key_clone,
        )
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()
        triton_times.append(time.time() - start_time)

    # Calculate statistics
@@ -103,7 +103,7 @@ def main(
        max_logits = torch.empty_like(exp_sums)

    def run_cuda_benchmark(num_iters: int, profile: bool = False) -> float:
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()
        if profile:
            torch.cuda.cudart().cudaProfilerStart()
        start_time = time.perf_counter()
@@ -173,7 +173,7 @@ def main(
                    )
            else:
                raise ValueError(f"Invalid version: {version}")
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()

        end_time = time.perf_counter()
        if profile:
@@ -28,7 +28,7 @@ def _time_cuda(
    # warmup
    for _ in range(warmup_iters):
        fn()
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    start = torch.Event(enable_timing=True)
    end = torch.Event(enable_timing=True)
@@ -37,7 +37,7 @@ def _time_cuda(
    for _ in range(bench_iters):
        fn()
    end.record()
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    return start.elapsed_time(end) / bench_iters  # ms/iter

@@ -29,7 +29,7 @@ def main(
    scale = torch.randn(1, 1, dtype=torch.float32) if static_scale else None

    def run_cuda_benchmark(num_iters: int, profile: bool = False) -> float:
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()
        if profile:
            torch.cuda.cudart().cudaProfilerStart()
        start_time = time.perf_counter()
@@ -39,7 +39,7 @@ def main(
                ops.scaled_int8_quant(x, scale)
            else:
                ops.scaled_fp8_quant(x, scale)
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()

        end_time = time.perf_counter()
        if profile:
@@ -84,16 +84,16 @@ def run_benchmark(
        g = torch.cuda.CUDAGraph()
        with torch.cuda.graph(g):
            function_under_test()
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()
        function_under_test = lambda: g.replay()

    def run_cuda_benchmark(n_iters: int) -> float:
        nonlocal key, value, key_cache, value_cache, slot_mapping
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()
        start = time.perf_counter()
        for _ in range(n_iters):
            function_under_test()
-            torch.accelerator.synchronize()
+            torch.cuda.synchronize()
        end = time.perf_counter()
        return (end - start) / n_iters

@@ -104,7 +104,7 @@ def run_benchmark(

    # free tensors to mitigate OOM when sweeping
    del key, value, key_cache, value_cache, slot_mapping
-    torch.accelerator.empty_cache()
+    torch.cuda.empty_cache()

    return lat

@@ -109,16 +109,16 @@ def run_benchmark(
        g = torch.cuda.CUDAGraph()
        with torch.cuda.graph(g):
            function_under_test()
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()
        function_under_test = lambda: g.replay()

    def run_cuda_benchmark(n_iters: int) -> float:
        nonlocal key, value, key_cache, value_cache, slot_mapping
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()
        start = time.perf_counter()
        for _ in range(n_iters):
            function_under_test()
-            torch.accelerator.synchronize()
+            torch.cuda.synchronize()
        end = time.perf_counter()
        return (end - start) / n_iters

@@ -129,7 +129,7 @@ def run_benchmark(

    # free tensors to mitigate OOM when sweeping
    del key, value, key_cache, value_cache, slot_mapping
-    torch.accelerator.empty_cache()
+    torch.cuda.empty_cache()

    return lat

@@ -5,7 +5,6 @@ import itertools

 import torch

-from vllm.benchmarks.lib.utils import default_vllm_config
 from vllm.model_executor.layers.rotary_embedding import get_rope
 from vllm.triton_utils import triton
 from vllm.utils.argparse_utils import FlexibleArgumentParser
@@ -30,7 +29,6 @@ def get_benchmark(head_size, rotary_dim, is_neox_style, device):
            args={},
        )
    )
-    @default_vllm_config()
    def benchmark(batch_size, seq_len, num_heads, provider):
        dtype = torch.bfloat16
        max_position = 8192
@@ -251,7 +251,7 @@ def benchmark(
        kernel(
            y, tokens_per_expert, num_parallel_tokens=num_parallel_tokens, group_size=G
        )
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    start_event = torch.Event(enable_timing=True)
    end_event = torch.Event(enable_timing=True)
@@ -259,7 +259,7 @@ def benchmark(
    # Benchmark
    latencies: list[float] = []
    for _ in range(runs):
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()

        start_event.record()
        for i in range(iterations_per_run):
@@ -126,7 +126,7 @@ def benchmark_decode(
    )

    def time_fn(fn, warmup=10, trials=20):
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()
        start = torch.Event(enable_timing=True)
        end = torch.Event(enable_timing=True)
        times = []
@@ -136,7 +136,7 @@ def benchmark_decode(
            start.record()
            fn()
            end.record()
-            torch.accelerator.synchronize()
+            torch.cuda.synchronize()
            times.append(start.elapsed_time(end))  # ms
        return sum(times) / len(times), torch.std(torch.tensor(times))

@@ -138,7 +138,7 @@ def benchmark_prefill(
    )

    def time_fn(fn, warmup=10, trials=20):
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()
        start = torch.Event(enable_timing=True)
        end = torch.Event(enable_timing=True)
        times = []
@@ -148,7 +148,7 @@ def benchmark_prefill(
            start.record()
            fn()
            end.record()
-            torch.accelerator.synchronize()
+            torch.cuda.synchronize()
            times.append(start.elapsed_time(end))  # ms
        return sum(times) / len(times), torch.std(torch.tensor(times))

@@ -177,18 +177,18 @@ def benchmark_config(
    def run():
        w8a8_block_matmul(A, B, As, Bs, block_size, config, out_dtype)

-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()
    # JIT complication & warmup
    for _ in range(5):
        run()
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    start_event = torch.Event(enable_timing=True)
    end_event = torch.Event(enable_timing=True)

    latencies: list[float] = []
    for i in range(num_iters):
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()
        start_event.record()
        run()
        end_event.record()
@@ -35,7 +35,7 @@ def benchmark_shape(
    B = torch.randn((n, k), device="cuda", dtype=torch.bfloat16)

    # Reference result in BF16
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()
    C_ref = A @ B.t()

    # Pre-quantize B for all implementations
@@ -121,14 +121,14 @@ def benchmark_shape(
        # Warmup
        for _ in range(warmup):
            func()
-            torch.accelerator.synchronize()
+            torch.cuda.synchronize()

        # Timing loop
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()
        start = time.time()
        for _ in range(repeat):
            func()
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()
        end = time.time()

        # Calculate timing and TFLOPS
@@ -7,7 +7,7 @@ First start serving your model
 ```bash
 export MODEL_PATH=/models/meta-llama/Meta-Llama-3.1-8B-Instruct/

-vllm serve $MODEL_PATH --served-model-name Llama
+vllm serve $MODEL_PATH --served-model-name Llama --disable-log-requests
 ```

 The variable `MODEL_PATH` should be a path to the model files (e.g. downloaded from huggingface).
@@ -13,16 +13,28 @@ endif()
 #
 # Define environment variables for special configurations
 #
-set(ENABLE_X86_ISA $ENV{VLLM_CPU_X86})
+set(ENABLE_AVX2 $ENV{VLLM_CPU_AVX2})
+set(ENABLE_AVX512 $ENV{VLLM_CPU_AVX512})
+set(ENABLE_AVX512BF16 $ENV{VLLM_CPU_AVX512BF16})
+set(ENABLE_AVX512VNNI $ENV{VLLM_CPU_AVX512VNNI})
+set(ENABLE_AMXBF16 $ENV{VLLM_CPU_AMXBF16})
 set(ENABLE_ARM_BF16 $ENV{VLLM_CPU_ARM_BF16})

 include_directories("${CMAKE_SOURCE_DIR}/csrc")

+
 set (ENABLE_NUMA TRUE)

 #
 # Check the compile flags
 #
+
+if (CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64")
+    list(APPEND CXX_COMPILE_FLAGS
+        "-mf16c"
+    )
+endif()
+
 if(MACOSX_FOUND)
    list(APPEND CXX_COMPILE_FLAGS
        "-DVLLM_CPU_EXTENSION")
@@ -66,6 +78,18 @@ function(check_sysctl TARGET OUT)
    endif()
 endfunction()

+
+function (is_avx512_disabled OUT)
+    set(DISABLE_AVX512 $ENV{VLLM_CPU_DISABLE_AVX512})
+    if(DISABLE_AVX512 AND DISABLE_AVX512 STREQUAL "true")
+        set(${OUT} ON PARENT_SCOPE)
+    else()
+        set(${OUT} OFF PARENT_SCOPE)
+    endif()
+endfunction()
+
+is_avx512_disabled(AVX512_DISABLED)
+
 if (MACOSX_FOUND AND CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64")
    message(STATUS "Apple Silicon Detected")
    set(APPLE_SILICON_FOUND TRUE)
@@ -73,6 +97,8 @@ if (MACOSX_FOUND AND CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64")
    check_sysctl(hw.optional.neon ASIMD_FOUND)
    check_sysctl(hw.optional.arm.FEAT_BF16 ARM_BF16_FOUND)
 else()
+    find_isa(${CPUINFO} "avx2" AVX2_FOUND)
+    find_isa(${CPUINFO} "avx512f" AVX512_FOUND)
    find_isa(${CPUINFO} "Power11" POWER11_FOUND)
    find_isa(${CPUINFO} "POWER10" POWER10_FOUND)
    find_isa(${CPUINFO} "POWER9" POWER9_FOUND)
@@ -82,32 +108,77 @@ else()
    find_isa(${CPUINFO} "v" RVV_FOUND) # Check for RISC-V RVV support

    # Support cross-compilation by allowing override via environment variables
+    if (ENABLE_AVX2)
+        set(AVX2_FOUND ON)
+        message(STATUS "AVX2 support enabled via VLLM_CPU_AVX2 environment variable")
+    endif()
+    if (ENABLE_AVX512)
+        set(AVX512_FOUND ON)
+        message(STATUS "AVX512 support enabled via VLLM_CPU_AVX512 environment variable")
+    endif()
    if (ENABLE_ARM_BF16)
        set(ARM_BF16_FOUND ON)
        message(STATUS "ARM BF16 support enabled via VLLM_CPU_ARM_BF16 environment variable")
    endif()
 endif()

-if (CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|amd64" OR ENABLE_X86_ISA)
-    set(ENABLE_X86_ISA ON)
-    if (NOT (CMAKE_CXX_COMPILER_ID STREQUAL "GNU" AND
-            CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL 12.3))
-        message(FATAL_ERROR "X86 backend requires gcc/g++ >= 12.3")
-    endif()
-    list(APPEND CXX_COMPILE_FLAGS "-mf16c")
-    list(APPEND CXX_COMPILE_FLAGS_AVX512 ${CXX_COMPILE_FLAGS})
-    list(APPEND CXX_COMPILE_FLAGS_AVX2 ${CXX_COMPILE_FLAGS})
-    list(APPEND CXX_COMPILE_FLAGS_AVX512
+if (AVX512_FOUND AND NOT AVX512_DISABLED)
+    list(APPEND CXX_COMPILE_FLAGS
        "-mavx512f"
        "-mavx512vl"
        "-mavx512bw"
-        "-mavx512dq"
-        "-mavx512bf16"
-        "-mavx512vnni"
-        "-mamx-bf16"
-        "-mamx-tile")
-    list(APPEND CXX_COMPILE_FLAGS_AVX2
-        "-mavx2")
+        "-mavx512dq")
+
+    find_isa(${CPUINFO} "avx512_bf16" AVX512BF16_FOUND)
+    if (AVX512BF16_FOUND OR ENABLE_AVX512BF16)
+        if (CMAKE_CXX_COMPILER_ID STREQUAL "GNU" AND
+            CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL 12.3)
+            list(APPEND CXX_COMPILE_FLAGS "-mavx512bf16")
+            set(ENABLE_AVX512BF16 ON)
+        else()
+            set(ENABLE_AVX512BF16 OFF)
+            message(WARNING "Disable AVX512-BF16 ISA support, requires gcc/g++ >= 12.3")
+        endif()
+    else()
+        set(ENABLE_AVX512BF16 OFF)
+        message(WARNING "Disable AVX512-BF16 ISA support, no avx512_bf16 found in local CPU flags." " If cross-compilation is required, please set env VLLM_CPU_AVX512BF16=1.")
+    endif()
+
+    find_isa(${CPUINFO} "avx512_vnni" AVX512VNNI_FOUND)
+    if (AVX512VNNI_FOUND OR ENABLE_AVX512VNNI)
+        if (CMAKE_CXX_COMPILER_ID STREQUAL "GNU" AND
+            CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL 12.3)
+            list(APPEND CXX_COMPILE_FLAGS "-mavx512vnni")
+            set(ENABLE_AVX512VNNI ON)
+        else()
+            set(ENABLE_AVX512VNNI OFF)
+            message(WARNING "Disable AVX512-VNNI ISA support, requires gcc/g++ >= 12.3")
+        endif()
+    else()
+        set(ENABLE_AVX512VNNI OFF)
+        message(WARNING "Disable AVX512-VNNI ISA support, no avx512_vnni found in local CPU flags." " If cross-compilation is required, please set env VLLM_CPU_AVX512VNNI=1.")
+    endif()
+
+    find_isa(${CPUINFO} "amx_bf16" AMXBF16_FOUND)
+    if (AMXBF16_FOUND OR ENABLE_AMXBF16)
+        if (CMAKE_CXX_COMPILER_ID STREQUAL "GNU" AND
+            CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL 12.3)
+            list(APPEND CXX_COMPILE_FLAGS "-mamx-bf16" "-mamx-tile")
+            set(ENABLE_AMXBF16 ON)
+            add_compile_definitions(-DCPU_CAPABILITY_AMXBF16)
+        else()
+            set(ENABLE_AMXBF16 OFF)
+            message(WARNING "Disable AMX_BF16 ISA support, requires gcc/g++ >= 12.3")
+        endif()
+    else()
+        set(ENABLE_AMXBF16 OFF)
+        message(WARNING "Disable AMX_BF16 ISA support, no amx_bf16 found in local CPU flags." " If cross-compilation is required, please set env VLLM_CPU_AMXBF16=1.")
+    endif()
+    
+elseif (AVX2_FOUND)
+    list(APPEND CXX_COMPILE_FLAGS "-mavx2")
+    message(WARNING "vLLM CPU backend using AVX2 ISA")
+    
 elseif (POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND)
    message(STATUS "PowerPC detected")
    if (POWER9_FOUND)
@@ -148,12 +219,12 @@ elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
        list(APPEND CXX_COMPILE_FLAGS "-march=rv64gc")
    endif()
 else()
-    message(FATAL_ERROR "vLLM CPU backend requires X86, Power9+ ISA, S390X ISA, ARMv8 or RISC-V support.")
+    message(FATAL_ERROR "vLLM CPU backend requires AVX512, AVX2, Power9+ ISA, S390X ISA, ARMv8 or RISC-V support.")
 endif()


-# Build oneDNN for GEMM kernels
-if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND)
+# Build oneDNN for GEMM kernels (only for x86-AVX512 /ARM platforms)
+if ((AVX512_FOUND AND NOT AVX512_DISABLED) OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND)
    # Fetch and build Arm Compute Library (ACL) as oneDNN's backend for AArch64
    # TODO [fadara01]: remove this once ACL can be fetched and built automatically as a dependency of oneDNN
    set(ONEDNN_AARCH64_USE_ACL OFF CACHE BOOL "")
@@ -242,24 +313,13 @@ if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND
        )
    else()
        message(STATUS "Downloading oneDNN from GitHub")
-        if(ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
-            message(STATUS "aarch64 detected: using pinned oneDNN commit 9c5be1cc59e368aebf0909e6cf20f981ea61462a")
-            FetchContent_Declare(
-                oneDNN
-                GIT_REPOSITORY https://github.com/oneapi-src/oneDNN.git
-                GIT_TAG        9c5be1cc59e368aebf0909e6cf20f981ea61462a
-                GIT_PROGRESS   TRUE
-                GIT_SHALLOW    FALSE
-            )
-        else()
-            FetchContent_Declare(
-                oneDNN
-                GIT_REPOSITORY https://github.com/oneapi-src/oneDNN.git
-                GIT_TAG        v3.10
-                GIT_PROGRESS   TRUE
-                GIT_SHALLOW    TRUE
-            )
-        endif()
+        FetchContent_Declare(
+            oneDNN
+            GIT_REPOSITORY https://github.com/oneapi-src/oneDNN.git
+            GIT_TAG v3.10
+            GIT_PROGRESS TRUE
+            GIT_SHALLOW TRUE
+        )
    endif()

    set(ONEDNN_LIBRARY_TYPE "STATIC")
@@ -269,21 +329,13 @@ if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND
    set(ONEDNN_ENABLE_WORKLOAD "INFERENCE")
    set(ONEDNN_ENABLE_PRIMITIVE "MATMUL;REORDER")
    set(ONEDNN_BUILD_GRAPH "OFF")
-    set(ONEDNN_ENABLE_JIT_PROFILING "ON")
+    set(ONEDNN_ENABLE_JIT_PROFILING "OFF")
    set(ONEDNN_ENABLE_ITT_TASKS "OFF")
-    set(ONEDNN_ENABLE_MAX_CPU_ISA "ON")
-    set(ONEDNN_ENABLE_CPU_ISA_HINTS "ON")
-    set(ONEDNN_VERBOSE "ON")
+    set(ONEDNN_ENABLE_MAX_CPU_ISA "OFF")
+    set(ONEDNN_ENABLE_CPU_ISA_HINTS "OFF")
+    set(ONEDNN_VERBOSE "OFF")
    set(CMAKE_POLICY_DEFAULT_CMP0077 NEW)

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

-# TODO: Refactor this
-if (ENABLE_X86_ISA)
-    message(STATUS "CPU extension (AVX512) compile flags: ${CXX_COMPILE_FLAGS_AVX512}")
-    message(STATUS "CPU extension (AVX2) compile flags: ${CXX_COMPILE_FLAGS_AVX2}")
-else()
-    message(STATUS "CPU extension compile flags: ${CXX_COMPILE_FLAGS}")
-endif()
+message(STATUS "CPU extension compile flags: ${CXX_COMPILE_FLAGS}")

 if(ENABLE_NUMA)
    list(APPEND LIBS numa)
@@ -344,6 +390,25 @@ set(VLLM_EXT_SRC
    "csrc/cpu/cpu_attn.cpp"
    "csrc/cpu/torch_bindings.cpp")

+if (AVX512_FOUND AND NOT AVX512_DISABLED)
+    set(VLLM_EXT_SRC
+        "csrc/cpu/shm.cpp"
+        "csrc/cpu/cpu_wna16.cpp"
+        "csrc/cpu/cpu_fused_moe.cpp"
+        ${VLLM_EXT_SRC})
+    if (ENABLE_AVX512BF16 AND ENABLE_AVX512VNNI)
+        set(VLLM_EXT_SRC
+            "csrc/cpu/sgl-kernels/gemm.cpp"
+            "csrc/cpu/sgl-kernels/gemm_int8.cpp"
+            "csrc/cpu/sgl-kernels/gemm_fp8.cpp"
+            "csrc/cpu/sgl-kernels/moe.cpp"
+            "csrc/cpu/sgl-kernels/moe_int8.cpp"
+            "csrc/cpu/sgl-kernels/moe_fp8.cpp"
+            ${VLLM_EXT_SRC})
+        add_compile_definitions(-DCPU_CAPABILITY_AVX512)
+    endif()
+endif()
+
 if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
    set(VLLM_EXT_SRC
        "csrc/cpu/shm.cpp"
@@ -356,83 +421,21 @@ if(USE_ONEDNN)
        ${VLLM_EXT_SRC})
 endif()

-if (ENABLE_X86_ISA)
-    set(VLLM_EXT_SRC_AVX512
-        "csrc/cpu/sgl-kernels/gemm.cpp"
-        "csrc/cpu/sgl-kernels/gemm_int8.cpp"
-        "csrc/cpu/sgl-kernels/gemm_fp8.cpp"
-        "csrc/cpu/sgl-kernels/moe.cpp"
-        "csrc/cpu/sgl-kernels/moe_int8.cpp"
-        "csrc/cpu/sgl-kernels/moe_fp8.cpp"
-        "csrc/cpu/shm.cpp"
-        "csrc/cpu/cpu_wna16.cpp"
-        "csrc/cpu/cpu_fused_moe.cpp"
-        "csrc/cpu/utils.cpp"
-        "csrc/cpu/cpu_attn.cpp"
-        "csrc/cpu/dnnl_kernels.cpp"
-        "csrc/cpu/torch_bindings.cpp"
-        # TODO: Remove these files
-        "csrc/cpu/activation.cpp"
-        "csrc/cpu/layernorm.cpp"
-        "csrc/cpu/mla_decode.cpp"
-        "csrc/cpu/pos_encoding.cpp"
-        "csrc/moe/dynamic_4bit_int_moe_cpu.cpp") 
+message(STATUS "CPU extension source files: ${VLLM_EXT_SRC}")

-    set(VLLM_EXT_SRC_AVX2 
-        "csrc/cpu/utils.cpp"
-        "csrc/cpu/cpu_attn.cpp"
-        "csrc/cpu/torch_bindings.cpp"
-        # TODO: Remove these files
-        "csrc/cpu/activation.cpp"
-        "csrc/cpu/layernorm.cpp"
-        "csrc/cpu/mla_decode.cpp"
-        "csrc/cpu/pos_encoding.cpp"
-        "csrc/moe/dynamic_4bit_int_moe_cpu.cpp") 
+#
+# Define extension targets
+#

-    message(STATUS "CPU extension (AVX512) source files: ${VLLM_EXT_SRC_AVX512}")
-    message(STATUS "CPU extension (AVX2) source files: ${VLLM_EXT_SRC_AVX2}")
-
-    define_extension_target(
-        _C
-        DESTINATION vllm
-        LANGUAGE CXX
-        SOURCES ${VLLM_EXT_SRC_AVX512}
-        LIBRARIES ${LIBS}
-        COMPILE_FLAGS ${CXX_COMPILE_FLAGS_AVX512}
-        USE_SABI 3
-        WITH_SOABI
-    )
-
-    # For SGL kernels
-    target_compile_definitions(_C PRIVATE "-DCPU_CAPABILITY_AVX512")
-    # For AMX kernels
-    target_compile_definitions(_C PRIVATE "-DCPU_CAPABILITY_AMXBF16")
-
-    define_extension_target(
-        _C_AVX2
-        DESTINATION vllm
-        LANGUAGE CXX
-        SOURCES ${VLLM_EXT_SRC_AVX2}
-        LIBRARIES ${LIBS}
-        COMPILE_FLAGS ${CXX_COMPILE_FLAGS_AVX2}
-        USE_SABI 3
-        WITH_SOABI
-    )
-else()
-    message(STATUS "CPU extension source files: ${VLLM_EXT_SRC}")
-    #
-    # Define extension targets
-    #
-    define_extension_target(
-        _C
-        DESTINATION vllm
-        LANGUAGE CXX
-        SOURCES ${VLLM_EXT_SRC}
-        LIBRARIES ${LIBS}
-        COMPILE_FLAGS ${CXX_COMPILE_FLAGS}
-        USE_SABI 3
-        WITH_SOABI
-    )
-endif()
+define_extension_target(
+    _C
+    DESTINATION vllm
+    LANGUAGE CXX
+    SOURCES ${VLLM_EXT_SRC}
+    LIBRARIES ${LIBS}
+    COMPILE_FLAGS ${CXX_COMPILE_FLAGS}
+    USE_SABI 3
+    WITH_SOABI
+)

 message(STATUS "Enabling C extension.")
@@ -17,8 +17,7 @@ endif()
 # They should be identical but if they aren't, this is a massive footgun.
 #
 # The vllm-flash-attn install rules are nested under vllm to make sure the library gets installed in the correct place.
-# To only install vllm-flash-attn, use --component _vllm_fa2_C (for FA2), --component _vllm_fa3_C (for FA3),
-# or --component _vllm_fa4_cutedsl_C (for FA4 CuteDSL Python files).
+# To only install vllm-flash-attn, use --component _vllm_fa2_C (for FA2) or --component _vllm_fa3_C (for FA3).
 # If no component is specified, vllm-flash-attn is still installed.

 # If VLLM_FLASH_ATTN_SRC_DIR is set, vllm-flash-attn is installed from that directory instead of downloading.
@@ -39,16 +38,22 @@ else()
  FetchContent_Declare(
          vllm-flash-attn
          GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
-          GIT_TAG 140c00c0241bb60cc6e44e7c1be9998d4b20d8d2
+          GIT_TAG 5824e6e2008271063c3229ab3e7032bd74abbbc6
          GIT_PROGRESS TRUE
          # Don't share the vllm-flash-attn build between build types
          BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
  )
 endif()

+
+# Ensure the vllm/vllm_flash_attn directory exists before installation
+install(CODE "file(MAKE_DIRECTORY \"\${CMAKE_INSTALL_PREFIX}/vllm/vllm_flash_attn\")" ALL_COMPONENTS)
+
 # Make sure vllm-flash-attn install rules are nested under vllm/
-# ALL_COMPONENTS ensures the save/modify/restore runs exactly once regardless
-# of how many components are being installed, avoiding double-append of /vllm/.
+# This is here to support installing all components under the same prefix with cmake --install.
+# setup.py installs every component separately but uses the same prefix for all.
+# ALL_COMPONENTS is used to avoid duplication for FA2 and FA3,
+# and these statements don't hurt when installing neither component.
 install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY FALSE)" ALL_COMPONENTS)
 install(CODE "set(OLD_CMAKE_INSTALL_PREFIX \"\${CMAKE_INSTALL_PREFIX}\")" ALL_COMPONENTS)
 install(CODE "set(CMAKE_INSTALL_PREFIX \"\${CMAKE_INSTALL_PREFIX}/vllm/\")" ALL_COMPONENTS)
@@ -57,48 +62,22 @@ install(CODE "set(CMAKE_INSTALL_PREFIX \"\${CMAKE_INSTALL_PREFIX}/vllm/\")" ALL_
 FetchContent_MakeAvailable(vllm-flash-attn)
 message(STATUS "vllm-flash-attn is available at ${vllm-flash-attn_SOURCE_DIR}")

-# Restore the install prefix after FA's install rules
+# Restore the install prefix
 install(CODE "set(CMAKE_INSTALL_PREFIX \"\${OLD_CMAKE_INSTALL_PREFIX}\")" ALL_COMPONENTS)
 install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY TRUE)" ALL_COMPONENTS)

-# Install shared Python files for both FA2 and FA3 components
-foreach(_FA_COMPONENT _vllm_fa2_C _vllm_fa3_C)
-  # Ensure the vllm/vllm_flash_attn directory exists before installation
-  install(CODE "file(MAKE_DIRECTORY \"\${CMAKE_INSTALL_PREFIX}/vllm/vllm_flash_attn\")"
-    COMPONENT ${_FA_COMPONENT})
+# Copy over the vllm-flash-attn python files (duplicated for fa2 and fa3, in
+# case only one is built, in the case both are built redundant work is done)
+install(
+  DIRECTORY ${vllm-flash-attn_SOURCE_DIR}/vllm_flash_attn/
+  DESTINATION vllm/vllm_flash_attn
+  COMPONENT _vllm_fa2_C
+  FILES_MATCHING PATTERN "*.py"
+)

-  # Copy vllm_flash_attn python files (except __init__.py and flash_attn_interface.py
-  # which are source-controlled in vllm)
-  install(
-    DIRECTORY ${vllm-flash-attn_SOURCE_DIR}/vllm_flash_attn/
-    DESTINATION vllm/vllm_flash_attn
-    COMPONENT ${_FA_COMPONENT}
-    FILES_MATCHING PATTERN "*.py"
-    PATTERN "__init__.py" EXCLUDE
-    PATTERN "flash_attn_interface.py" EXCLUDE
-  )
-
-endforeach()
-
-#
-# FA4 CuteDSL component
-# This is a Python-only component that copies the flash_attn/cute directory
-# and transforms imports to match our package structure.
-#
-add_custom_target(_vllm_fa4_cutedsl_C)
-
-# Copy flash_attn/cute directory (needed for FA4) and transform imports
-# The cute directory uses flash_attn.cute imports internally, which we replace
-# with vllm.vllm_flash_attn.cute to match our package structure.
-install(CODE "
-  file(GLOB_RECURSE CUTE_PY_FILES \"${vllm-flash-attn_SOURCE_DIR}/flash_attn/cute/*.py\")
-  foreach(SRC_FILE \${CUTE_PY_FILES})
-    file(RELATIVE_PATH REL_PATH \"${vllm-flash-attn_SOURCE_DIR}/flash_attn/cute\" \${SRC_FILE})
-    set(DST_FILE \"\${CMAKE_INSTALL_PREFIX}/vllm/vllm_flash_attn/cute/\${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.vllm_flash_attn.cute\" FILE_CONTENTS \"\${FILE_CONTENTS}\")
-    file(WRITE \${DST_FILE} \"\${FILE_CONTENTS}\")
-  endforeach()
-" COMPONENT _vllm_fa4_cutedsl_C)
+install(
+  DIRECTORY ${vllm-flash-attn_SOURCE_DIR}/vllm_flash_attn/
+  DESTINATION vllm/vllm_flash_attn
+  COMPONENT _vllm_fa3_C
+  FILES_MATCHING PATTERN "*.py"
+)
@@ -5,11 +5,115 @@
 #include <cmath>

 #include "cuda_compat.h"
-#include "cuda_vec_utils.cuh"
 #include "dispatch_utils.h"

 namespace vllm {

+struct alignas(32) u32x8_t {
+  uint32_t u0, u1, u2, u3, u4, u5, u6, u7;
+};
+
+__device__ __forceinline__ void ld256(u32x8_t& val, const u32x8_t* ptr) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000
+  asm volatile("ld.global.nc.v8.u32 {%0,%1,%2,%3,%4,%5,%6,%7}, [%8];\n"
+               : "=r"(val.u0), "=r"(val.u1), "=r"(val.u2), "=r"(val.u3),
+                 "=r"(val.u4), "=r"(val.u5), "=r"(val.u6), "=r"(val.u7)
+               : "l"(ptr));
+#else
+  const uint4* uint_ptr = reinterpret_cast<const uint4*>(ptr);
+  uint4 top_half = __ldg(&uint_ptr[0]);
+  uint4 bottom_half = __ldg(&uint_ptr[1]);
+  val.u0 = top_half.x;
+  val.u1 = top_half.y;
+  val.u2 = top_half.z;
+  val.u3 = top_half.w;
+  val.u4 = bottom_half.x;
+  val.u5 = bottom_half.y;
+  val.u6 = bottom_half.z;
+  val.u7 = bottom_half.w;
+#endif
+}
+
+__device__ __forceinline__ void st256(u32x8_t& val, u32x8_t* ptr) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000
+  asm volatile("st.global.v8.u32 [%0], {%1,%2,%3,%4,%5,%6,%7,%8};\n"
+               :
+               : "l"(ptr), "r"(val.u0), "r"(val.u1), "r"(val.u2), "r"(val.u3),
+                 "r"(val.u4), "r"(val.u5), "r"(val.u6), "r"(val.u7)
+               : "memory");
+#else
+  uint4* uint_ptr = reinterpret_cast<uint4*>(ptr);
+  uint_ptr[0] = make_uint4(val.u0, val.u1, val.u2, val.u3);
+  uint_ptr[1] = make_uint4(val.u4, val.u5, val.u6, val.u7);
+#endif
+}
+
+template <bool support_256>
+struct VecTraits;
+
+template <>
+struct VecTraits<true> {
+  static constexpr int ARCH_MAX_VEC_SIZE = 32;
+  using vec_t = u32x8_t;
+};
+
+template <>
+struct VecTraits<false> {
+  static constexpr int ARCH_MAX_VEC_SIZE = 16;
+  using vec_t = int4;
+};
+
+template <typename T>
+struct PackedTraits;
+
+template <>
+struct PackedTraits<c10::BFloat16> {
+  using packed_t = __nv_bfloat162;
+};
+
+template <>
+struct PackedTraits<c10::Half> {
+  using packed_t = __half2;
+};
+
+template <>
+struct PackedTraits<float> {
+  using packed_t = float2;
+};
+
+template <typename packed_t>
+__device__ __forceinline__ float2 cast_to_float2(const packed_t& val) {
+  if constexpr (std::is_same_v<packed_t, __nv_bfloat162>) {
+    return __bfloat1622float2(val);
+  } else if constexpr (std::is_same_v<packed_t, __half2>) {
+    return __half22float2(val);
+  } else if constexpr (std::is_same_v<packed_t, float2>) {
+    return float2(val);
+  }
+}
+
+template <typename packed_t>
+__device__ __forceinline__ packed_t cast_to_packed(const float2& val) {
+  if constexpr (std::is_same_v<packed_t, __nv_bfloat162>) {
+    return __float22bfloat162_rn(val);
+  } else if constexpr (std::is_same_v<packed_t, __half2>) {
+    return __float22half2_rn(val);
+  } else if constexpr (std::is_same_v<packed_t, float2>) {
+    return float2(val);
+  }
+}
+
+template <typename packed_t>
+__device__ __forceinline__ packed_t packed_mul(const packed_t& x,
+                                               const packed_t& y) {
+  if constexpr (std::is_same_v<packed_t, __nv_bfloat162> ||
+                std::is_same_v<packed_t, __half2>) {
+    return __hmul2(x, y);
+  } else if constexpr (std::is_same_v<packed_t, float2>) {
+    return make_float2(x.x * y.x, x.y * y.y);
+  }
+}
+
 template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
          bool act_first>
 __device__ __forceinline__ scalar_t compute(const scalar_t& x,
@@ -25,6 +129,16 @@ __device__ __forceinline__ packed_t packed_compute(const packed_t& x,
                   : packed_mul(x, PACKED_ACT_FN(y));
 }

+// Check if all pointers are 16-byte aligned for int4 vectorized access
+__host__ __device__ __forceinline__ bool is_16byte_aligned(const void* ptr) {
+  return (reinterpret_cast<uintptr_t>(ptr) & 15) == 0;
+}
+
+// Check if all pointers are 16-byte aligned for longlong4_32a vectorized access
+__host__ __device__ __forceinline__ bool is_32byte_aligned(const void* ptr) {
+  return (reinterpret_cast<uintptr_t>(ptr) & 31) == 0;
+}
+
 // Activation and gating kernel template.
 template <typename scalar_t, typename packed_t,
          scalar_t (*ACT_FN)(const scalar_t&),
@@ -39,32 +153,36 @@ __global__ void act_and_mul_kernel(
  scalar_t* out_ptr = out + blockIdx.x * d;

  if constexpr (use_vec) {
-    using cuda_t = typename CUDATypeConverter<scalar_t>::Type;
-    using pvec_t = PackedVec<cuda_t, use_256b>;
+    // Fast path: 128-bit/256-bit vectorized loop
+    using vec_t = typename VecTraits<use_256b>::vec_t;
+    constexpr int ARCH_MAX_VEC_SIZE = VecTraits<use_256b>::ARCH_MAX_VEC_SIZE;
+    constexpr int VEC_SIZE = ARCH_MAX_VEC_SIZE / sizeof(packed_t);

-    const pvec_t* x_vec = reinterpret_cast<const pvec_t*>(x_ptr);
-    const pvec_t* y_vec = reinterpret_cast<const pvec_t*>(y_ptr);
-    pvec_t* out_vec = reinterpret_cast<pvec_t*>(out_ptr);
-    const int num_vecs = d / 2 / pvec_t::NUM_ELTS;
+    const vec_t* x_vec = reinterpret_cast<const vec_t*>(x_ptr);
+    const vec_t* y_vec = reinterpret_cast<const vec_t*>(y_ptr);
+    vec_t* out_vec = reinterpret_cast<vec_t*>(out_ptr);
+    const int num_vecs = d / 2 / VEC_SIZE;

    for (int i = threadIdx.x; i < num_vecs; i += blockDim.x) {
-      pvec_t x, y;
+      vec_t x, y;
      if constexpr (use_256b) {
        ld256(x, &x_vec[i]);
        ld256(y, &y_vec[i]);
      } else {
-        ld128(x, &x_vec[i]);
-        ld128(y, &y_vec[i]);
+        x = VLLM_LDG(&x_vec[i]);
+        y = VLLM_LDG(&y_vec[i]);
      }
+      auto* xp = reinterpret_cast<packed_t*>(&x);
+      auto* yp = reinterpret_cast<packed_t*>(&y);
 #pragma unroll
-      for (int j = 0; j < pvec_t::NUM_ELTS; j++) {
-        x.elts[j] = packed_compute<packed_t, PACKED_ACT_FN, act_first>(
-            x.elts[j], y.elts[j]);
+      for (int j = 0; j < VEC_SIZE; j++) {
+        xp[j] =
+            packed_compute<packed_t, PACKED_ACT_FN, act_first>(xp[j], yp[j]);
      }
      if constexpr (use_256b) {
        st256(x, &out_vec[i]);
      } else {
-        st128(x, &out_vec[i]);
+        out_vec[i] = x;
      }
    }
  } else {
@@ -152,54 +270,51 @@ packed_gelu_tanh_kernel(const packed_t& val) {
 // Launch activation and gating kernel.
 // Use ACT_FIRST (bool) indicating whether to apply the activation function
 // first.
-#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST)        \
-  auto dtype = input.scalar_type();                                            \
-  int d = input.size(-1) / 2;                                                  \
-  int64_t num_tokens = input.numel() / input.size(-1);                         \
-  if (num_tokens == 0) {                                                       \
-    return;                                                                    \
-  }                                                                            \
-  dim3 grid(num_tokens);                                                       \
-  int cc_major = at::cuda::getCurrentDeviceProperties()->major;                \
-  int support_vec =                                                            \
-      (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128)            \
-          ? vllm::VecTraits<true>::ARCH_MAX_VEC_SIZE                           \
-          : vllm::VecTraits<false>::ARCH_MAX_VEC_SIZE;                         \
-  int vec_size = support_vec / at::elementSize(dtype);                         \
-  const bool use_vec = (d % vec_size == 0);                                    \
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));            \
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();                \
-  if (use_vec) {                                                               \
-    dim3 block(std::min(d / vec_size, 1024));                                  \
-    if (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) {         \
-      VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] {          \
-        vllm::act_and_mul_kernel<                                              \
-            scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,      \
-            KERNEL<scalar_t>,                                                  \
-            PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
-            ACT_FIRST, true, true><<<grid, block, 0, stream>>>(                \
-            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d);          \
-      });                                                                      \
-    } else {                                                                   \
-      VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] {          \
-        vllm::act_and_mul_kernel<                                              \
-            scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,      \
-            KERNEL<scalar_t>,                                                  \
-            PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
-            ACT_FIRST, true, false><<<grid, block, 0, stream>>>(               \
-            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d);          \
-      });                                                                      \
-    }                                                                          \
-  } else {                                                                     \
-    dim3 block(std::min(d, 1024));                                             \
-    VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] {            \
-      vllm::act_and_mul_kernel<                                                \
-          scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,        \
-          KERNEL<scalar_t>,                                                    \
-          PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>,   \
-          ACT_FIRST, false><<<grid, block, 0, stream>>>(                       \
-          out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d);            \
-    });                                                                        \
+#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST)     \
+  auto dtype = input.scalar_type();                                         \
+  int d = input.size(-1) / 2;                                               \
+  int64_t num_tokens = input.numel() / input.size(-1);                      \
+  if (num_tokens == 0) {                                                    \
+    return;                                                                 \
+  }                                                                         \
+  dim3 grid(num_tokens);                                                    \
+  int cc_major = at::cuda::getCurrentDeviceProperties()->major;             \
+  int support_vec = (cc_major >= 10 && num_tokens > 128) ? 32 : 16;         \
+  int vec_size = support_vec / at::elementSize(dtype);                      \
+  const bool use_vec = (d % vec_size == 0);                                 \
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));         \
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();             \
+  if (use_vec) {                                                            \
+    dim3 block(std::min(d / vec_size, 1024));                               \
+    if (cc_major >= 10 && num_tokens > 128) {                               \
+      VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] {       \
+        vllm::act_and_mul_kernel<                                           \
+            scalar_t, typename vllm::PackedTraits<scalar_t>::packed_t,      \
+            KERNEL<scalar_t>,                                               \
+            PACKED_KERNEL<typename vllm::PackedTraits<scalar_t>::packed_t>, \
+            ACT_FIRST, true, true><<<grid, block, 0, stream>>>(             \
+            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d);       \
+      });                                                                   \
+    } else {                                                                \
+      VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] {       \
+        vllm::act_and_mul_kernel<                                           \
+            scalar_t, typename vllm::PackedTraits<scalar_t>::packed_t,      \
+            KERNEL<scalar_t>,                                               \
+            PACKED_KERNEL<typename vllm::PackedTraits<scalar_t>::packed_t>, \
+            ACT_FIRST, true, false><<<grid, block, 0, stream>>>(            \
+            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d);       \
+      });                                                                   \
+    }                                                                       \
+  } else {                                                                  \
+    dim3 block(std::min(d, 1024));                                          \
+    VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] {         \
+      vllm::act_and_mul_kernel<                                             \
+          scalar_t, typename vllm::PackedTraits<scalar_t>::packed_t,        \
+          KERNEL<scalar_t>,                                                 \
+          PACKED_KERNEL<typename vllm::PackedTraits<scalar_t>::packed_t>,   \
+          ACT_FIRST, false><<<grid, block, 0, stream>>>(                    \
+          out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d);         \
+    });                                                                     \
  }

 void silu_and_mul(torch::Tensor& out,    // [..., d]
@@ -261,31 +376,35 @@ __global__ void act_and_mul_kernel_with_param(
  scalar_t* out_ptr = out + blockIdx.x * d;

  if constexpr (use_vec) {
-    using cuda_t = typename CUDATypeConverter<scalar_t>::Type;
-    using pvec_t = PackedVec<cuda_t, use_256b>;
+    // Fast path: 128-bit/256-bit vectorized loop
+    using vec_t = typename VecTraits<use_256b>::vec_t;
+    constexpr int ARCH_MAX_VEC_SIZE = VecTraits<use_256b>::ARCH_MAX_VEC_SIZE;
+    constexpr int VEC_SIZE = ARCH_MAX_VEC_SIZE / sizeof(packed_t);

-    const pvec_t* x_vec = reinterpret_cast<const pvec_t*>(x_ptr);
-    const pvec_t* y_vec = reinterpret_cast<const pvec_t*>(y_ptr);
-    pvec_t* out_vec = reinterpret_cast<pvec_t*>(out_ptr);
-    const int num_vecs = d / 2 / pvec_t::NUM_ELTS;
+    const vec_t* x_vec = reinterpret_cast<const vec_t*>(x_ptr);
+    const vec_t* y_vec = reinterpret_cast<const vec_t*>(y_ptr);
+    vec_t* out_vec = reinterpret_cast<vec_t*>(out_ptr);
+    const int num_vecs = d / 2 / VEC_SIZE;

    for (int i = threadIdx.x; i < num_vecs; i += blockDim.x) {
-      pvec_t x, y;
+      vec_t x, y;
      if constexpr (use_256b) {
        ld256(x, &x_vec[i]);
        ld256(y, &y_vec[i]);
      } else {
-        ld128(x, &x_vec[i]);
-        ld128(y, &y_vec[i]);
+        x = VLLM_LDG(&x_vec[i]);
+        y = VLLM_LDG(&y_vec[i]);
      }
+      auto* xp = reinterpret_cast<packed_t*>(&x);
+      auto* yp = reinterpret_cast<packed_t*>(&y);
 #pragma unroll
-      for (int j = 0; j < pvec_t::NUM_ELTS; j++) {
-        x.elts[j] = packed_mul(PACKED_ACT_FN(x.elts[j], param), y.elts[j]);
+      for (int j = 0; j < VEC_SIZE; j++) {
+        xp[j] = packed_mul(PACKED_ACT_FN(xp[j], param), yp[j]);
      }
      if constexpr (use_256b) {
        st256(x, &out_vec[i]);
      } else {
-        st128(x, &out_vec[i]);
+        out_vec[i] = x;
      }
    }
  } else {
@@ -378,24 +497,21 @@ __global__ void swigluoai_and_mul_kernel(
  }                                                                            \
  dim3 grid(num_tokens);                                                       \
  int cc_major = at::cuda::getCurrentDeviceProperties()->major;                \
-  int support_vec =                                                            \
-      (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128)            \
-          ? vllm::VecTraits<true>::ARCH_MAX_VEC_SIZE                           \
-          : vllm::VecTraits<false>::ARCH_MAX_VEC_SIZE;                         \
+  int support_vec = (cc_major >= 10 && num_tokens > 128) ? 32 : 16;            \
  int vec_size = support_vec / at::elementSize(dtype);                         \
  const bool use_vec = (d % vec_size == 0);                                    \
  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));            \
  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();                \
  if (use_vec) {                                                               \
    dim3 block(std::min(d / vec_size, 1024));                                  \
-    if (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) {         \
+    if (cc_major >= 10 && num_tokens > 128) {                                  \
      VLLM_DISPATCH_FLOATING_TYPES(                                            \
          dtype, "act_and_mul_kernel_with_param", [&] {                        \
            vllm::act_and_mul_kernel_with_param<                               \
-                scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,  \
+                scalar_t, typename vllm::PackedTraits<scalar_t>::packed_t,     \
                KERNEL<scalar_t>,                                              \
                PACKED_KERNEL<                                                 \
-                    typename vllm::PackedTypeConverter<scalar_t>::Type>,       \
+                    typename vllm::PackedTraits<scalar_t>::packed_t>,          \
                true, true><<<grid, block, 0, stream>>>(                       \
                out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d,       \
                PARAM);                                                        \
@@ -404,10 +520,10 @@ __global__ void swigluoai_and_mul_kernel(
      VLLM_DISPATCH_FLOATING_TYPES(                                            \
          dtype, "act_and_mul_kernel_with_param", [&] {                        \
            vllm::act_and_mul_kernel_with_param<                               \
-                scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,  \
+                scalar_t, typename vllm::PackedTraits<scalar_t>::packed_t,     \
                KERNEL<scalar_t>,                                              \
                PACKED_KERNEL<                                                 \
-                    typename vllm::PackedTypeConverter<scalar_t>::Type>,       \
+                    typename vllm::PackedTraits<scalar_t>::packed_t>,          \
                true, false><<<grid, block, 0, stream>>>(                      \
                out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d,       \
                PARAM);                                                        \
@@ -417,9 +533,9 @@ __global__ void swigluoai_and_mul_kernel(
    dim3 block(std::min(d, 1024));                                             \
    VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel_with_param", [&] { \
      vllm::act_and_mul_kernel_with_param<                                     \
-          scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,        \
+          scalar_t, typename vllm::PackedTraits<scalar_t>::packed_t,           \
          KERNEL<scalar_t>,                                                    \
-          PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>,   \
+          PACKED_KERNEL<typename vllm::PackedTraits<scalar_t>::packed_t>,      \
          false><<<grid, block, 0, stream>>>(                                  \
          out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, PARAM);     \
    });                                                                        \
@@ -511,17 +627,14 @@ __global__ void activation_kernel(
  }                                                                      \
  dim3 grid(num_tokens);                                                 \
  int cc_major = at::cuda::getCurrentDeviceProperties()->major;          \
-  int support_vec =                                                      \
-      (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128)      \
-          ? vllm::VecTraits<true>::ARCH_MAX_VEC_SIZE                     \
-          : vllm::VecTraits<false>::ARCH_MAX_VEC_SIZE;                   \
+  int support_vec = (cc_major >= 10 && num_tokens > 128) ? 32 : 16;      \
  int vec_size = support_vec / at::elementSize(dtype);                   \
  const bool use_vec = (d % vec_size == 0);                              \
  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));      \
  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();          \
  if (use_vec) {                                                         \
    dim3 block(std::min(d / vec_size, 1024));                            \
-    if (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) {   \
+    if (cc_major >= 10 && num_tokens > 128) {                            \
      VLLM_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] {     \
        vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, true, true>  \
            <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(),       \
@@ -1305,8 +1305,7 @@ void indexer_k_quant_and_cache(
  const at::cuda::OptionalCUDAGuard device_guard(device_of(k));
  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();

-  static const std::string kv_cache_dtype = "fp8_e4m3";
-  DISPATCH_BY_KV_CACHE_DTYPE(k.dtype(), kv_cache_dtype,
+  DISPATCH_BY_KV_CACHE_DTYPE(k.dtype(), "fp8_e4m3",
                             CALL_INDEXER_K_QUANT_AND_CACHE);
 }

@@ -16,8 +16,6 @@ torch::Tensor get_scheduler_metadata(
    isa = cpu_attention::ISA::VEC16;
  } else if (isa_hint == "neon") {
    isa = cpu_attention::ISA::NEON;
-  } else if (isa_hint == "vxe") {
-    isa = cpu_attention::ISA::VXE;
  } else {
    TORCH_CHECK(false, "Unsupported CPU attention ISA hint: " + isa_hint);
  }
@@ -102,8 +100,6 @@ void cpu_attn_reshape_and_cache(
      return cpu_attention::ISA::VEC16;
    } else if (isa == "neon") {
      return cpu_attention::ISA::NEON;
-    } else if (isa == "vxe") {
-      return cpu_attention::ISA::VXE;
    } else {
      TORCH_CHECK(false, "Invalid ISA type: " + isa);
    }
@@ -420,7 +420,7 @@ class AttentionImpl<ISA::AMX, scalar_t, head_dim> {
      const int64_t block_size, const int64_t block_size_stride) {
    // For AMX 2D tiles, size of each line is 64 bytes
    constexpr int64_t amx_tile_row_size = AMX_TILE_ROW_BYTES;
-    // For AMX B matrix, N always is 16
+    // For AMX B martix, N always is 16
    constexpr int64_t amx_b_tile_n_size = AMX_TILE_ROW_BYTES / 4;
    constexpr int64_t amx_b_tile_k_size = amx_tile_row_size / sizeof(scalar_t);
    // For now suppose block_size is divisible by amx_tile_column_num
@@ -12,7 +12,7 @@
 #include "cpu/utils.hpp"

 namespace cpu_attention {
-enum class ISA { AMX, VEC, VEC16, NEON, VXE };
+enum class ISA { AMX, VEC, VEC16, NEON };

 template <ISA isa, typename scalar_t, int64_t head_dim>
 class AttentionImpl {};
@@ -1,386 +0,0 @@
-#ifndef CPU_ATTN_VXE_HPP
-#define CPU_ATTN_VXE_HPP
-
-#include "cpu_attn_impl.hpp"
-#include <vecintrin.h>
-#include <type_traits>
-
-namespace cpu_attention {
-
-namespace {
-
-// s390x Vector = 16 bytes (128 bits)
-#define BLOCK_SIZE_ALIGNMENT 32
-#define HEAD_SIZE_ALIGNMENT 32
-#define MAX_Q_HEAD_NUM_PER_ITER 16
-
-template <typename kv_cache_t>
-FORCE_INLINE void load_row8_B_as_f32(const kv_cache_t* p, __vector float& b0,
-                                     __vector float& b1);
-
-// [1] Float Specialization
-template <>
-FORCE_INLINE void load_row8_B_as_f32<float>(const float* p, __vector float& b0,
-                                            __vector float& b1) {
-  // Explicitly cast to long long for offset, and float* for pointer
-  b0 = vec_xl((long long)0, const_cast<float*>(p));
-  b1 = vec_xl((long long)0, const_cast<float*>(p + 4));
-}
-
-// [2] BFloat16 Specialization (Big Endian Fix)
-template <>
-FORCE_INLINE void load_row8_B_as_f32<c10::BFloat16>(const c10::BFloat16* p,
-                                                    __vector float& b0,
-                                                    __vector float& b1) {
-  // 1. Load 8 BF16s (16 bytes) into one vector
-  // Explicit cast to unsigned short* for vec_xl to return vector unsigned short
-  __vector unsigned short raw = vec_xl((long long)0, (unsigned short*)p);
-
-  // 2. Prepare Zero vector
-  __vector unsigned short zeros = vec_splat_u16(0);
-
-  // 3. Merge High/Low to expand BF16 -> Float32
-  // On Big Endian, a float is [BF16_bits | 16_zero_bits]
-  b0 = (__vector float)vec_mergeh(raw, zeros);
-  b1 = (__vector float)vec_mergel(raw, zeros);
-}
-
-template <>
-FORCE_INLINE void load_row8_B_as_f32<c10::Half>(const c10::Half* p,
-                                                __vector float& b0,
-                                                __vector float& b1) {
-  alignas(16) float tmp[8];
-
-  // Manual unroll / conversion
-  tmp[0] = static_cast<float>(p[0]);
-  tmp[1] = static_cast<float>(p[1]);
-  tmp[2] = static_cast<float>(p[2]);
-  tmp[3] = static_cast<float>(p[3]);
-  tmp[4] = static_cast<float>(p[4]);
-  tmp[5] = static_cast<float>(p[5]);
-  tmp[6] = static_cast<float>(p[6]);
-  tmp[7] = static_cast<float>(p[7]);
-
-  // Explicit arguments for intrinsic: (long long offset, float* ptr)
-  b0 = vec_xl((long long)0, (float*)tmp);
-  b1 = vec_xl((long long)0, (float*)(tmp + 4));
-}
-
-template <int32_t M, typename kv_cache_t>
-FORCE_INLINE void gemm_micro_s390x_Mx8_Ku4(
-    const float* __restrict A,       // [M x K]
-    const kv_cache_t* __restrict B,  // [K x 8]
-    float* __restrict C,             // [M x 8]
-    int64_t lda, int64_t ldb, int64_t ldc, int32_t K, bool accumulate) {
-  static_assert(1 <= M && M <= 8, "M must be in [1,8]");
-
-// Helper macros to unroll codegen for M rows
-#define ROWS_APPLY(OP) OP(0) OP(1) OP(2) OP(3) OP(4) OP(5) OP(6) OP(7)
-#define IF_M(i) if constexpr (M > (i))
-
-  // 1. Define A pointers
-#define DECL_A(i) const float* a##i = A + (i) * lda;
-  ROWS_APPLY(DECL_A)
-#undef DECL_A
-
-  // 2. Define Accumulators (2 vectors covers 8 columns)
-#define DECL_ACC(i) __vector float acc##i##_0, acc##i##_1;
-  ROWS_APPLY(DECL_ACC)
-#undef DECL_ACC
-
-  // 3. Initialize Accumulators (Load C or Zero)
-#define INIT_ACC(i)                                                    \
-  IF_M(i) {                                                            \
-    if (accumulate) {                                                  \
-      acc##i##_0 =                                                     \
-          vec_xl((long long)0, const_cast<float*>(C + (i) * ldc + 0)); \
-      acc##i##_1 =                                                     \
-          vec_xl((long long)0, const_cast<float*>(C + (i) * ldc + 4)); \
-    } else {                                                           \
-      acc##i##_0 = vec_splats(0.0f);                                   \
-      acc##i##_1 = vec_splats(0.0f);                                   \
-    }                                                                  \
-  }
-  ROWS_APPLY(INIT_ACC)
-#undef INIT_ACC
-
-  int32_t k = 0;
-
-  for (; k + 3 < K; k += 4) {
-    // Load 4 values of A for each Row M: A[k...k+3]
-#define LOAD_A4(i)        \
-  __vector float a##i##v; \
-  IF_M(i) a##i##v = vec_xl((long long)0, const_cast<float*>(a##i + k));
-    ROWS_APPLY(LOAD_A4)
-#undef LOAD_A4
-
-    // Helper: FMA for specific lane L of A
-    // s390x: vec_madd(b, vec_splat(a, lane), acc)
-#define FMAS_LANE(i, aiv, L)                        \
-  IF_M(i) {                                         \
-    __vector float a_broad = vec_splat(aiv, L);     \
-    acc##i##_0 = vec_madd(b0, a_broad, acc##i##_0); \
-    acc##i##_1 = vec_madd(b1, a_broad, acc##i##_1); \
-  }
-
-    // Unroll K=0..3
-    {
-      __vector float b0, b1;
-      load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 0) * ldb, b0, b1);
-#define STEP_K0(i) FMAS_LANE(i, a##i##v, 0)
-      ROWS_APPLY(STEP_K0)
-#undef STEP_K0
-    }
-    {
-      __vector float b0, b1;
-      load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 1) * ldb, b0, b1);
-#define STEP_K1(i) FMAS_LANE(i, a##i##v, 1)
-      ROWS_APPLY(STEP_K1)
-#undef STEP_K1
-    }
-    {
-      __vector float b0, b1;
-      load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 2) * ldb, b0, b1);
-#define STEP_K2(i) FMAS_LANE(i, a##i##v, 2)
-      ROWS_APPLY(STEP_K2)
-#undef STEP_K2
-    }
-
-    {
-      __vector float b0, b1;
-      load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 3) * ldb, b0, b1);
-#define STEP_K3(i) FMAS_LANE(i, a##i##v, 3)
-      ROWS_APPLY(STEP_K3)
-#undef STEP_K3
-    }
-#undef FMAS_LANE
-  }
-
-  for (; k < K; ++k) {
-    __vector float b0, b1;
-    load_row8_B_as_f32<kv_cache_t>(B + (int64_t)k * ldb, b0, b1);
-#define TAIL_ROW(i)                              \
-  IF_M(i) {                                      \
-    __vector float ai = vec_splats(*(a##i + k)); \
-    acc##i##_0 = vec_madd(b0, ai, acc##i##_0);   \
-    acc##i##_1 = vec_madd(b1, ai, acc##i##_1);   \
-  }
-    ROWS_APPLY(TAIL_ROW)
-#undef TAIL_ROW
-  }
-
-#define STORE_ROW(i)                           \
-  IF_M(i) {                                    \
-    vec_xst(acc##i##_0, 0, C + (i) * ldc + 0); \
-    vec_xst(acc##i##_1, 0, C + (i) * ldc + 4); \
-  }
-  ROWS_APPLY(STORE_ROW)
-#undef STORE_ROW
-
-#undef ROWS_APPLY
-#undef IF_M
-}
-
-template <int32_t N, typename kv_cache_t>
-FORCE_INLINE void gemm_macro_s390x_Mx8_Ku4(const float* __restrict A,
-                                           const kv_cache_t* __restrict B,
-                                           float* __restrict C, int32_t M,
-                                           int32_t K, int64_t lda, int64_t ldb,
-                                           int64_t ldc, bool accumulate) {
-  static_assert(N % 8 == 0, "N must be a multiple of 8");
-  for (int32_t m = 0; m < M;) {
-    int32_t mb = (M - m >= 8) ? 8 : (M - m >= 4) ? 4 : (M - m >= 2) ? 2 : 1;
-    const float* Ab = A + m * lda;
-    float* Cb = C + m * ldc;
-
-    for (int32_t n = 0; n < N; n += 8) {
-      const kv_cache_t* Bn = B + n;
-      float* Cn = Cb + n;
-      switch (mb) {
-        case 8:
-          gemm_micro_s390x_Mx8_Ku4<8, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc, K,
-                                                  accumulate);
-          break;
-        case 4:
-          gemm_micro_s390x_Mx8_Ku4<4, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc, K,
-                                                  accumulate);
-          break;
-        case 2:
-          gemm_micro_s390x_Mx8_Ku4<2, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc, K,
-                                                  accumulate);
-          break;
-        default:
-          gemm_micro_s390x_Mx8_Ku4<1, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc, K,
-                                                  accumulate);
-          break;
-      }
-    }
-    m += mb;
-  }
-}
-
-template <typename kv_cache_t>
-class TileGemmS390X {
- public:
-  template <AttentionGemmPhase phase, int32_t k_size>
-  FORCE_INLINE static void gemm(const int32_t m_size,
-                                float* __restrict__ a_tile,
-                                kv_cache_t* __restrict__ b_tile,
-                                float* __restrict__ c_tile, const int64_t lda,
-                                const int64_t ldb, const int64_t ldc,
-                                const int32_t block_size,
-                                const int32_t dynamic_k_size,
-                                const bool accum_c) {
-    if constexpr (phase == AttentionGemmPhase::QK) {
-      gemm_macro_s390x_Mx8_Ku4<BLOCK_SIZE_ALIGNMENT, kv_cache_t>(
-          a_tile, b_tile, c_tile, m_size, k_size, lda, ldb, ldc, accum_c);
-    } else {
-      gemm_macro_s390x_Mx8_Ku4<HEAD_SIZE_ALIGNMENT, kv_cache_t>(
-          a_tile, b_tile, c_tile, m_size, dynamic_k_size, lda, ldb, ldc,
-          accum_c);
-    }
-  }
-};
-
-}  // namespace
-
-template <typename scalar_t, int64_t head_dim>
-class AttentionImpl<ISA::VXE, scalar_t, head_dim> {
- public:
-  using query_t = scalar_t;
-  using q_buffer_t = float;
-  using kv_cache_t = scalar_t;
-  using logits_buffer_t = float;
-  using partial_output_buffer_t = float;
-  using prob_buffer_t = float;
-
-  constexpr static int64_t BlockSizeAlignment = BLOCK_SIZE_ALIGNMENT;
-  constexpr static int64_t HeadDimAlignment = HEAD_SIZE_ALIGNMENT;
-  constexpr static int64_t MaxQHeadNumPerIteration = MAX_Q_HEAD_NUM_PER_ITER;
-  constexpr static int64_t HeadDim = head_dim;
-  constexpr static ISA ISAType = ISA::VXE;
-  constexpr static bool scale_on_logits =
-      false;  // Scale is applied to Q during copy
-
- public:
-  AttentionImpl() {}
-
-  template <template <typename tile_gemm_t> typename attention>
-  FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) {
-    attention<TileGemmS390X<kv_cache_t>> attention_iteration;
-    attention_iteration(CPU_ATTENTION_PARAMS);
-  }
-
-  // Strides for Memory Layout
-  constexpr static int64_t k_cache_token_group_stride(
-      const int32_t block_size) {
-    return BlockSizeAlignment;  // [head_dim, block_size] layout
-  }
-
-  constexpr static int64_t v_cache_token_group_stride(
-      const int32_t block_size) {
-    return head_dim * BlockSizeAlignment;
-  }
-
-  constexpr static int64_t v_cache_head_group_stride(const int32_t block_size) {
-    return HeadDimAlignment;
-  }
-
-  static void copy_q_heads_tile(scalar_t* __restrict__ src,
-                                float* __restrict__ q_buffer,
-                                const int32_t q_num,
-                                const int32_t q_heads_per_kv,
-                                const int64_t q_num_stride,
-                                const int64_t q_head_stride, float scale) {
-    __vector float scale_vec = vec_splats(scale);
-    constexpr bool is_bf16 = std::is_same<scalar_t, c10::BFloat16>::value;
-
-    // Process 8 elements at a time (32 bytes of float output)
-    for (int32_t i = 0; i < q_num; ++i) {
-      for (int32_t h = 0; h < q_heads_per_kv; ++h) {
-        scalar_t* curr_src = src + i * q_num_stride + h * q_head_stride;
-        float* curr_dst =
-            q_buffer + i * q_heads_per_kv * head_dim + h * head_dim;
-
-        int32_t d = 0;
-        for (; d <= head_dim - 8; d += 8) {
-          if constexpr (is_bf16) {
-            __vector float v0, v1;
-            // Reuse our Big-Endian-Safe loader
-            load_row8_B_as_f32<scalar_t>(curr_src + d, v0, v1);
-
-            v0 = vec_mul(v0, scale_vec);
-            v1 = vec_mul(v1, scale_vec);
-
-            vec_xst(v0, 0, curr_dst + d);
-            vec_xst(v1, 0, curr_dst + d + 4);
-          } else {
-            __vector float v0 = vec_xl((long long)0, (float*)curr_src + d);
-            __vector float v1 = vec_xl((long long)0, (float*)curr_src + d + 4);
-
-            v0 = vec_mul(v0, scale_vec);
-            v1 = vec_mul(v1, scale_vec);
-
-            vec_xst(v0, 0, curr_dst + d);
-            vec_xst(v1, 0, curr_dst + d + 4);
-          }
-        }
-
-        for (; d < head_dim; ++d) {
-          float val = static_cast<float>(curr_src[d]);
-          curr_dst[d] = val * scale;
-        }
-      }
-    }
-  }
-
-  static void reshape_and_cache(
-      const scalar_t* __restrict__ key, const scalar_t* __restrict__ value,
-      scalar_t* __restrict__ key_cache, scalar_t* __restrict__ value_cache,
-      const int64_t* __restrict__ slot_mapping, const int64_t token_num,
-      const int64_t key_token_num_stride, const int64_t value_token_num_stride,
-      const int64_t head_num, const int64_t key_head_num_stride,
-      const int64_t value_head_num_stride, const int64_t num_blocks,
-      const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
-      const int64_t block_size, const int64_t block_size_stride) {
-#pragma omp parallel for collapse(2)
-    for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
-      for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
-        const int64_t pos = slot_mapping[token_idx];
-        if (pos < 0) continue;
-
-        const int64_t block_idx = pos / block_size;
-        const int64_t block_offset = pos % block_size;
-
-        {
-          const scalar_t* key_src = key + token_idx * key_token_num_stride +
-                                    head_idx * key_head_num_stride;
-          scalar_t* key_dst = key_cache + block_idx * num_blocks_stride +
-                              head_idx * cache_head_num_stride + block_offset;
-
-          for (int64_t i = 0, j = 0; i < head_dim; ++i, j += block_size) {
-            key_dst[j] = key_src[i];
-          }
-        }
-
-        {
-          const scalar_t* val_src = value + token_idx * value_token_num_stride +
-                                    head_idx * value_head_num_stride;
-          scalar_t* val_dst = value_cache + block_idx * num_blocks_stride +
-                              head_idx * cache_head_num_stride +
-                              block_offset * head_dim;
-
-          std::memcpy(val_dst, val_src, sizeof(scalar_t) * head_dim);
-        }
-      }
-    }
-  }
-};
-
-}  // namespace cpu_attention
-
-#undef BLOCK_SIZE_ALIGNMENT
-#undef HEAD_SIZE_ALIGNMENT
-#undef MAX_Q_HEAD_NUM_PER_ITER
-
-#endif
@@ -237,10 +237,13 @@ W8A8MatMulPrimitiveHandler::W8A8MatMulPrimitiveHandler(const Args& args)
  };
  dnnl::memory::desc original_b_md({b_k_size_, b_n_size_}, b_type_,
                                   {b_k_stride_, b_n_stride_});
-
+#ifdef __aarch64__
  // dummy M size for prepacking weights
  // Prepacking weights improves performance and avoid runtime reorders
  constexpr dnnl_dim_t kProbeM = 128;
+#else
+  constexpr dnnl_dim_t kProbeM = DNNL_RUNTIME_DIM_VAL;
+#endif

  prepack_weight(args.b_ptr, original_b_md,
                 create_primitive_desc(
@@ -408,19 +411,21 @@ MatMulPrimitiveHandler::MatMulPrimitiveHandler(const Args& args)
  dnnl::memory::desc original_b_md({b_k_size_, b_n_size_}, b_type_,
                                   {b_k_stride_, b_n_stride_});

-  // dummy M size for prepacking weights
-  // Prepacking weights improves performance and avoid runtime reorders
-  constexpr dnnl_dim_t kProbeM = 128;
-
  prepack_weight(args.b_ptr, original_b_md,
                 create_primitive_desc(
-                     MSizeCacheKey{// Use a concrete M so oneDNN's kernel
-                                   // selector can choose an optimally blocked
-                                   // weight layout.
-                                   .a_m_size = kProbeM,
-                                   .a_m_stride = b_k_size_,
-                                   .use_bias = false,
-                                   .bias_type = dnnl::memory::data_type::undef},
+                     MSizeCacheKey{
+#ifdef VLLM_USE_ACL
+                         // Arm Compute Library (ACL) backend for oneDNN does
+                         // not support runtime
+                         // dimensions, so we set M to a default value
+                         .a_m_size = 128,
+                         .a_m_stride = b_k_size_,
+#else
+                         .a_m_size = DNNL_RUNTIME_DIM_VAL,
+                         .a_m_stride = DNNL_RUNTIME_DIM_VAL,
+#endif
+                         .use_bias = false,
+                         .bias_type = dnnl::memory::data_type::undef},
                     true)
                     .weights_desc());
  init_runtime_memory_cache(args);
@@ -19,11 +19,10 @@ ISA_TYPES = {
    "VEC": 1,
    "VEC16": 2,
    "NEON": 3,
-    "VXE": 4,
 }

 # ISAs supported for head_dims divisible by 32
-ISA_FOR_32 = ["AMX", "NEON", "VEC", "VEC16", "VXE"]
+ISA_FOR_32 = ["AMX", "NEON", "VEC", "VEC16"]

 # ISAs supported for head_dims divisible by 16 only
 ISA_FOR_16 = ["VEC16"]
@@ -119,10 +118,6 @@ def generate_header_file() -> str:
  #include "cpu_attn_neon.hpp"
 #endif

-#ifdef __s390x__
-  #include "cpu_attn_vxe.hpp"
-#endif
-
 """

    header += generate_helper_function()
@@ -168,25 +163,6 @@ def generate_header_file() -> str:
    } \\
  }()

-"""
-
-    # s390x with VXE
-    header += """#elif defined(__s390x__)
-#define CPU_ATTN_DISPATCH(HEAD_DIM, ISA_TYPE, ...) \\
-  [&] { \\
-    int64_t encoded_params = encode_cpu_attn_params(HEAD_DIM, ISA_TYPE); \\
-    switch (encoded_params) { \\
-"""
-    header += generate_cases_for_isa_group(["VXE", "VEC", "VEC16"])
-    header += """
-      default: { \\
-        TORCH_CHECK(false, "Unsupported CPU attention configuration: head_dim=" + \\
-                    std::to_string(HEAD_DIM) + " isa=" + \\
-                    std::to_string(static_cast<int>(ISA_TYPE))); \\
-      } \\
-    } \\
-  }()
-
 """

    # Fallback: VEC and VEC16 only
@@ -206,7 +182,7 @@ def generate_header_file() -> str:
    } \\
  }()

-#endif  /* CPU_CAPABILITY_AMXBF16 / __aarch64__ / __s390x__ */
+#endif  /* CPU_CAPABILITY_AMXBF16 / __aarch64__ */

 #endif  // CPU_ATTN_DISPATCH_GENERATED_H
 """
@@ -4,10 +4,6 @@

 #include <torch/library.h>

-// Note: overwrite the external definition for sharing same name between
-// libraries use different ISAs.
-#define TORCH_EXTENSION_NAME _C
-
 std::string init_cpu_threads_env(const std::string& cpu_ids);

 void release_dnnl_matmul_handler(int64_t handler);
@@ -328,12 +324,19 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      "str act, str isa) -> ()");
  ops.impl("cpu_fused_moe", torch::kCPU, &cpu_fused_moe);
 #endif
-  ops.def("init_cpu_threads_env(str cpu_ids) -> str", &init_cpu_threads_env);
-  ops.def(
+}
+
+TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _utils), utils) {
+  // CPU utils
+  utils.def("init_cpu_threads_env(str cpu_ids) -> str", &init_cpu_threads_env);
+}
+
+TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cpu), cpu_ops) {
+  cpu_ops.def(
      "mla_decode_kvcache("
      "   Tensor! out, Tensor query, Tensor kv_cache,"
      "   float scale, Tensor block_tables, Tensor seq_lens) -> ()");
-  ops.impl("mla_decode_kvcache", torch::kCPU, &mla_decode_kvcache);
+  cpu_ops.impl("mla_decode_kvcache", torch::kCPU, &mla_decode_kvcache);
 }

 REGISTER_EXTENSION(TORCH_EXTENSION_NAME)
--- a/Show More
+++ b/Show More