[Release] [CI] Optim release pipeline (#33156)

Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
(cherry picked from commit f9d03599ef)

.buildkite/ci_config.yaml

@@ -1,8 +1,7 @@
 name: vllm_ci
 job_dirs:
-  - ".buildkite/image_build"
   - ".buildkite/test_areas"
-  - ".buildkite/hardware_tests"
+  - ".buildkite/image_build"
 run_all_patterns:
   - "docker/Dockerfile"
   - "CMakeLists.txt"

.buildkite/hardware_tests/amd.yaml

@@ -1,30 +0,0 @@
-group: Hardware - AMD Build 
-steps:
-  - label: "AMD: :docker: build image"
-    key: image-build-amd
-    depends_on: []
-    device: amd_cpu
-    no_plugin: true
-    commands:
-    - >
-      docker build
-      --build-arg max_jobs=16
-      --build-arg REMOTE_VLLM=1
-      --build-arg ARG_PYTORCH_ROCM_ARCH='gfx90a;gfx942'
-      --build-arg VLLM_BRANCH=$BUILDKITE_COMMIT
-      --tag "rocm/vllm-ci:${BUILDKITE_COMMIT}"
-      -f docker/Dockerfile.rocm
-      --target test
-      --no-cache
-      --progress plain .
-    - docker push "rocm/vllm-ci:${BUILDKITE_COMMIT}"
-    env:
-      DOCKER_BUILDKIT: "1"
-    retry:
-      automatic:
-        - exit_status: -1  # Agent was lost
-          limit: 1
-        - exit_status: -10  # Agent was lost
-          limit: 1
-        - exit_status: 1  # Machine occasionally fail
-          limit: 1

.buildkite/hardware_tests/ascend_npu.yaml

@@ -1,10 +0,0 @@
-group: Hardware
-depends_on: ~
-steps:
-  - label: "Ascend NPU Test"
-    soft_fail: true
-    timeout_in_minutes: 20
-    no_plugin: true
-    device: ascend_npu
-    commands: 
-    - bash .buildkite/scripts/hardware_ci/run-npu-test.sh

.buildkite/hardware_tests/cpu.yaml

@@ -1,100 +0,0 @@
-group: CPU
-depends_on: []
-steps:
-- label: CPU-Kernel Tests
-  depends_on: []
-  soft_fail: true
-  device: intel_cpu
-  no_plugin: true
-  source_file_dependencies:
-  - csrc/cpu/
-  - cmake/cpu_extension.cmake
-  - CMakeLists.txt
-  - vllm/_custom_ops.py
-  - tests/kernels/attention/test_cpu_attn.py
-  - tests/kernels/moe/test_cpu_fused_moe.py
-  - tests/kernels/test_onednn.py
-  commands:
-    - |
-      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 20m "
-      pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
-      pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py
-      pytest -x -v -s tests/kernels/test_onednn.py"
-
-- label: CPU-Language Generation and Pooling Model Tests
-  depends_on: []
-  soft_fail: true
-  device: intel_cpu
-  no_plugin: true
-  source_file_dependencies:
-  - csrc/cpu/
-  - vllm/
-  - tests/models/language/generation/
-  - tests/models/language/pooling/
-  commands:
-    - |
-      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
-      pytest -x -v -s tests/models/language/generation -m cpu_model
-      pytest -x -v -s tests/models/language/pooling -m cpu_model"
-
-- label: CPU-Quantization Model Tests
-  depends_on: []
-  soft_fail: true
-  device: intel_cpu
-  no_plugin: true
-  source_file_dependencies:
-  - csrc/cpu/
-  - vllm/model_executor/layers/quantization/cpu_wna16.py
-  - vllm/model_executor/layers/quantization/gptq_marlin.py
-  - vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_int8.py
-  - vllm/model_executor/layers/quantization/kernels/scaled_mm/cpu.py
-  - vllm/model_executor/layers/quantization/kernels/mixed_precision/cpu.py
-  - tests/quantization/test_compressed_tensors.py
-  - tests/quantization/test_cpu_wna16.py
-  commands:
-    - |
-      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 20m "
-      pytest -x -v -s tests/quantization/test_compressed_tensors.py::test_compressed_tensors_w8a8_logprobs
-      pytest -x -v -s tests/quantization/test_cpu_wna16.py"
-      
-- label: CPU-Distributed Tests
-  depends_on: []
-  soft_fail: true
-  device: intel_cpu
-  no_plugin: true
-  source_file_dependencies:
-  - csrc/cpu/shm.cpp
-  - vllm/v1/worker/cpu_worker.py
-  - vllm/v1/worker/gpu_worker.py
-  - vllm/v1/worker/cpu_model_runner.py
-  - vllm/v1/worker/gpu_model_runner.py
-  - vllm/platforms/cpu.py
-  - vllm/distributed/parallel_state.py
-  - vllm/distributed/device_communicators/cpu_communicator.py
-  commands:
-    - |
-      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 10m "
-      bash .buildkite/scripts/hardware_ci/run-cpu-distributed-smoke-test.sh"
-
-- label: CPU-Multi-Modal Model Tests %N
-  depends_on: []
-  soft_fail: true
-  device: intel_cpu
-  no_plugin: true
-  source_file_dependencies:
-  # - vllm/
-  - vllm/model_executor/layers/rotary_embedding
-  - tests/models/multimodal/generation/
-  commands:
-    - |
-      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 45m "
-      pytest -x -v -s tests/models/multimodal/generation --ignore=tests/models/multimodal/generation/test_pixtral.py -m cpu_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB"
-  parallelism: 2
-
-- label: "Arm CPU Test"
-  depends_on: []
-  soft_fail: true
-  device: arm_cpu
-  no_plugin: true
-  commands: 
-  - bash .buildkite/scripts/hardware_ci/run-cpu-test-arm.sh

.buildkite/hardware_tests/gh200.yaml

@@ -1,10 +0,0 @@
-group: Hardware
-steps:
-  - label: "GH200 Test"
-    soft_fail: true
-    device: gh200
-    no_plugin: true
-    optional: true
-    commands: 
-    - nvidia-smi 
-    - bash .buildkite/scripts/hardware_ci/run-gh200-test.sh

.buildkite/hardware_tests/intel.yaml

@@ -1,17 +0,0 @@
-group: Hardware
-depends_on: ~
-steps:
-  - label: "Intel HPU Test"
-    soft_fail: true
-    device: intel_hpu
-    no_plugin: true
-    commands: 
-    - bash .buildkite/scripts/hardware_ci/run-hpu-test.sh
-
-  - label: "Intel GPU Test"
-    depends_on: []
-    soft_fail: true
-    device: intel_gpu
-    no_plugin: true
-    commands: 
-    - bash .buildkite/scripts/hardware_ci/run-xpu-test.sh

.buildkite/image_build/image_build.sh

@@ -1,256 +1,56 @@
 #!/bin/bash
-set -euo pipefail
+set -e
 
-# replace invalid characters in Docker image tags and truncate to 128 chars
-clean_docker_tag() {
-    local input="$1"
-    echo "$input" | sed 's/[^a-zA-Z0-9._-]/_/g' | cut -c1-128
-}
-
-print_usage_and_exit() {
-    echo "Usage: $0 <registry> <repo> <commit> <branch> <vllm_use_precompiled> <vllm_merge_base_commit> <cache_from> <cache_to>"
-    exit 1
-}
-
-print_instance_info() {
-    echo ""
-    echo "=== Debug: Instance Information ==="
-    # Get IMDSv2 token
-    if TOKEN=$(curl -s -X PUT "http://169.254.169.254/latest/api/token" \
-            -H "X-aws-ec2-metadata-token-ttl-seconds: 21600" 2>/dev/null); then
-        AMI_ID=$(curl -s -H "X-aws-ec2-metadata-token: $TOKEN" \
-            http://169.254.169.254/latest/meta-data/ami-id 2>/dev/null || echo "unknown")
-        INSTANCE_TYPE=$(curl -s -H "X-aws-ec2-metadata-token: $TOKEN" \
-            http://169.254.169.254/latest/meta-data/instance-type 2>/dev/null || echo "unknown")
-        INSTANCE_ID=$(curl -s -H "X-aws-ec2-metadata-token: $TOKEN" \
-            http://169.254.169.254/latest/meta-data/instance-id 2>/dev/null || echo "unknown")
-        AZ=$(curl -s -H "X-aws-ec2-metadata-token: $TOKEN" \
-            http://169.254.169.254/latest/meta-data/placement/availability-zone 2>/dev/null || echo "unknown")
-        echo "AMI ID:        ${AMI_ID}"
-        echo "Instance Type: ${INSTANCE_TYPE}"
-        echo "Instance ID:   ${INSTANCE_ID}"
-        echo "AZ:            ${AZ}"
-    else
-        echo "Not running on EC2 or IMDS not available"
-    fi
-    # Check for warm cache AMI (marker file baked into custom AMI)
-    if [[ -f /etc/vllm-ami-info ]]; then
-        echo "Cache:         warm (custom vLLM AMI)"
-        cat /etc/vllm-ami-info
-    else
-        echo "Cache:         cold (standard AMI)"
-    fi
-    echo "==================================="
-    echo ""
-}
-
-setup_buildx_builder() {
-    echo "--- :buildkite: Setting up buildx builder"
-    if [[ -S "${BUILDKIT_SOCKET}" ]]; then
-        # Custom AMI with standalone buildkitd - use remote driver for warm cache
-        echo "✅ Found local buildkitd socket at ${BUILDKIT_SOCKET}"
-        echo "Using remote driver to connect to buildkitd (warm cache available)"
-        if docker buildx inspect baked-vllm-builder >/dev/null 2>&1; then
-            echo "Using existing baked-vllm-builder"
-            docker buildx use baked-vllm-builder
-        else
-            echo "Creating baked-vllm-builder with remote driver"
-            docker buildx create \
-                --name baked-vllm-builder \
-                --driver remote \
-                --use \
-                "unix://${BUILDKIT_SOCKET}"
-        fi
-        docker buildx inspect --bootstrap
-    elif docker buildx inspect "${BUILDER_NAME}" >/dev/null 2>&1; then
-        # Existing builder available
-        echo "Using existing builder: ${BUILDER_NAME}"
-        docker buildx use "${BUILDER_NAME}"
-        docker buildx inspect --bootstrap
-    else
-        # No local buildkitd, no existing builder - create new docker-container builder
-        echo "No local buildkitd found, using docker-container driver"
-        docker buildx create --name "${BUILDER_NAME}" --driver docker-container --use
-        docker buildx inspect --bootstrap
-    fi
-
-    # builder info
-    echo "Active builder:"
-    docker buildx ls | grep -E '^\*|^NAME' || docker buildx ls
-}
-
-check_and_skip_if_image_exists() {
-    if [[ -n "${IMAGE_TAG:-}" ]]; then
-        echo "--- :mag: Checking if image exists"
-        if docker manifest inspect "${IMAGE_TAG}" >/dev/null 2>&1; then
-            echo "Image already exists: ${IMAGE_TAG}"
-            echo "Skipping build"
-            exit 0
-        fi
-        echo "Image not found, proceeding with build"
-    fi
-}
-
-ecr_login() {
-    aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY"
-    aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com
-}
-
-prepare_cache_tags() {
-    # resolve and set: CACHE_TO, CACHE_FROM, CACHE_FROM_BASE_BRANCH, CACHE_FROM_MAIN
-    TEST_CACHE_ECR="936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-test-cache"
-    MAIN_CACHE_ECR="936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-postmerge-cache"
-
-    if [[ "$BUILDKITE_PULL_REQUEST" == "false" ]]; then
-        if [[ "$BUILDKITE_BRANCH" == "main" ]]; then
-            cache="${MAIN_CACHE_ECR}:latest"
-        else
-            clean_branch=$(clean_docker_tag "$BUILDKITE_BRANCH")
-            cache="${TEST_CACHE_ECR}:${clean_branch}"
-        fi
-        CACHE_TO="$cache"
-        CACHE_FROM="$cache"
-        CACHE_FROM_BASE_BRANCH="$cache"
-    else
-        CACHE_TO="${TEST_CACHE_ECR}:pr-${BUILDKITE_PULL_REQUEST}"
-        CACHE_FROM="${TEST_CACHE_ECR}:pr-${BUILDKITE_PULL_REQUEST}"
-        if [[ "$BUILDKITE_PULL_REQUEST_BASE_BRANCH" == "main" ]]; then
-            CACHE_FROM_BASE_BRANCH="${MAIN_CACHE_ECR}:latest"
-        else
-            clean_base=$(clean_docker_tag "$BUILDKITE_PULL_REQUEST_BASE_BRANCH")
-            CACHE_FROM_BASE_BRANCH="${TEST_CACHE_ECR}:${clean_base}"
-        fi
-    fi
-
-    CACHE_FROM_MAIN="${MAIN_CACHE_ECR}:latest"
-    export CACHE_TO CACHE_FROM CACHE_FROM_BASE_BRANCH CACHE_FROM_MAIN
-}
-
-resolve_parent_commit() {
-    if [[ -z "${PARENT_COMMIT:-}" ]]; then
-        PARENT_COMMIT=$(git rev-parse HEAD~1 2>/dev/null || echo "")
-        if [[ -n "${PARENT_COMMIT}" ]]; then
-            echo "Computed parent commit for cache fallback: ${PARENT_COMMIT}"
-            export PARENT_COMMIT
-        else
-            echo "Could not determine parent commit (may be first commit in repo)"
-        fi
-    else
-        echo "Using provided PARENT_COMMIT: ${PARENT_COMMIT}"
-    fi
-}
-
-print_bake_config() {
-    echo "--- :page_facing_up: Resolved bake configuration"
-    BAKE_CONFIG_FILE="bake-config-build-${BUILDKITE_BUILD_NUMBER:-local}.json"
-    docker buildx bake -f "${VLLM_BAKE_FILE_PATH}" -f "${CI_HCL_PATH}" --print "${TARGET}" | tee "${BAKE_CONFIG_FILE}" || true
-    echo "Saved bake config to ${BAKE_CONFIG_FILE}"
-    echo "--- :arrow_down: Uploading bake config to Buildkite"
-    buildkite-agent artifact upload "${BAKE_CONFIG_FILE}"
-}
-
-#################################
-#         Main Script           #
-#################################
-print_instance_info
-
-if [[ $# -lt 7 ]]; then
-    print_usage_and_exit
+if [[ $# -lt 8 ]]; then
+  echo "Usage: $0 <registry> <repo> <commit> <branch> <vllm_use_precompiled> <vllm_merge_base_commit> <cache_from> <cache_to>"
+  exit 1
 fi
 
-# input args
 REGISTRY=$1
 REPO=$2
 BUILDKITE_COMMIT=$3
 BRANCH=$4
 VLLM_USE_PRECOMPILED=$5
 VLLM_MERGE_BASE_COMMIT=$6
-IMAGE_TAG=$7
-IMAGE_TAG_LATEST=${8:-} # only used for main branch, optional
+CACHE_FROM=$7
+CACHE_TO=$8
 
-# build config
-TARGET="test-ci"
-VLLM_BAKE_FILE_PATH="${VLLM_BAKE_FILE_PATH:-docker/docker-bake.hcl}"
-BUILDER_NAME="${BUILDER_NAME:-vllm-builder}"
-CI_HCL_URL="${CI_HCL_URL:-https://raw.githubusercontent.com/vllm-project/ci-infra/main/docker/ci.hcl}"
-CI_HCL_PATH="/tmp/ci.hcl"
-BUILDKIT_SOCKET="/run/buildkit/buildkitd.sock"
+# authenticate with AWS ECR
+aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin $REGISTRY
+aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com
 
-prepare_cache_tags
-ecr_login
+# docker buildx 
+docker buildx create --name vllm-builder --driver docker-container --use
+docker buildx inspect --bootstrap
+docker buildx ls
 
-# Environment info (for docs and human readers)
-#   VLLM_CI_BRANCH      - ci-infra branch to use (default: main)
-#   VLLM_BAKE_FILE_PATH      - Path to vLLM's bake file (default: docker/docker-bake.hcl)
-#   BUILDER_NAME        - Name for buildx builder (default: vllm-builder)
-#
-# Build configuration (exported as environment variables for bake):
-export BUILDKITE_COMMIT
-export PARENT_COMMIT
-export IMAGE_TAG
-export IMAGE_TAG_LATEST
-export CACHE_FROM
-export CACHE_FROM_BASE_BRANCH
-export CACHE_FROM_MAIN
-export CACHE_TO
-export VLLM_USE_PRECOMPILED
-export VLLM_MERGE_BASE_COMMIT
-
-# print args
-echo "--- :mag: Arguments"
-echo "REGISTRY: ${REGISTRY}"
-echo "REPO: ${REPO}"
-echo "BUILDKITE_COMMIT: ${BUILDKITE_COMMIT}"
-echo "BRANCH: ${BRANCH}"
-echo "VLLM_USE_PRECOMPILED: ${VLLM_USE_PRECOMPILED}"
-echo "VLLM_MERGE_BASE_COMMIT: ${VLLM_MERGE_BASE_COMMIT}"
-echo "IMAGE_TAG: ${IMAGE_TAG}"
-echo "IMAGE_TAG_LATEST: ${IMAGE_TAG_LATEST}"
-
-# print build configuration
-echo "--- :mag: Build configuration"
-echo "TARGET: ${TARGET}"
-echo "vLLM bake file: ${VLLM_BAKE_FILE_PATH}"
-echo "BUILDER_NAME: ${BUILDER_NAME}"
-echo "CI_HCL_URL: ${CI_HCL_URL}"
-echo "BUILDKIT_SOCKET: ${BUILDKIT_SOCKET}"
-
-echo "--- :mag: Cache tags"
-echo "CACHE_TO: ${CACHE_TO}"
-echo "CACHE_FROM: ${CACHE_FROM}"
-echo "CACHE_FROM_BASE_BRANCH: ${CACHE_FROM_BASE_BRANCH}"
-echo "CACHE_FROM_MAIN: ${CACHE_FROM_MAIN}"
-
-check_and_skip_if_image_exists
-
-echo "--- :docker: Setting up Docker buildx bake"
-echo "Target: ${TARGET}"
-echo "vLLM bake file: ${VLLM_BAKE_FILE_PATH}"
-echo "CI HCL path: ${CI_HCL_PATH}"
-
-if [[ ! -f "${VLLM_BAKE_FILE_PATH}" ]]; then
-    echo "Error: vLLM bake file not found at ${VLLM_BAKE_FILE_PATH}"
-    echo "Make sure you're running from the vLLM repository root"
-    exit 1
+# skip build if image already exists
+if [[ -z $(docker manifest inspect $REGISTRY/$REPO:$BUILDKITE_COMMIT) ]]; then
+  echo "Image not found, proceeding with build..."
+else
+  echo "Image found"
+  exit 0
 fi
 
-echo "--- :arrow_down: Downloading ci.hcl"
-curl -sSfL -o "${CI_HCL_PATH}" "${CI_HCL_URL}"
-echo "Downloaded to ${CI_HCL_PATH}"
-
-if [[ ! -f "${CI_HCL_PATH}" ]]; then
-    echo "Error: ci.hcl not found at ${CI_HCL_PATH}"
-    exit 1
+if [[ "${VLLM_USE_PRECOMPILED:-0}" == "1" ]]; then
+  merge_base_commit_build_args="--build-arg VLLM_MERGE_BASE_COMMIT=${VLLM_MERGE_BASE_COMMIT}"
+else
+  merge_base_commit_build_args=""
 fi
 
-setup_buildx_builder
-
-resolve_parent_commit
-export PARENT_COMMIT
-
-print_bake_config
-
-echo "--- :docker: Building ${TARGET}"
-docker --debug buildx bake -f "${VLLM_BAKE_FILE_PATH}" -f "${CI_HCL_PATH}" --progress plain "${TARGET}"
-
-echo "--- :white_check_mark: Build complete"
+# build
+docker buildx build --file docker/Dockerfile \
+  --build-arg max_jobs=16 \
+  --build-arg buildkite_commit=$BUILDKITE_COMMIT \
+  --build-arg USE_SCCACHE=1 \
+  --build-arg TORCH_CUDA_ARCH_LIST="8.0 8.9 9.0 10.0" \
+  --build-arg FI_TORCH_CUDA_ARCH_LIST="8.0 8.9 9.0a 10.0a" \
+  --build-arg VLLM_USE_PRECOMPILED="${VLLM_USE_PRECOMPILED:-0}" \
+  ${merge_base_commit_build_args} \
+  --cache-from type=registry,ref=${CACHE_FROM},mode=max \
+  --cache-to type=registry,ref=${CACHE_TO},mode=max \
+  --tag ${REGISTRY}/${REPO}:${BUILDKITE_COMMIT} \
+  $( [[ "${BRANCH}" == "main" ]] && echo "--tag ${REGISTRY}/${REPO}:latest" ) \
+  --push \
+  --target test \
+  --progress plain .

.buildkite/image_build/image_build.yaml

@@ -3,10 +3,8 @@ steps:
   - label: ":docker: Build image"
     key: image-build
     depends_on: []
-    timeout_in_minutes: 600
     commands:
-    - if [[ "$BUILDKITE_BRANCH" != "main" ]]; then .buildkite/image_build/image_build.sh $REGISTRY $REPO $BUILDKITE_COMMIT $BRANCH $VLLM_USE_PRECOMPILED $VLLM_MERGE_BASE_COMMIT $IMAGE_TAG; fi
-    - if [[ "$BUILDKITE_BRANCH" == "main" ]]; then .buildkite/image_build/image_build.sh $REGISTRY $REPO $BUILDKITE_COMMIT $BRANCH $VLLM_USE_PRECOMPILED $VLLM_MERGE_BASE_COMMIT $IMAGE_TAG $IMAGE_TAG_LATEST; fi
+    - .buildkite/image_build/image_build.sh $REGISTRY $REPO $BUILDKITE_COMMIT $BRANCH $VLLM_USE_PRECOMPILED $VLLM_MERGE_BASE_COMMIT $CACHE_FROM $CACHE_TO
     retry:
       automatic:
         - exit_status: -1  # Agent was lost
@@ -42,7 +40,7 @@ steps:
           limit: 2
         - exit_status: -10  # Agent was lost
           limit: 2
-
+  
   - label: ":docker: Build CPU arm64 image"
     key: cpu-arm64-image-build
     depends_on: []

.buildkite/lm-eval-harness/configs/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16.yaml

@@ -1,15 +0,0 @@
-model_name: "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16"
-tasks:
-- name: "gsm8k"
-  metrics:
-  - name: "exact_match,strict-match"
-    value: 0.695
-  - name: "exact_match,flexible-extract"
-    value: 0.447
-limit: 1319
-num_fewshot: 5
-max_model_len: 262144
-enforce_eager: false
-apply_chat_template: true
-fewshot_as_multiturn: true
-trust_remote_code: true

.buildkite/lm-eval-harness/configs/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8.yaml

@@ -1,19 +0,0 @@
-model_name: "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8"
-tasks:
-- name: "gsm8k"
-  metrics:
-  - name: "exact_match,strict-match"
-    value: 0.7142
-  - name: "exact_match,flexible-extract"
-    value: 0.4579
-env_vars:
-  VLLM_USE_FLASHINFER_MOE_FP8: "1"
-  VLLM_FLASHINFER_MOE_BACKEND: "throughput"
-limit: 1319
-num_fewshot: 5
-max_model_len: 262144
-kv_cache_dtype: fp8
-enforce_eager: false
-apply_chat_template: true
-fewshot_as_multiturn: true
-trust_remote_code: true

.buildkite/lm-eval-harness/configs/models-large-hopper.txt

@@ -1,2 +1 @@
 Qwen3-235B-A22B-Instruct-2507-FP8.yaml
-NVIDIA-Nemotron-3-Nano-30B-A3B-FP8.yaml

.buildkite/lm-eval-harness/configs/models-large.txt

@@ -3,4 +3,3 @@ Meta-Llama-3-70B-Instruct.yaml
 Mixtral-8x7B-Instruct-v0.1.yaml
 Qwen2-57B-A14-Instruct.yaml
 DeepSeek-V2-Lite-Chat.yaml
-NVIDIA-Nemotron-3-Nano-30B-A3B-BF16.yaml

.buildkite/performance-benchmarks/scripts/compare-json-results.py

@@ -9,10 +9,8 @@ import json
 import os
 from dataclasses import dataclass
 from importlib import util
-from pathlib import Path
 
 import pandas as pd
-import regex as re
 
 pd.options.display.float_format = "{:.2f}".format
 plotly_found = util.find_spec("plotly.express") is not None
@@ -277,131 +275,6 @@ def _apply_two_decimals(
     return styler.format({c: "{:.2f}" for c in num_cols}, na_rep="")
 
 
-# -----------------------------
-# Export helpers (Excel + CSV)
-# -----------------------------
-def _sanitize_sheet_name(name: str) -> str:
-    """
-    Excel sheet constraints:
-      - max 31 chars
-      - cannot contain: : \ / ? * [ ]
-      - cannot be empty
-    """
-    name = "sheet" if name is None else str(name)
-    name = re.sub(r"[:\\/?*\[\]]", "_", name)
-    name = name.strip().strip("'")
-    name = re.sub(r"\s+", " ", name)
-    if not name:
-        name = "sheet"
-    return name[:31]
-
-
-def _group_to_sheet_base(group_cols: list[str], gkey_tuple) -> str:
-    d = dict(zip(group_cols, gkey_tuple))
-    model = d.get("Model", "model")
-    model_short = str(model).split("/")[-1]
-    ilen = d.get("Input Len", "")
-    olen = d.get("Output Len", "")
-    lens = f"_{ilen}x{olen}" if ilen != "" and olen != "" else ""
-    return _sanitize_sheet_name(f"{model_short}{lens}")
-
-
-def _write_tables_to_excel_sheet(
-    writer: pd.ExcelWriter, sheet: str, blocks: list[tuple[str, pd.DataFrame]]
-):
-    startrow = 0
-    for title, df in blocks:
-        pd.DataFrame([[title]]).to_excel(
-            writer, sheet_name=sheet, index=False, header=False, startrow=startrow
-        )
-        startrow += 1
-        df.to_excel(writer, sheet_name=sheet, index=False, startrow=startrow)
-        startrow += len(df) + 3
-
-
-def _safe_filename(s: str) -> str:
-    s = re.sub(r"[^\w\-.]+", "_", str(s).strip())
-    return s[:180] if len(s) > 180 else s
-
-
-# -----------------------------
-# vLLM environment export helper
-# -----------------------------
-def _parse_vllm_env_txt(env_path: Path) -> pd.DataFrame:
-    """Parse vllm_env.txt into a flat table (Section, Key, Value).
-
-    Supports:
-      - section headers as standalone lines (no ':' or '=')
-      - key-value lines like 'OS: Ubuntu ...'
-      - env var lines like 'HF_HOME=/data/hf'
-    """
-    lines = env_path.read_text(encoding="utf-8", errors="replace").splitlines()
-    section = "General"
-    rows: list[dict] = []
-
-    def set_section(s: str):
-        nonlocal section
-        s = (s or "").strip()
-        if s:
-            section = s
-
-    for raw in lines:
-        stripped = raw.strip()
-        if not stripped:
-            continue
-        # divider lines like =====
-        if set(stripped) <= {"="}:
-            continue
-
-        # section header heuristic: short standalone line
-        if ":" not in stripped and "=" not in stripped and len(stripped) <= 64:
-            if stripped.lower().startswith("collecting environment information"):
-                continue
-            set_section(stripped)
-            continue
-
-        # env var style: KEY=VALUE (and not a URL with :)
-        if "=" in stripped and ":" not in stripped:
-            k, v = stripped.split("=", 1)
-            k = k.strip()
-            v = v.strip()
-            if k:
-                rows.append({"Section": section, "Key": k, "Value": v})
-            continue
-
-        # key: value
-        if ":" in stripped:
-            k, v = stripped.split(":", 1)
-            k = k.strip()
-            v = v.strip()
-            if k:
-                rows.append({"Section": section, "Key": k, "Value": v})
-            continue
-
-    return pd.DataFrame(rows, columns=["Section", "Key", "Value"])
-
-
-def _load_env_df_for_inputs(args, files: list[str]) -> pd.DataFrame | None:
-    """Load vllm_env.txt next to the *original* input JSON file.
-
-    Note: when only one -f is provided, the script may split JSON into ./splits/...,
-    but vllm_env.txt typically lives next to the original benchmark_results.json.
-    """
-    base_dir: Path | None = None
-    if getattr(args, "file", None):
-        base_dir = Path(args.file[0]).resolve().parent
-    elif files:
-        base_dir = Path(files[0]).resolve().parent
-    if base_dir is None:
-        return None
-
-    env_path = base_dir / "vllm_env.txt"
-    if not env_path.exists():
-        return None
-    df = _parse_vllm_env_txt(env_path)
-    return df
-
-
 # -----------------------------
 # Valid max concurrency summary helpers
 # -----------------------------
@@ -555,6 +428,7 @@ def build_valid_max_concurrency_summary_html(
 
     summary_df = pd.DataFrame(rows)
 
+    # --- Coerce numeric columns so Styler doesn't miss them due to object dtype ---
     for c in summary_df.columns:
         if c == "Configuration":
             continue
@@ -562,10 +436,12 @@ def build_valid_max_concurrency_summary_html(
 
     both_col = f"Max {conc_col} (Both)"
 
+    # --- Strict 2-decimal formatting for ALL non-Configuration columns ---
     formatters = {}
     for c in summary_df.columns:
         if c == "Configuration":
             continue
+        # default argument binds per-column formatter correctly
         formatters[c] = lambda v: "" if pd.isna(v) else f"{float(v):.2f}"
 
     styler = summary_df.style.format(formatters)
@@ -584,95 +460,6 @@ def build_valid_max_concurrency_summary_html(
     return title + styler.to_html(table_attributes='border="1" class="dataframe"')
 
 
-def build_valid_max_concurrency_summary_df(
-    tput_group_df: pd.DataFrame | None,
-    ttft_group_df: pd.DataFrame | None,
-    tpot_group_df: pd.DataFrame | None,
-    conc_col: str,
-    args,
-) -> pd.DataFrame | None:
-    if ttft_group_df is None and tpot_group_df is None:
-        return None
-
-    ttft_cols = (
-        _config_value_columns(ttft_group_df, conc_col)
-        if ttft_group_df is not None
-        else []
-    )
-    tpot_cols = (
-        _config_value_columns(tpot_group_df, conc_col)
-        if tpot_group_df is not None
-        else []
-    )
-    tput_cols = (
-        _config_value_columns(tput_group_df, conc_col)
-        if tput_group_df is not None
-        else []
-    )
-
-    if ttft_group_df is not None and tpot_group_df is not None:
-        cfg_cols = [c for c in ttft_cols if c in tpot_cols]
-        if tput_group_df is not None:
-            cfg_cols = [c for c in cfg_cols if c in tput_cols] or cfg_cols
-    else:
-        cfg_cols = ttft_cols or tpot_cols
-
-    if not cfg_cols:
-        cfg_cols = sorted(set(ttft_cols) | set(tpot_cols) | set(tput_cols), key=str)
-
-    rows = []
-    for cfg in cfg_cols:
-        ttft_max = (
-            _max_concurrency_ok(ttft_group_df, conc_col, cfg, args.ttft_max_ms)
-            if ttft_group_df is not None
-            else pd.NA
-        )
-        tpot_max = (
-            _max_concurrency_ok(tpot_group_df, conc_col, cfg, args.tpot_max_ms)
-            if tpot_group_df is not None
-            else pd.NA
-        )
-        both = (
-            pd.NA
-            if (pd.isna(ttft_max) or pd.isna(tpot_max))
-            else min(ttft_max, tpot_max)
-        )
-
-        tput_at_both = (
-            _value_at_concurrency(tput_group_df, conc_col, cfg, both)
-            if tput_group_df is not None
-            else pd.NA
-        )
-        ttft_at_both = (
-            _value_at_concurrency(ttft_group_df, conc_col, cfg, both)
-            if ttft_group_df is not None
-            else pd.NA
-        )
-        tpot_at_both = (
-            _value_at_concurrency(tpot_group_df, conc_col, cfg, both)
-            if tpot_group_df is not None
-            else pd.NA
-        )
-
-        rows.append(
-            {
-                "Configuration": cfg,
-                f"Max {conc_col} (TTFT ≤ {args.ttft_max_ms:g} ms)": ttft_max,
-                f"Max {conc_col} (TPOT ≤ {args.tpot_max_ms:g} ms)": tpot_max,
-                f"Max {conc_col} (Both)": both,
-                "Output Tput @ Both (tok/s)": tput_at_both,
-                "TTFT @ Both (ms)": ttft_at_both,
-                "TPOT @ Both (ms)": tpot_at_both,
-            }
-        )
-
-    df = pd.DataFrame(rows)
-    for c in df.columns:
-        if c != "Configuration":
-            df[c] = pd.to_numeric(df[c], errors="coerce")
-    return df
-
-
 # -----------------------------
 # Plot helper
 # -----------------------------
@@ -750,21 +537,6 @@ def build_parser() -> argparse.ArgumentParser:
         default=100.0,
         help="Reference limit for TPOT plots (ms)",
     )
-
-    # ---- NEW: export options ----
-    parser.add_argument(
-        "--excel-out",
-        type=str,
-        default="perf_comparison.xlsx",
-        help="Write one sheet per (Model, Dataset, Input Len, Output Len).",
-    )
-    parser.add_argument(
-        "--csv-out-dir",
-        type=str,
-        default="",
-        help="If set, write per-group per-metric CSVs into this directory.",
-    )
-
     return parser
 
 
@@ -885,6 +657,7 @@ def maybe_write_plot(
         markers=True,
     )
 
+    # Ensure plot hover + y tick labels are also 2 decimals.
     fig.update_traces(hovertemplate="%{y:.2f}<extra></extra>")
     fig.update_yaxes(tickformat=".2f")
 
@@ -957,151 +730,87 @@ def write_report_group_first(
         for metric_label, (df, _) in metric_cache.items()
     }
 
-    csv_dir = Path(args.csv_out_dir) if args.csv_out_dir else None
-    if csv_dir:
-        csv_dir.mkdir(parents=True, exist_ok=True)
+    with open("perf_comparison.html", "w", encoding="utf-8") as main_fh:
+        main_fh.write('<meta charset="utf-8">\n')
+        for gkey in group_keys:
+            gkey_tuple = normalize_group_key(gkey)
+            suffix = build_group_suffix(group_cols_canonical, gkey_tuple)
+            sub_path = group_filename(gkey_tuple)
+            group_header = (
+                '<div style="font-size: 1.4em; font-weight: 700; '
+                'margin: 18px 0 10px 0;">'
+                f"{_html.escape(suffix)}"
+                "</div>\n"
+            )
 
-    excel_path = args.excel_out or "perf_comparison.xlsx"
-    with pd.ExcelWriter(excel_path, engine="openpyxl") as xw:
-        # ---- Environment sheet (first) ----
-        env_sheet = _sanitize_sheet_name("Environment")
-        env_df = _load_env_df_for_inputs(args, files)
-        if env_df is None or env_df.empty:
-            pd.DataFrame(
-                [
-                    {
-                        "Section": "Environment",
-                        "Key": "vllm_env.txt",
-                        "Value": "NOT FOUND (or empty)",
-                    }
-                ]
-            ).to_excel(xw, sheet_name=env_sheet, index=False)
-        else:
-            env_df.to_excel(xw, sheet_name=env_sheet, index=False)
-        with open("perf_comparison.html", "w", encoding="utf-8") as main_fh:
-            main_fh.write('<meta charset="utf-8">\n')
-            for gkey in group_keys:
-                gkey_tuple = normalize_group_key(gkey)
-                suffix = build_group_suffix(group_cols_canonical, gkey_tuple)
-                sub_path = group_filename(gkey_tuple)
-                group_header = (
-                    '<div style="font-size: 1.4em; font-weight: 700; '
-                    'margin: 18px 0 10px 0;">'
-                    f"{_html.escape(suffix)}"
-                    "</div>\n"
+            main_fh.write(group_header)
+            with open(sub_path, "w", encoding="utf-8") as sub_fh:
+                sub_fh.write('<meta charset="utf-8">\n')
+                sub_fh.write(group_header)
+                tput_group_df = None
+                ttft_group_df = None
+                tpot_group_df = None
+                conc_col = args.xaxis
+
+                for metric_label in plan.data_cols:
+                    gb = metric_groupbys[metric_label]
+                    df_sorted, raw_data_cols = metric_cache[metric_label]
+
+                    try:
+                        group_df = gb.get_group(gkey)
+                    except KeyError:
+                        missing = (
+                            '<div style="font-size: 1.1em; font-weight: 600; '
+                            'margin: 10px 0;">'
+                            f"{_html.escape(metric_label)} — missing for this group"
+                            "</div>\n"
+                        )
+
+                        main_fh.write(missing)
+                        sub_fh.write(missing)
+                        continue
+
+                    if conc_col not in group_df.columns:
+                        conc_col = _find_concurrency_col(group_df)
+
+                    mn = metric_label.lower().strip()
+                    if "tok/s" in mn:
+                        tput_group_df = group_df
+                    elif "ttft" in mn:
+                        ttft_group_df = group_df
+                    elif mn in ("p99", "median") or "tpot" in mn:
+                        tpot_group_df = group_df
+
+                    display_group = group_df.drop(
+                        columns=group_cols_canonical, errors="ignore"
+                    )
+
+                    html = render_metric_table_html(
+                        display_group, metric_label, suffix, args
+                    )
+                    main_fh.write(html)
+                    sub_fh.write(html)
+
+                    maybe_write_plot(
+                        main_fh,
+                        sub_fh,
+                        group_df=group_df,
+                        raw_data_cols=raw_data_cols,
+                        metric_label=metric_label,
+                        y_axis_col=y_axis_col,
+                        args=args,
+                    )
+
+                summary_html = build_valid_max_concurrency_summary_html(
+                    tput_group_df=tput_group_df,
+                    ttft_group_df=ttft_group_df,
+                    tpot_group_df=tpot_group_df,
+                    conc_col=conc_col,
+                    args=args,
                 )
-
-                main_fh.write(group_header)
-
-                sheet = _group_to_sheet_base(group_cols_canonical, gkey_tuple)
-                sheet_base = sheet
-                dedup_i = 1
-                while sheet in xw.sheets:
-                    dedup_i += 1
-                    sheet = _sanitize_sheet_name(f"{sheet_base}_{dedup_i}")
-
-                excel_blocks: list[tuple[str, pd.DataFrame]] = []
-
-                with open(sub_path, "w", encoding="utf-8") as sub_fh:
-                    sub_fh.write('<meta charset="utf-8">\n')
-                    sub_fh.write(group_header)
-                    tput_group_df = None
-                    ttft_group_df = None
-                    tpot_group_df = None
-                    conc_col = args.xaxis
-
-                    for metric_label in plan.data_cols:
-                        gb = metric_groupbys[metric_label]
-                        df_sorted, raw_data_cols = metric_cache[metric_label]
-
-                        try:
-                            group_df = gb.get_group(gkey)
-                        except KeyError:
-                            missing = (
-                                '<div style="font-size: 1.1em; font-weight: 600; '
-                                'margin: 10px 0;">'
-                                f"{_html.escape(metric_label)} — missing for this group"
-                                "</div>\n"
-                            )
-                            main_fh.write(missing)
-                            sub_fh.write(missing)
-                            continue
-
-                        if conc_col not in group_df.columns:
-                            conc_col = _find_concurrency_col(group_df)
-
-                        mn = metric_label.lower().strip()
-                        if "tok/s" in mn:
-                            tput_group_df = group_df
-                        elif "ttft" in mn:
-                            ttft_group_df = group_df
-                        elif mn in ("p99", "median") or "tpot" in mn:
-                            tpot_group_df = group_df
-
-                        display_group = group_df.drop(
-                            columns=group_cols_canonical, errors="ignore"
-                        )
-
-                        html = render_metric_table_html(
-                            display_group, metric_label, suffix, args
-                        )
-                        main_fh.write(html)
-                        sub_fh.write(html)
-
-                        maybe_write_plot(
-                            main_fh,
-                            sub_fh,
-                            group_df=group_df,
-                            raw_data_cols=raw_data_cols,
-                            metric_label=metric_label,
-                            y_axis_col=y_axis_col,
-                            args=args,
-                        )
-
-                        excel_blocks.append(
-                            (metric_label, display_group.reset_index(drop=True))
-                        )
-                        if csv_dir:
-                            fn = _safe_filename(
-                                f"{sheet}__{metric_label}".replace(" ", "_").replace(
-                                    "/", "_"
-                                )
-                            )
-                            display_group.to_csv(csv_dir / f"{fn}.csv", index=False)
-
-                    summary_html = build_valid_max_concurrency_summary_html(
-                        tput_group_df=tput_group_df,
-                        ttft_group_df=ttft_group_df,
-                        tpot_group_df=tpot_group_df,
-                        conc_col=conc_col,
-                        args=args,
-                    )
-                    if summary_html:
-                        main_fh.write(summary_html)
-                        sub_fh.write(summary_html)
-
-                    summary_df = build_valid_max_concurrency_summary_df(
-                        tput_group_df=tput_group_df,
-                        ttft_group_df=ttft_group_df,
-                        tpot_group_df=tpot_group_df,
-                        conc_col=conc_col,
-                        args=args,
-                    )
-                    if summary_df is not None:
-                        excel_blocks.append(
-                            ("Valid Max Concurrency Summary", summary_df)
-                        )
-                        if csv_dir:
-                            fn = _safe_filename(
-                                f"{sheet}__Valid_Max_Concurrency_Summary"
-                            )
-                            summary_df.to_csv(csv_dir / f"{fn}.csv", index=False)
-
-                _write_tables_to_excel_sheet(xw, sheet, excel_blocks)
-
-    print(f"Wrote Excel: {excel_path}")
-    if csv_dir:
-        print(f"Wrote CSVs under: {csv_dir}")
+                if summary_html:
+                    main_fh.write(summary_html)
+                    sub_fh.write(summary_html)
 
 
 def main():

.buildkite/performance-benchmarks/scripts/convert-results-json-to-markdown.py

@@ -393,7 +393,7 @@ if __name__ == "__main__":
     with open(results_folder / md_file, "w") as f:
         results = read_markdown(
             "../.buildkite/performance-benchmarks/"
-            "performance-benchmarks-descriptions.md"
+            + "performance-benchmarks-descriptions.md"
         )
         results = results.format(
             latency_tests_markdown_table=latency_md_table,

.buildkite/performance-benchmarks/scripts/run-performance-benchmarks.sh

@@ -1,4 +1,6 @@
 #!/bin/bash
+
+# This script should be run inside the CI process
 # This script assumes that we are already inside the vllm/ directory
 # Benchmarking results will be available inside vllm/benchmarks/results/
 
@@ -7,11 +9,6 @@
 set -x
 set -o pipefail
 
-# Environment-driven debug controls (like ON_CPU=1)
-DRY_RUN="${DRY_RUN:-0}"
-MODEL_FILTER="${MODEL_FILTER:-}"
-DTYPE_FILTER="${DTYPE_FILTER:-}"
-
 check_gpus() {
   if command -v nvidia-smi; then
     # check the number of GPUs and GPU type.
@@ -28,9 +25,9 @@ check_gpus() {
     echo "Need at least 1 GPU to run benchmarking."
     exit 1
   fi
-
+  
   declare -g arch_suffix=''
-
+  
   if command -v nvidia-smi; then
     declare -g gpu_type=$(nvidia-smi --query-gpu=name --format=csv,noheader | awk '{print $2}')
   elif command -v amd-smi; then
@@ -115,12 +112,13 @@ json2envs() {
 }
 
 wait_for_server() {
+  # wait for vllm server to start
+  # return 1 if vllm server crashes
   local timeout_val="1200"
   timeout "$timeout_val" bash -c '
-    until curl -sf http://localhost:8000/v1/models >/dev/null; do
+    until curl -X POST localhost:8000/v1/completions; do
       sleep 1
-    done
-  '
+    done' && return 0 || return 1
 }
 
 kill_processes_launched_by_current_bash() {
@@ -183,20 +181,19 @@ upload_to_buildkite() {
   $BUILDKITE_AGENT_COMMAND artifact upload "$RESULTS_FOLDER/*"
 }
 
-run_benchmark_tests() {
-  # run benchmark tests using `vllm bench <test_type>` command
-  # $1: test type (latency or throughput)
-  # $2: a json file specifying test cases
+run_latency_tests() {
+  # run latency tests using `vllm bench latency` command
+  # $1: a json file specifying latency test cases
 
-  local test_type=$1
-  local test_file=$2
+  local latency_test_file
+  latency_test_file=$1
 
-  # Iterate over tests
-  jq -c '.[]' "$test_file" | while read -r params; do
+  # Iterate over latency tests
+  jq -c '.[]' "$latency_test_file" | while read -r params; do
     # get the test name, and append the GPU type back to it.
     test_name=$(echo "$params" | jq -r '.test_name')
-    if [[ ! "$test_name" =~ ^${test_type}_ ]]; then
-      echo "In ${test_type}-test.json, test_name must start with \"${test_type}_\"."
+    if [[ ! "$test_name" =~ ^latency_ ]]; then
+      echo "In latency-test.json, test_name must start with \"latency_\"."
       exit 1
     fi
 
@@ -207,15 +204,15 @@ run_benchmark_tests() {
     fi
 
     # get arguments
-    bench_params=$(echo "$params" | jq -r '.parameters')
-    bench_args=$(json2args "$bench_params")
-    bench_environment_variables=$(echo "$params" | jq -r '.environment_variables')
-    bench_envs=$(json2envs "$bench_environment_variables")
+    latency_params=$(echo "$params" | jq -r '.parameters')
+    latency_args=$(json2args "$latency_params")
+    latency_environment_variables=$(echo "$params" | jq -r '.environment_variables')
+    latency_envs=$(json2envs "$latency_environment_variables")
 
     # check if there is enough GPU to run the test
-    tp=$(echo "$bench_params" | jq -r '.tensor_parallel_size')
+    tp=$(echo "$latency_params" | jq -r '.tensor_parallel_size')
     if [[ "$ON_CPU" == "1" ]]; then
-      pp=$(echo "$bench_params" | jq -r '.pipeline_parallel_size // 1')
+      pp=$(echo "$latency_params" | jq -r '.pipeline_parallel_size // 1')
       world_size=$(($tp*$pp))
       if [[ $numa_count -lt $world_size  && -z "${REMOTE_HOST}" ]]; then
         echo "Required world-size $world_size but only $numa_count NUMA nodes found. Skip testcase $test_name."
@@ -228,42 +225,118 @@ run_benchmark_tests() {
       fi
     fi
 
-    bench_command=" $bench_envs vllm bench $test_type \
+    latency_command=" $latency_envs vllm bench latency \
       --output-json $RESULTS_FOLDER/${test_name}.json \
-      $bench_args"
+      $latency_args"
 
     echo "Running test case $test_name"
-    echo "${test_type^} command: $bench_command"
+    echo "Latency command: $latency_command"
 
-    # recording benchmarking command and GPU command
+    # recoding benchmarking command ang GPU command
     jq_output=$(jq -n \
-      --arg command "$bench_command" \
+      --arg latency "$latency_command" \
       --arg gpu "$gpu_type" \
-      --arg test_type "$test_type" \
       '{
-        ($test_type + "_command"): $command,
+        latency_command: $latency,
         gpu_type: $gpu
       }')
     echo "$jq_output" >"$RESULTS_FOLDER/$test_name.commands"
 
     # run the benchmark
-    eval "$bench_command"
+    eval "$latency_command"
 
     kill_gpu_processes
 
   done
 }
 
-run_latency_tests() { run_benchmark_tests "latency" "$1"; }
-run_startup_tests() { run_benchmark_tests "startup" "$1"; }
-run_throughput_tests() { run_benchmark_tests "throughput" "$1"; }
+run_throughput_tests() {
+  # run throughput tests using `vllm bench throughput`
+  # $1: a json file specifying throughput test cases
 
-merge_serving_tests_stream() {
-  # Emit merged serving test objects, optionally filtered by MODEL_FILTER/DTYPE_FILTER in DRY_RUN mode.
-  # This helper does NOT modify JSON; it only filters the stream in dry-run mode.
-  local serving_test_file="$1"
-  # shellcheck disable=SC2016
-  local merged='
+  local throughput_test_file
+  throughput_test_file=$1
+
+  # Iterate over throughput tests
+  jq -c '.[]' "$throughput_test_file" | while read -r params; do
+    # get the test name, and append the GPU type back to it.
+    test_name=$(echo "$params" | jq -r '.test_name')
+    if [[ ! "$test_name" =~ ^throughput_ ]]; then
+      echo "In throughput-test.json, test_name must start with \"throughput_\"."
+      exit 1
+    fi
+
+    # if TEST_SELECTOR is set, only run the test cases that match the selector
+    if [[ -n "$TEST_SELECTOR" ]] && [[ ! "$test_name" =~ $TEST_SELECTOR ]]; then
+      echo "Skip test case $test_name."
+      continue
+    fi
+
+    # get arguments
+    throughput_params=$(echo "$params" | jq -r '.parameters')
+    throughput_args=$(json2args "$throughput_params")
+    throughput_environment_variables=$(echo "$params" | jq -r '.environment_variables')
+    throughput_envs=$(json2envs "$throughput_environment_variables")
+
+    # check if there is enough GPU to run the test
+    tp=$(echo "$throughput_params" | jq -r '.tensor_parallel_size')
+    if [[ "$ON_CPU" == "1" ]]; then
+      pp=$(echo "$throughput_params" | jq -r '.pipeline_parallel_size // 1')
+      world_size=$(($tp*$pp))
+      if [[ $numa_count -lt $world_size  && -z "${REMOTE_HOST}" ]]; then
+        echo "Required world-size $world_size but only $numa_count NUMA nodes found. Skip testcase $test_name."
+        continue
+      fi
+    else
+      if [[ $gpu_count -lt $tp ]]; then
+        echo "Required tensor-parallel-size $tp but only $gpu_count GPU found. Skip testcase $test_name."
+        continue
+      fi
+    fi
+
+    throughput_command=" $throughput_envs vllm bench throughput \
+      --output-json $RESULTS_FOLDER/${test_name}.json \
+      $throughput_args"
+
+    echo "Running test case $test_name"
+    echo "Throughput command: $throughput_command"
+    # recoding benchmarking command ang GPU command
+    jq_output=$(jq -n \
+      --arg command "$throughput_command" \
+      --arg gpu "$gpu_type" \
+      '{
+        throughput_command: $command,
+        gpu_type: $gpu
+      }')
+    echo "$jq_output" >"$RESULTS_FOLDER/$test_name.commands"
+
+    # run the benchmark
+    eval "$throughput_command"
+
+    kill_gpu_processes
+
+  done
+}
+
+run_serving_tests() {
+  # run serving tests using `vllm bench serve` command
+  # $1: a json file specifying serving test cases
+  #
+  # Supported JSON formats:
+  # 1) Plain format: top-level array
+  #    [ { "test_name": "...", "server_parameters": {...}, ... }, ... ]
+  #
+  # 2) Default parameters field + plain format tests
+  #    {
+  #      "defaults": { ... },
+  #      "tests": [ { "test_name": "...", "server_parameters": {...}, ... }, ... ]
+  #    }
+
+  local serving_test_file
+  serving_test_file=$1
+
+  # Iterate over serving tests
+  jq -c '
     if type == "array" then
       # Plain format: test cases array
       .[]
@@ -285,50 +358,7 @@ merge_serving_tests_stream() {
     else
       error("Unsupported serving test file format: must be array or object with .tests")
     end
-  '
-
-  jq -c "$merged" "$serving_test_file" | \
-  if [[ "${DRY_RUN:-0}" == "1" && ( "${MODEL_FILTER}${DTYPE_FILTER}" != "" ) ]]; then
-    jq -c --arg model "$MODEL_FILTER" --arg dtype "$DTYPE_FILTER" '
-      select((($model|length)==0)
-             or ((.server_parameters.model // "") == $model)
-             or ((.client_parameters.model // "") == $model))
-      | select((($dtype|length)==0) or ((.server_parameters.dtype // "") == $dtype))
-    '
-  else
-    cat
-  fi
-}
-
-run_serving_tests() {
-  # run serving tests using `vllm bench serve` command
-  # $1: a json file specifying serving test cases
-  #
-  # Supported JSON formats:
-  # 1) Plain format: top-level array
-  #    [ { "test_name": "...", "server_parameters": {...}, ... }, ... ]
-  #
-  # 2) Default parameters field + plain format tests
-  #    {
-  #      "defaults": { ... },
-  #      "tests": [ { "test_name": "...", "server_parameters": {...}, ... }, ... ]
-  #    }
-
-  local serving_test_file
-  serving_test_file=$1
-
-  # In dry-run mode, if filters are provided but no tests match, fail fast.
-  if [[ "${DRY_RUN:-0}" == "1" && ( "${MODEL_FILTER}${DTYPE_FILTER}" != "" ) ]]; then
-    local count
-    count=$(merge_serving_tests_stream "$serving_test_file" | wc -l | tr -d ' ')
-    if [[ "$count" -eq 0 ]]; then
-      echo "No matching serving tests found in $serving_test_file for model='$MODEL_FILTER' dtype='$DTYPE_FILTER'." >&2
-      return 0
-    fi
-  fi
-
-  # Iterate over serving tests (merged + optional filtered stream)
-  merge_serving_tests_stream "$serving_test_file" | while read -r params; do
+  ' "$serving_test_file" | while read -r params; do
     # get the test name, and append the GPU type back to it.
     test_name=$(echo "$params" | jq -r '.test_name')
     if [[ ! "$test_name" =~ ^serving_ ]]; then
@@ -397,7 +427,7 @@ run_serving_tests() {
     echo "Server command: $server_command"
     # support remote vllm server
     client_remote_args=""
-    if [[ -z "${REMOTE_HOST}" && "${DRY_RUN:-0}" != "1" ]]; then
+    if [[ -z "${REMOTE_HOST}" ]]; then
       bash -c "$server_command" &
       server_pid=$!
       # wait until the server is alive
@@ -408,9 +438,6 @@ run_serving_tests() {
         echo ""
         echo "vLLM failed to start within the timeout period."
       fi
-    elif [[ "${DRY_RUN:-0}" == "1" ]]; then
-        # dry-run: don't start server
-        echo "Dry Run."
     else
       server_command="Using Remote Server $REMOTE_HOST $REMOTE_PORT"
       if [[ ${REMOTE_PORT} ]]; then
@@ -420,39 +447,34 @@ run_serving_tests() {
       fi
     fi
 
-    # save the compilation mode and optimization level on the serving results
-    # whenever they are set
-    compilation_config_mode=$(echo "$server_params" | jq -r '."compilation_config.mode" // empty')
-    optimization_level=$(echo "$server_params" | jq -r '.optimization_level // empty')
-
     # iterate over different QPS
     for qps in $qps_list; do
       # remove the surrounding single quote from qps
       if [[ "$qps" == *"inf"* ]]; then
+        echo "qps was $qps"
         qps="inf"
+        echo "now qps is $qps"
       fi
 
       # iterate over different max_concurrency
       for max_concurrency in $max_concurrency_list; do
         new_test_name=$test_name"_qps_"$qps"_concurrency_"$max_concurrency
         echo " new test name $new_test_name"
-        # pass the tensor parallel size, the compilation mode, and the optimization
-        # level to the client so that they can be used on the benchmark dashboard
+        # pass the tensor parallel size to the client so that it can be displayed
+        # on the benchmark dashboard
         client_command="vllm bench serve \
           --save-result \
           --result-dir $RESULTS_FOLDER \
           --result-filename ${new_test_name}.json \
           --request-rate $qps \
           --max-concurrency $max_concurrency \
-          --metadata tensor_parallel_size=$tp compilation_config.mode=$compilation_config_mode optimization_level=$optimization_level \
+          --metadata "tensor_parallel_size=$tp" \
           $client_args $client_remote_args "
 
         echo "Running test case $test_name with qps $qps"
         echo "Client command: $client_command"
 
-        if [[ "${DRY_RUN:-0}" != "1" ]]; then
-          bash -c "$client_command"
-        fi
+        bash -c "$client_command"
 
         # record the benchmarking commands
         jq_output=$(jq -n \
@@ -470,15 +492,12 @@ run_serving_tests() {
     done
 
     # clean up
-    if [[ "${DRY_RUN:-0}" != "1" ]]; then
-      kill -9 $server_pid
-      kill_gpu_processes
-    fi
+    kill -9 $server_pid
+    kill_gpu_processes
   done
 }
 
 main() {
-
   local ARCH
   ARCH=''
   if [[ "$ON_CPU" == "1" ]]; then
@@ -488,13 +507,7 @@ main() {
      check_gpus
      ARCH="$arch_suffix"
   fi
-
-  # DRY_RUN does not execute vLLM; do not require HF_TOKEN.
-  if [[ "${DRY_RUN:-0}" != "1" ]]; then
-    check_hf_token
-  else
-    echo "DRY_RUN=1 -> skip HF_TOKEN validation"
-  fi
+  check_hf_token
 
   # dependencies
   (which wget && which curl) || (apt-get update && apt-get install -y wget curl)
@@ -515,18 +528,12 @@ main() {
 
   # dump vllm info via vllm collect-env
   env_output=$(vllm collect-env)
+
   echo "$env_output" >"$RESULTS_FOLDER/vllm_env.txt"
 
   # benchmarking
-  run_serving_tests $QUICK_BENCHMARK_ROOT/tests/"${SERVING_JSON:-serving-tests$ARCH.json}" || exit $?
-
-  if [[ "${DRY_RUN:-0}" == "1" ]]; then
-    echo "DRY_RUN=1 -> skip latency/startup/throughput suites"
-    exit 0
-  fi
-
+  run_serving_tests $QUICK_BENCHMARK_ROOT/tests/"${SERVING_JSON:-serving-tests$ARCH.json}"
   run_latency_tests $QUICK_BENCHMARK_ROOT/tests/"${LATENCY_JSON:-latency-tests$ARCH.json}"
-  run_startup_tests $QUICK_BENCHMARK_ROOT/tests/"${STARTUP_JSON:-startup-tests$ARCH.json}"
   run_throughput_tests $QUICK_BENCHMARK_ROOT/tests/"${THROUGHPUT_JSON:-throughput-tests$ARCH.json}"
 
   # postprocess benchmarking results

.buildkite/performance-benchmarks/tests/serving-tests-cpu-embed.json

@@ -1,41 +0,0 @@
-{
-  "defaults": {
-    "qps_list": [
-      "inf"
-    ],
-    "max_concurrency_list": [
-      32,
-      64,
-      128
-    ],
-    "server_environment_variables": {
-      "VLLM_RPC_TIMEOUT": 100000,
-      "VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1,
-      "VLLM_ENGINE_ITERATION_TIMEOUT_S": 120,
-      "VLLM_CPU_SGL_KERNEL": 1,
-      "VLLM_CPU_KVCACHE_SPACE": 40
-    },
-    "server_parameters": {
-      "dtype": "bfloat16",
-      "model": "jinaai/jina-embeddings-v3",
-      "trust_remote_code": ""
-    },
-    "client_parameters": {
-      "model": "jinaai/jina-embeddings-v3",
-      "backend": "openai-embeddings",
-      "endpoint": "/v1/embeddings",
-      "dataset_name": "sharegpt",
-      "dataset_path": "ShareGPT_V3_unfiltered_cleaned_split.json",
-      "num_prompts": 200
-    }
-  },
-  "tests": [
-    {
-      "test_name": "serving_jina_embed_v3_tp1_sharegpt",
-      "server_parameters": {
-        "tensor_parallel_size": 1
-      },
-      "client_parameters": {}
-    }
-  ]
-}

.buildkite/performance-benchmarks/tests/serving-tests-cpu-text.json

@@ -1,283 +0,0 @@
-{
-  "defaults": {
-    "qps_list": [
-      "inf"
-    ],
-    "max_concurrency_list": [12, 16, 24, 32, 64, 128, 200],
-    "server_environment_variables": {
-      "VLLM_RPC_TIMEOUT": 100000,
-      "VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1,
-      "VLLM_ENGINE_ITERATION_TIMEOUT_S": 120,
-      "VLLM_CPU_SGL_KERNEL": 1,
-      "VLLM_CPU_KVCACHE_SPACE": 40
-    },
-    "server_parameters": {
-      "model": "meta-llama/Llama-3.1-8B-Instruct",
-      "tensor_parallel_size": 1,
-      "dtype": "bfloat16",
-      "distributed_executor_backend": "mp",
-      "block_size": 128,
-      "trust_remote_code": "",
-      "disable_log_stats": "",
-      "max_num_batched_tokens": 2048,
-      "max_num_seqs": 256
-    },
-    "client_parameters": {
-      "model": "meta-llama/Llama-3.1-8B-Instruct",
-      "backend": "vllm",
-      "ignore-eos": "",
-      "num_prompts": 200
-    }
-  },
-  "tests": [
-    {
-      "test_name": "serving_llama8B_tp1_sharegpt",
-      "server_parameters": {
-        "tensor_parallel_size": 1
-      },
-      "client_parameters": {
-        "dataset_name": "sharegpt",
-        "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json"
-      }
-    },
-    {
-      "test_name": "serving_llama8B_tp2_sharegpt",
-      "server_parameters": {
-        "tensor_parallel_size": 2
-      },
-      "client_parameters": {
-        "dataset_name": "sharegpt",
-        "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json"
-      }
-    },
-    {
-      "test_name": "serving_llama8B_tp1_random_128_128",
-      "server_parameters": {
-        "tensor_parallel_size": 1
-      },
-      "client_parameters": {
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_llama8B_tp2_random_128_128",
-      "server_parameters": {
-        "tensor_parallel_size": 2
-      },
-      "client_parameters": {
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_llama8B_tp4_random_128_128",
-      "server_parameters": {
-        "tensor_parallel_size": 4
-      },
-      "client_parameters": {
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_llama8B_tp1_random_128_2048",
-      "server_parameters": {
-        "tensor_parallel_size": 1
-      },
-      "client_parameters": {
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 2048
-      }
-    },
-    {
-      "test_name": "serving_llama8B_tp2_random_128_2048",
-      "server_parameters": {
-        "tensor_parallel_size": 2
-      },
-      "client_parameters": {
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 2048
-      }
-    },
-    {
-      "test_name": "serving_llama8B_tp4_random_128_2048",
-      "server_parameters": {
-        "tensor_parallel_size": 4
-      },
-      "client_parameters": {
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 2048
-      }
-    },
-    {
-      "test_name": "serving_llama8B_tp1_random_2048_128",
-      "server_parameters": {
-        "tensor_parallel_size": 1
-      },
-      "client_parameters": {
-        "dataset_name": "random",
-        "random-input-len": 2048,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_llama8B_tp2_random_2048_128",
-      "server_parameters": {
-        "tensor_parallel_size": 2
-      },
-      "client_parameters": {
-        "dataset_name": "random",
-        "random-input-len": 2048,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_llama8B_tp4_random_2048_128",
-      "server_parameters": {
-        "tensor_parallel_size": 4
-      },
-      "client_parameters": {
-        "dataset_name": "random",
-        "random-input-len": 2048,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_llama8B_int4_tp1_random_128_128",
-      "server_parameters": {
-        "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
-        "tensor_parallel_size": 1
-      },
-      "client_parameters": {
-        "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_llama8B_int4_tp2_random_128_128",
-      "server_parameters": {
-        "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
-        "tensor_parallel_size": 2
-      },
-      "client_parameters": {
-        "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_llama8B_int4_tp4_random_128_128",
-      "server_parameters": {
-        "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
-        "tensor_parallel_size": 4
-      },
-      "client_parameters": {
-        "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_llama3B_tp1_random_128_128",
-      "server_parameters": {
-        "model": "meta-llama/Llama-3.2-3B-Instruct",
-        "tensor_parallel_size": 1
-      },
-      "client_parameters": {
-        "model": "meta-llama/Llama-3.2-3B-Instruct",
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_granite2B_tp1_random_128_128",
-      "server_parameters": {
-        "model": "ibm-granite/granite-3.2-2b-instruct",
-        "tensor_parallel_size": 1
-      },
-      "client_parameters": {
-        "model": "ibm-granite/granite-3.2-2b-instruct",
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_qwen1.7B_tp1_random_128_128",
-      "server_parameters": {
-        "model": "Qwen/Qwen3-1.7B",
-        "tensor_parallel_size": 1
-      },
-      "client_parameters": {
-        "model": "Qwen/Qwen3-1.7B",
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_qwen4B_tp1_random_128_128",
-      "server_parameters": {
-        "model": "Qwen/Qwen3-4B",
-        "tensor_parallel_size": 1
-      },
-      "client_parameters": {
-        "model": "Qwen/Qwen3-4B",
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_qwen8B_tp1_random_128_128",
-      "server_parameters": {
-        "model": "Qwen/Qwen3-8B",
-        "tensor_parallel_size": 1
-      },
-      "client_parameters": {
-        "model": "Qwen/Qwen3-8B",
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_glm9B_tp1_random_128_128",
-      "server_parameters": {
-        "model": "zai-org/glm-4-9b-hf",
-        "tensor_parallel_size": 1
-      },
-      "client_parameters": {
-        "model": "zai-org/glm-4-9b-hf",
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_gemma7B_tp1_random_128_128",
-      "server_parameters": {
-        "model": "google/gemma-7b",
-        "tensor_parallel_size": 1
-      },
-      "client_parameters": {
-        "model": "google/gemma-7b",
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 128
-      }
-    }
-  ]
-}

.buildkite/performance-benchmarks/tests/serving-tests-cpu.json

@@ -148,6 +148,136 @@
         "random-input-len": 2048,
         "random-output-len": 128
       }
+    },
+    {
+      "test_name": "serving_llama8B_int4_tp1_random_128_128",
+      "server_parameters": {
+        "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
+        "tensor_parallel_size": 1
+      },
+      "client_parameters": {
+        "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
+        "dataset_name": "random",
+        "random-input-len": 128,
+        "random-output-len": 128
+      }
+    },
+    {
+      "test_name": "serving_llama8B_int4_tp2_random_128_128",
+      "server_parameters": {
+        "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
+        "tensor_parallel_size": 2
+      },
+      "client_parameters": {
+        "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
+        "dataset_name": "random",
+        "random-input-len": 128,
+        "random-output-len": 128
+      }
+    },
+    {
+      "test_name": "serving_llama8B_int4_tp4_random_128_128",
+      "server_parameters": {
+        "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
+        "tensor_parallel_size": 4
+      },
+      "client_parameters": {
+        "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
+        "dataset_name": "random",
+        "random-input-len": 128,
+        "random-output-len": 128
+      }
+    },
+    {
+      "test_name": "serving_llama3B_tp1_random_128_128",
+      "server_parameters": {
+        "model": "meta-llama/Llama-3.2-3B-Instruct",
+        "tensor_parallel_size": 1
+      },
+      "client_parameters": {
+        "model": "meta-llama/Llama-3.2-3B-Instruct",
+        "dataset_name": "random",
+        "random-input-len": 128,
+        "random-output-len": 128
+      }
+    },
+    {
+      "test_name": "serving_granite2B_tp1_random_128_128",
+      "server_parameters": {
+        "model": "ibm-granite/granite-3.2-2b-instruct",
+        "tensor_parallel_size": 1
+      },
+      "client_parameters": {
+        "model": "ibm-granite/granite-3.2-2b-instruct",
+        "dataset_name": "random",
+        "random-input-len": 128,
+        "random-output-len": 128
+      }
+    },
+    {
+      "test_name": "serving_qwen1.7B_tp1_random_128_128",
+      "server_parameters": {
+        "model": "Qwen/Qwen3-1.7B",
+        "tensor_parallel_size": 1
+      },
+      "client_parameters": {
+        "model": "Qwen/Qwen3-1.7B",
+        "dataset_name": "random",
+        "random-input-len": 128,
+        "random-output-len": 128
+      }
+    },
+    {
+      "test_name": "serving_qwen4B_tp1_random_128_128",
+      "server_parameters": {
+        "model": "Qwen/Qwen3-4B",
+        "tensor_parallel_size": 1
+      },
+      "client_parameters": {
+        "model": "Qwen/Qwen3-4B",
+        "dataset_name": "random",
+        "random-input-len": 128,
+        "random-output-len": 128
+      }
+    },
+    {
+      "test_name": "serving_qwen8B_tp1_random_128_128",
+      "server_parameters": {
+        "model": "Qwen/Qwen3-8B",
+        "tensor_parallel_size": 1
+      },
+      "client_parameters": {
+        "model": "Qwen/Qwen3-8B",
+        "dataset_name": "random",
+        "random-input-len": 128,
+        "random-output-len": 128
+      }
+    },
+    {
+      "test_name": "serving_glm9B_tp1_random_128_128",
+      "server_parameters": {
+        "model": "zai-org/glm-4-9b-hf",
+        "tensor_parallel_size": 1
+      },
+      "client_parameters": {
+        "model": "zai-org/glm-4-9b-hf",
+        "dataset_name": "random",
+        "random-input-len": 128,
+        "random-output-len": 128
+      }
+    },
+    {
+      "test_name": "serving_gemma7B_tp1_random_128_128",
+      "server_parameters": {
+        "model": "google/gemma-7b",
+        "tensor_parallel_size": 1
+      },
+      "client_parameters": {
+        "model": "google/gemma-7b",
+        "dataset_name": "random",
+        "random-input-len": 128,
+        "random-output-len": 128
+      }
     }
   ]
 }

.buildkite/release-pipeline.yaml

@@ -176,6 +176,23 @@ steps:
         env:
           DOCKER_BUILDKIT: "1"
 
+      - block: "Build release image for x86_64 ROCm"
+        key: block-rocm-release-image-build
+        depends_on: ~
+
+      - label: "Build release image - x86_64 - ROCm"
+        depends_on: block-rocm-release-image-build
+        id: build-release-image-rocm
+        agents:
+          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"
+          # Build base image first
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --tag rocm/vllm-dev:base-$BUILDKITE_COMMIT --target final --progress plain -f docker/Dockerfile.rocm_base ."
+          # Build vLLM ROCm image using the base
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg BASE_IMAGE=rocm/vllm-dev:base-$BUILDKITE_COMMIT --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-rocm --target vllm-openai --progress plain -f docker/Dockerfile.rocm ."
+          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-rocm"
+
   - group: "Publish release images"
     key: "publish-release-images"
     steps:
@@ -257,14 +274,14 @@ steps:
           - input-release-version
           - build-wheels
 
-      - label: "Upload release wheels to PyPI"
+      - label: "Upload release wheels to PyPI and GitHub"
         depends_on:
           - block-upload-release-wheels
         id: upload-release-wheels
         agents:
           queue: small_cpu_queue_postmerge
         commands:
-          - "bash .buildkite/scripts/upload-release-wheels-pypi.sh"
+          - "bash .buildkite/scripts/upload-release-wheels.sh"
 
   # =============================================================================
   # ROCm Release Pipeline (x86_64 only)
@@ -459,7 +476,7 @@ steps:
       S3_BUCKET: "vllm-wheels"
 
   # ROCm Job 2: Build vLLM ROCm Wheel
-  - label: ":python: Build vLLM ROCm Wheel - x86_64"
+  - label: ":python: Build vLLM ROCm Wheel"
     id: build-rocm-vllm-wheel
     depends_on:
       - step: build-rocm-base-wheels
@@ -649,7 +666,7 @@ steps:
       VARIANT: "rocm700"
 
   # ROCm Job 5: Build ROCm Release Docker Image
-  - label: ":docker: Build release image - x86_64 - ROCm"
+  - label: ":rocm: :docker: Build ROCm Release Docker Image"
     id: build-rocm-release-image
     depends_on:
       - step: build-rocm-base-wheels

.buildkite/scripts/annotate-release.sh

@@ -11,36 +11,29 @@ fi
 buildkite-agent annotate --style 'info' --context 'release-workflow' << EOF
 To download the wheel (by commit):
 \`\`\`
-aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux_2_31_x86_64.whl .
-aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux_2_31_aarch64.whl .
+aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux1_x86_64.whl .
+aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux2014_aarch64.whl .
 
-(Optional) For CUDA 13.0:
-aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu130-cp38-abi3-manylinux_2_35_x86_64.whl .
-aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu130-cp38-abi3-manylinux_2_35_aarch64.whl .
-
-(Optional) For CPU:
-aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cpu-cp38-abi3-manylinux_2_35_x86_64.whl .
-aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cpu-cp38-abi3-manylinux_2_35_aarch64.whl .
+aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu129-cp38-abi3-manylinux1_x86_64.whl .
+aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu129-cp38-abi3-manylinux1_x86_64.whl .
 \`\`\`
 
+To download the wheel (by version):
+\`\`\`
+aws s3 cp s3://vllm-wheels/${RELEASE_VERSION}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux1_x86_64.whl .
+aws s3 cp s3://vllm-wheels/${RELEASE_VERSION}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux2014_aarch64.whl .
+
+aws s3 cp s3://vllm-wheels/${RELEASE_VERSION}+cu129/vllm-${RELEASE_VERSION}+cu129-cp38-abi3-manylinux1_x86_64.whl .
+aws s3 cp s3://vllm-wheels/${RELEASE_VERSION}+cu130/vllm-${RELEASE_VERSION}+cu130-cp38-abi3-manylinux1_x86_64.whl .
+\`\`\`
 
 To download and upload the image:
 
 \`\`\`
-# Download images:
-
 docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64
 docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64
-docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu130
-docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu130
 docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm-base
 docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm
-docker pull public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:v${RELEASE_VERSION}
-docker pull public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:v${RELEASE_VERSION}
-
-# Tag and push images:
-
-## CUDA
 
 docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64 vllm/vllm-openai:x86_64
 docker tag vllm/vllm-openai:x86_64 vllm/vllm-openai:latest-x86_64
@@ -48,70 +41,28 @@ docker tag vllm/vllm-openai:x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64
 docker push vllm/vllm-openai:latest-x86_64
 docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64
 
-docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu130 vllm/vllm-openai:x86_64-cu130
-docker tag vllm/vllm-openai:x86_64-cu130 vllm/vllm-openai:latest-x86_64-cu130
-docker tag vllm/vllm-openai:x86_64-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu130
-docker push vllm/vllm-openai:latest-x86_64-cu130
-docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu130
-
 docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64 vllm/vllm-openai:aarch64
 docker tag vllm/vllm-openai:aarch64 vllm/vllm-openai:latest-aarch64
 docker tag vllm/vllm-openai:aarch64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
 docker push vllm/vllm-openai:latest-aarch64
 docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
 
-docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu130 vllm/vllm-openai:aarch64-cu130
-docker tag vllm/vllm-openai:aarch64-cu130 vllm/vllm-openai:latest-aarch64-cu130
-docker tag vllm/vllm-openai:aarch64-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu130
-docker push vllm/vllm-openai:latest-aarch64-cu130
-docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu130
-
-## ROCm
-
-docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}
-docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT} vllm/vllm-openai-rocm:latest
-docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT} vllm/vllm-openai-rocm:v${RELEASE_VERSION}
-docker push vllm/vllm-openai-rocm:latest
-docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}
-
 docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm-base vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base
 docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base vllm/vllm-openai-rocm:latest-base
 docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base
 docker push vllm/vllm-openai-rocm:latest-base
 docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base
 
-## CPU
-
-docker tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:v${RELEASE_VERSION} vllm/vllm-openai-cpu:x86_64
-docker tag vllm/vllm-openai-cpu:x86_64 vllm/vllm-openai-cpu:latest-x86_64
-docker tag vllm/vllm-openai-cpu:x86_64 vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64
-docker push vllm/vllm-openai-cpu:latest-x86_64
-docker push vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64
-
-docker tag public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:v${RELEASE_VERSION} vllm/vllm-openai-cpu:arm64
-docker tag vllm/vllm-openai-cpu:arm64 vllm/vllm-openai-cpu:latest-arm64
-docker tag vllm/vllm-openai-cpu:arm64 vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64
-docker push vllm/vllm-openai-cpu:latest-arm64
-docker push vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64
-
-# Create multi-arch manifest:
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}
+docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT} vllm/vllm-openai-rocm:latest
+docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT} vllm/vllm-openai-rocm:v${RELEASE_VERSION}
+docker push vllm/vllm-openai-rocm:latest
+docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}
 
 docker manifest rm vllm/vllm-openai:latest
 docker manifest create vllm/vllm-openai:latest vllm/vllm-openai:latest-x86_64 vllm/vllm-openai:latest-aarch64
 docker manifest create vllm/vllm-openai:v${RELEASE_VERSION} vllm/vllm-openai:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
 docker manifest push vllm/vllm-openai:latest
 docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}
-
-docker manifest rm vllm/vllm-openai:latest-cu130
-docker manifest create vllm/vllm-openai:latest-cu130 vllm/vllm-openai:latest-x86_64-cu130 vllm/vllm-openai:latest-aarch64-cu130
-docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu130
-docker manifest push vllm/vllm-openai:latest-cu130
-docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu130
-
-docker manifest rm vllm/vllm-openai-cpu:latest || true
-docker manifest create vllm/vllm-openai-cpu:latest vllm/vllm-openai-cpu:latest-x86_64 vllm/vllm-openai-cpu:latest-arm64
-docker manifest create vllm/vllm-openai-cpu:v${RELEASE_VERSION} vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64
-docker manifest push vllm/vllm-openai-cpu:latest
-docker manifest push vllm/vllm-openai-cpu:v${RELEASE_VERSION}
 \`\`\`
-EOF
+EOF 

.buildkite/scripts/generate-nightly-index.py

@@ -112,7 +112,7 @@ def parse_from_filename(file: str) -> WheelFileInfo:
 
 def generate_project_list(subdir_names: list[str], comment: str = "") -> str:
     """
-    Generate project list HTML content linking to each project & variant subdirectory.
+    Generate project list HTML content linking to each project & variant sub-directory.
     """
     href_tags = []
     for name in sorted(subdir_names):
@@ -168,23 +168,23 @@ def generate_index_and_metadata(
         comment (str | None): Optional comment to include in the generated HTML files.
 
     First, parse all wheel files to extract metadata.
-    We need to collect all wheel files for each variant, and generate an index for it (in a subdirectory).
+    We need to collect all wheel files for each variant, and generate an index for it (in a sub-directory).
     The index for the default variant (if any) is generated in the root index directory.
 
     If `default_variant` is provided, all wheels must have variant suffixes, and the default variant index
     is purely a copy of the corresponding variant index, with only the links adjusted.
     Otherwise, all wheels without variant suffixes are treated as the default variant.
 
-    If `alias_to_default` is provided, an additional alias subdirectory is created, it has the same content
+    If `alias_to_default` is provided, an additional alias sub-directory is created, it has the same content
     as the default variant index, but the links are adjusted accordingly.
 
     Index directory structure:
         index_base_dir/ (hosted at wheels.vllm.ai/{nightly,$commit,$version}/)
-            index.html  # project list, linking to "vllm/" and other packages, and all variant subdirectories
+            index.html  # project list, linking to "vllm/" and other packages, and all variant sub-directories
             vllm/
                 index.html # package index, pointing to actual files in wheel_base_dir (relative path)
                 metadata.json # machine-readable metadata for all wheels in this package
-            cpu/ # cpu variant subdirectory
+            cpu/ # cpu variant sub-directory
                 index.html
                 vllm/
                     index.html
@@ -194,7 +194,7 @@ def generate_index_and_metadata(
                 vllm/
                     index.html
                     metadata.json
-            cu130/ # cu130 variant subdirectory
+            cu130/ # cu130 variant sub-directory
                 index.html
                 vllm/
                     index.html

.buildkite/scripts/hardware_ci/run-amd-test.sh

@@ -44,17 +44,6 @@ cleanup_docker() {
   fi
 }
 
-cleanup_network() {
-  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 docker-net; then
-    docker network rm docker-net
-  fi
-}
-
 # Call the cleanup docker function
 cleanup_docker
 
@@ -87,7 +76,7 @@ mkdir -p "${HF_CACHE}"
 HF_MOUNT="/root/.cache/huggingface"
 
 commands=$@
-echo "Raw commands: $commands"
+echo "Commands:$commands"
 
 commands=${commands//"pytest -v -s basic_correctness/test_basic_correctness.py"/"pytest -v -s basic_correctness/test_basic_correctness.py"}
 
@@ -169,9 +158,6 @@ if [[ $commands == *" entrypoints/llm "* ]]; then
   --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
@@ -179,6 +165,7 @@ echo "Final commands: $commands"
 # --ignore=entrypoints/openai/test_models.py <= Fails on MI250 but passes on MI300 as of 2025-03-13
 
 
+PARALLEL_JOB_COUNT=8
 MYPYTHONPATH=".."
 
 # Test that we're launching on the machine that has
@@ -189,33 +176,53 @@ if [[ -z "$render_gid" ]]; then
   exit 1
 fi
 
-if [[ $commands == *"VLLM_TEST_GROUP_NAME=mi325_4-2-node-tests-4-gpus-in-total"* ]]; then
-
-  export DCKR_VER=$(docker --version | sed 's/Docker version \(.*\), build .*/\1/')
-
-  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 "Failed to parse node commands! Exiting."
-      cleanup_network
-      exit 111
+# check if the command contains shard flag, we will run all shards in parallel because the host have 8 GPUs.
+if [[ $commands == *"--shard-id="* ]]; then
+  # assign job count as the number of shards used
+  commands=$(echo "$commands" | sed -E "s/--num-shards[[:blank:]]*=[[:blank:]]*[0-9]*/--num-shards=${PARALLEL_JOB_COUNT} /g" | sed 's/ \\ / /g')
+  for GPU in $(seq 0 $(($PARALLEL_JOB_COUNT-1))); do
+    # assign shard-id for each shard
+    commands_gpu=$(echo "$commands" | sed -E "s/--shard-id[[:blank:]]*=[[:blank:]]*[0-9]*/--shard-id=${GPU} /g" | sed 's/ \\ / /g')
+    echo "Shard ${GPU} commands:$commands_gpu"
+    echo "Render devices: $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES"
+    docker run \
+        --device /dev/kfd $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES \
+        --network=host \
+        --shm-size=16gb \
+        --group-add "$render_gid" \
+        --rm \
+        -e HIP_VISIBLE_DEVICES="${GPU}" \
+        -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}_${GPU}" \
+        "${image_name}" \
+        /bin/bash -c "${commands_gpu}" \
+        |& while read -r line; do echo ">>Shard $GPU: $line"; done &
+    PIDS+=($!)
+  done
+  #wait for all processes to finish and collect exit codes
+  for pid in "${PIDS[@]}"; do
+    wait "${pid}"
+    STATUS+=($?)
+  done
+  at_least_one_shard_with_tests=0
+  for st in "${STATUS[@]}"; do
+    if [[ ${st} -ne 0 ]] && [[ ${st} -ne 5 ]]; then
+      echo "One of the processes failed with $st"
+      exit "${st}"
+    elif [[ ${st} -eq 5 ]]; then
+      echo "Shard exited with status 5 (no tests collected) - treating as success"
+    else # This means st is 0
+      at_least_one_shard_with_tests=1
+    fi
+  done
+  if [[ ${#STATUS[@]} -gt 0 && ${at_least_one_shard_with_tests} -eq 0 ]]; then
+    echo "All shards reported no tests collected. Failing the build."
+    exit 1
   fi
 else
   echo "Render devices: $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES"

.buildkite/scripts/hardware_ci/run-cpu-distributed-smoke-test.sh

@@ -1,26 +0,0 @@
-#!/bin/bash
-set -euox pipefail
-
-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; 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 \
-    --endpoint /v1/completions
-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; 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 \
-    --endpoint /v1/completions
-kill -s SIGTERM $server_pid &

.buildkite/scripts/hardware_ci/run-cpu-test.sh

@@ -2,19 +2,119 @@
 
 # This script build the CPU docker image and run the offline inference inside the container.
 # It serves a sanity check for compilation and basic model usage.
-set -euox pipefail
+set -ex
 
 # allow to bind to different cores
 CORE_RANGE=${CORE_RANGE:-48-95}
+# used for TP/PP E2E test
+OMP_CORE_RANGE=${OMP_CORE_RANGE:-48-95}
 NUMA_NODE=${NUMA_NODE:-1}
-IMAGE_NAME="cpu-test-$NUMA_NODE"
-TIMEOUT_VAL=$1
-TEST_COMMAND=$2
 
-# building the docker image
-echo "--- :docker: Building Docker image"
-docker build --progress plain --tag "$IMAGE_NAME" --target vllm-test -f docker/Dockerfile.cpu .
+export CMAKE_BUILD_PARALLEL_LEVEL=32
+
+# Setup cleanup
+remove_docker_container() {
+    set -e;
+    docker rm -f cpu-test-"$NUMA_NODE" cpu-test-"$NUMA_NODE"-avx2 || true;
+}
+trap remove_docker_container EXIT
+remove_docker_container
+
+# Try building the docker image
+numactl -C "$CORE_RANGE" -N "$NUMA_NODE" docker build --progress plain --tag cpu-test-"$NUMA_NODE" --target vllm-test -f docker/Dockerfile.cpu .
+numactl -C "$CORE_RANGE" -N "$NUMA_NODE" docker build --progress plain --build-arg VLLM_CPU_DISABLE_AVX512="true" --tag cpu-test-"$NUMA_NODE"-avx2 --target vllm-test -f docker/Dockerfile.cpu .
 
 # Run the image, setting --shm-size=4g for tensor parallel.
-docker run --rm --cpuset-cpus=$CORE_RANGE --cpuset-mems=$NUMA_NODE -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN -e VLLM_CPU_KVCACHE_SPACE=16 -e VLLM_CPU_CI_ENV=1 -e VLLM_CPU_SIM_MULTI_NUMA=1 --shm-size=4g $IMAGE_NAME \
-        timeout $TIMEOUT_VAL bash -c "set -euox pipefail; echo \"--- Print packages\"; pip list; echo \"--- Running tests\"; ${TEST_COMMAND}"
+docker run -itd --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" --entrypoint /bin/bash -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN --env VLLM_CPU_KVCACHE_SPACE=16 --env VLLM_CPU_CI_ENV=1 -e E2E_OMP_THREADS="$OMP_CORE_RANGE" --shm-size=4g --name cpu-test-"$NUMA_NODE" cpu-test-"$NUMA_NODE"
+docker run -itd --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" --entrypoint /bin/bash -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN --env VLLM_CPU_KVCACHE_SPACE=16 --env VLLM_CPU_CI_ENV=1 -e E2E_OMP_THREADS="$OMP_CORE_RANGE" --shm-size=4g --name cpu-test-"$NUMA_NODE"-avx2 cpu-test-"$NUMA_NODE"-avx2
+
+function cpu_tests() {
+  set -e
+  export NUMA_NODE=$2
+
+  # list packages
+  docker exec cpu-test-"$NUMA_NODE"-avx2 bash -c "
+    set -e
+    pip list"
+
+  docker exec cpu-test-"$NUMA_NODE" bash -c "
+    set -e
+    pip list"
+
+  # offline inference
+  docker exec cpu-test-"$NUMA_NODE"-avx2 bash -c "
+    set -e
+    python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m"
+
+  # Run kernel tests
+  docker exec cpu-test-"$NUMA_NODE" bash -c "
+    set -e
+    pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
+    pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py
+    pytest -x -v -s tests/kernels/test_onednn.py"
+
+  # Run basic model test
+  docker exec cpu-test-"$NUMA_NODE" bash -c "
+    set -e
+    # Note: disable until supports V1
+    # pytest -x -v -s tests/kernels/attention/test_cache.py -m cpu_model
+    # pytest -x -v -s tests/kernels/attention/test_mla_decode_cpu.py -m cpu_model
+
+    pytest -x -v -s tests/models/language/generation -m cpu_model
+    VLLM_CPU_SGL_KERNEL=1 pytest -x -v -s tests/models/language/generation -m cpu_model
+
+    pytest -x -v -s tests/models/language/pooling -m cpu_model
+    pytest -x -v -s tests/models/multimodal/generation \
+                --ignore=tests/models/multimodal/generation/test_pixtral.py \
+                -m cpu_model"
+
+  # Run compressed-tensor test
+  docker exec cpu-test-"$NUMA_NODE" bash -c "
+    set -e
+    pytest -x -s -v \
+    tests/quantization/test_compressed_tensors.py::test_compressed_tensors_w8a8_logprobs"
+
+  # Run AWQ/GPTQ test
+  docker exec cpu-test-"$NUMA_NODE" bash -c "
+    set -e
+    pytest -x -s -v \
+    tests/quantization/test_cpu_wna16.py"
+
+  # Run multi-lora tests
+  docker exec cpu-test-"$NUMA_NODE" bash -c "
+    set -e
+    pytest -x -s -v \
+    tests/lora/test_qwenvl.py"
+
+  # online serving: tp+pp
+  docker exec cpu-test-"$NUMA_NODE" bash -c '
+    set -e
+    VLLM_CPU_OMP_THREADS_BIND=$E2E_OMP_THREADS VLLM_CPU_SGL_KERNEL=1 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; 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 \
+      --endpoint /v1/completions
+    kill -s SIGTERM $server_pid &'
+
+  # online serving: tp+dp
+  docker exec cpu-test-"$NUMA_NODE" bash -c '
+    set -e
+    VLLM_CPU_OMP_THREADS_BIND=$E2E_OMP_THREADS VLLM_CPU_SGL_KERNEL=1 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; 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 \
+      --endpoint /v1/completions
+    kill -s SIGTERM $server_pid &'
+}
+
+# All of CPU tests are expected to be finished less than 40 mins.
+export -f cpu_tests
+timeout 2.5h bash -c "cpu_tests $CORE_RANGE $NUMA_NODE"

.buildkite/scripts/hardware_ci/run-hpu-test.sh

@@ -5,9 +5,7 @@
 set -exuo pipefail
 
 # Try building the docker image
-image_name="hpu/upstream-vllm-ci:${BUILDKITE_COMMIT}"
-container_name="hpu-upstream-vllm-ci-${BUILDKITE_COMMIT}-container"
-cat <<EOF | docker build -t ${image_name} -f - .
+cat <<EOF | docker build -t hpu-plugin-v1-test-env -f - .
 FROM gaudi-base-image:latest
 
 COPY ./ /workspace/vllm
@@ -17,8 +15,7 @@ WORKDIR /workspace/vllm
 ENV no_proxy=localhost,127.0.0.1
 ENV PT_HPU_ENABLE_LAZY_COLLECTIVES=true
 
-RUN bash -c 'pip install -r <(sed "/^torch/d" requirements/build.txt)'
-RUN VLLM_TARGET_DEVICE=empty pip install --no-build-isolation -e .
+RUN VLLM_TARGET_DEVICE=empty pip install .
 RUN pip install git+https://github.com/vllm-project/vllm-gaudi.git
 
 # install development dependencies (for testing)
@@ -39,20 +36,15 @@ EOF
 # functions, while other platforms only need one remove_docker_container
 # function.
 EXITCODE=1
-remove_docker_containers() { docker rm -f ${container_name} || true; }
+remove_docker_containers() { docker rm -f hpu-plugin-v1-test || true; }
 trap 'remove_docker_containers; exit $EXITCODE;' EXIT
 remove_docker_containers
 
 echo "Running HPU plugin v1 test"
-docker run --rm --runtime=habana --name=${container_name} --network=host \
+docker run --rm --runtime=habana --name=hpu-plugin-v1-test --network=host \
   -e HABANA_VISIBLE_DEVICES=all \
-  -e VLLM_SKIP_WARMUP=true \
-  -e PT_HPU_ENABLE_LAZY_COLLECTIVES=true \
-  -e PT_HPU_LAZY_MODE=1 \
-  "${image_name}" \
-  /bin/bash -c '
-  cd vllm; timeout 120s python -u examples/offline_inference/basic/generate.py --model facebook/opt-125m
-'
+  hpu-plugin-v1-test-env \
+  /bin/bash "/workspace/vllm-gaudi/tests/upstream_tests/ci_tests.sh"
 
 EXITCODE=$?
 if [ $EXITCODE -eq 0 ]; then

.buildkite/scripts/hardware_ci/run-xpu-test.sh

@@ -38,18 +38,15 @@ docker run \
     python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 -O3 -cc.cudagraph_mode=NONE
     python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend ray
     python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend mp
+    python3 examples/offline_inference/basic/generate.py --model Intel/Qwen2.5-0.5B-W4A16-G128-AutoRound-LLMC-TEST-ONLY --enforce-eager
     python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --attention-backend=TRITON_ATTN
-    python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --quantization fp8
-    python3 examples/offline_inference/basic/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4  --block-size 64 --enforce-eager
-    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
+    pytest -v -s v1/core
     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
     pytest -v -s v1/structured_output
-    pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_tree_attention.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py
+    pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_tree_attention.py --ignore=v1/spec_decode/test_speculators_eagle3.py
     pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_nixl_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py
     pytest -v -s v1/test_serial_utils.py
 '

.buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh

@@ -43,6 +43,7 @@ trap cleanup EXIT
 
 for BACK in "${BACKENDS[@]}"; do
   VLLM_DEEP_GEMM_WARMUP=skip \
+  VLLM_ALL2ALL_BACKEND=$BACK \
   vllm serve "$MODEL" \
     --enforce-eager \
     --tensor-parallel-size 2 \
@@ -51,7 +52,6 @@ for BACK in "${BACKENDS[@]}"; do
     --enable-eplb \
     --trust-remote-code \
     --max-model-len 2048 \
-    --all2all-backend $BACK \
     --port $PORT &
   SERVER_PID=$!
   wait_for_server $PORT

.buildkite/scripts/upload-release-wheels.sh

@@ -7,19 +7,17 @@ SUBPATH=$BUILDKITE_COMMIT
 S3_COMMIT_PREFIX="s3://$BUCKET/$SUBPATH/"
 
 RELEASE_VERSION=$(buildkite-agent meta-data get release-version)
-GIT_VERSION=$(git describe --exact-match --tags $BUILDKITE_COMMIT 2>/dev/null)
-
 echo "Release version from Buildkite: $RELEASE_VERSION"
-
-if [[ -z "$GIT_VERSION" ]]; then
+GIT_VERSION=$(git describe --exact-match --tags $BUILDKITE_COMMIT 2>/dev/null)
+if [ -z "$GIT_VERSION" ]; then
     echo "[FATAL] Not on a git tag, cannot create release."
     exit 1
 else
     echo "Git version for commit $BUILDKITE_COMMIT: $GIT_VERSION"
 fi
 # sanity check for version mismatch
-if [[ "$RELEASE_VERSION" != "$GIT_VERSION" ]]; then
-  if [[ "$FORCE_RELEASE_IGNORE_VERSION_MISMATCH" == "true" ]]; then
+if [ "$RELEASE_VERSION" != "$GIT_VERSION" ]; then
+  if [ "$FORCE_RELEASE_IGNORE_VERSION_MISMATCH" == "true" ]; then
     echo "[WARNING] Force release and ignore version mismatch"
   else
     echo "[FATAL] Release version from Buildkite does not match Git version."
@@ -29,7 +27,7 @@ fi
 PURE_VERSION=${RELEASE_VERSION#v} # remove leading 'v'
 
 # check pypi token
-if [[ -z "$PYPI_TOKEN" ]]; then
+if [ -z "$PYPI_TOKEN" ]; then
   echo "[FATAL] PYPI_TOKEN is not set."
   exit 1
 else
@@ -37,8 +35,41 @@ else
   export TWINE_PASSWORD="$PYPI_TOKEN"
 fi
 
+# check github token
+if [ -z "$GITHUB_TOKEN" ]; then
+  echo "[FATAL] GITHUB_TOKEN is not set."
+  exit 1
+else
+  export GH_TOKEN="$GITHUB_TOKEN"
+fi
+
 set -x # avoid printing secrets above
 
+# download gh CLI from github
+# Get latest gh CLI version from GitHub API
+GH_VERSION=$(curl -s https://api.github.com/repos/cli/cli/releases/latest | grep '"tag_name":' | sed -E 's/.*"([^"]+)".*/\1/' | sed 's/^v//')
+if [ -z "$GH_VERSION" ]; then
+  echo "[FATAL] Failed to get latest gh CLI version from GitHub"
+  exit 1
+fi
+echo "Downloading gh CLI version: $GH_VERSION"
+GH_TARBALL="gh_${GH_VERSION}_linux_amd64.tar.gz"
+GH_URL="https://github.com/cli/cli/releases/download/v${GH_VERSION}/${GH_TARBALL}"
+GH_INSTALL_DIR="/tmp/gh-install"
+mkdir -p "$GH_INSTALL_DIR"
+pushd "$GH_INSTALL_DIR"
+curl -L -o "$GH_TARBALL" "$GH_URL"
+tar -xzf "$GH_TARBALL"
+GH_BIN=$(realpath $(find . -name "gh" -type f -executable | head -n 1))
+if [ -z "$GH_BIN" ]; then
+  echo "[FATAL] Failed to find gh CLI executable"
+  exit 1
+fi
+echo "gh CLI downloaded successfully, version: $($GH_BIN --version)"
+echo "Last 5 releases on GitHub:" # as a sanity check of gh and GH_TOKEN
+command "$GH_BIN" release list --limit 5
+popd
+
 # install twine from pypi
 python3 -m venv /tmp/vllm-release-env
 source /tmp/vllm-release-env/bin/activate
@@ -58,13 +89,16 @@ echo "Wheels copied to local directory"
 git archive --format=tar.gz --output="$DIST_DIR/vllm-${PURE_VERSION}.tar.gz" $BUILDKITE_COMMIT
 ls -la $DIST_DIR
 
+
 # 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
+if [ -z "$PYPI_WHEEL_FILES" ]; then
   echo "No default variant wheels found, quitting..."
   exit 1
 fi
-
 python3 -m twine check $PYPI_WHEEL_FILES
-python3 -m twine upload --non-interactive --verbose $PYPI_WHEEL_FILES
+python3 -m twine --non-interactive --verbose upload $PYPI_WHEEL_FILES
 echo "Wheels uploaded to PyPI"
+
+# create release on GitHub with the release version and all wheels
+command "$GH_BIN" release create $GIT_VERSION -d --latest --notes-from-tag --verify-tag $DIST_DIR/*.whl

.buildkite/test-amd.yaml

@@ -70,7 +70,6 @@ steps:
   - vllm/
   - tests/test_inputs.py
   - tests/test_outputs.py
-  - tests/test_pooling_params.py
   - tests/multimodal
   - tests/renderers
   - tests/standalone_tests/lazy_imports.py
@@ -83,7 +82,6 @@ steps:
   - python3 standalone_tests/lazy_imports.py
   - pytest -v -s test_inputs.py
   - pytest -v -s test_outputs.py
-  - pytest -v -s test_pooling_params.py
   - pytest -v -s -m 'cpu_test' multimodal
   - pytest -v -s renderers
   - pytest -v -s tokenizers_
@@ -132,7 +130,7 @@ steps:
   - tests/entrypoints/
   commands:
   - pytest -v -s entrypoints/openai/tool_parsers
-  - pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/openai --ignore=entrypoints/rpc --ignore=entrypoints/instrumentator --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py  --ignore=entrypoints/pooling
+  - pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/openai --ignore=entrypoints/rpc --ignore=entrypoints/sleep --ignore=entrypoints/instrumentator --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py  --ignore=entrypoints/pooling
 
 - label: Entrypoints Integration Test (LLM) # 30min
   timeout_in_minutes: 40
@@ -179,14 +177,14 @@ steps:
   torch_nightly: true
   source_file_dependencies:
   - vllm/
+  - tests/entrypoints/sleep
   - tests/entrypoints/rpc
-  - tests/entrypoints/instrumentator
   - tests/tool_use
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/instrumentator
-  - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc
+  - pytest -v -s entrypoints/sleep
   - pytest -v -s tool_use
+  - PYTHONPATH=/vllm-workspace  pytest -v -s entrypoints/rpc
 
 - label: Entrypoints Integration Test (Pooling)
   timeout_in_minutes: 50
@@ -233,7 +231,6 @@ steps:
   - tests/compile/fullgraph/test_basic_correctness.py
   - examples/offline_inference/rlhf.py
   - examples/offline_inference/rlhf_colocate.py
-  - examples/offline_inference/new_weight_syncing/
   - tests/examples/offline_inference/data_parallel.py
   - tests/v1/distributed
   - tests/v1/engine/test_engine_core_client.py
@@ -269,16 +266,10 @@ steps:
   - pytest -v -s distributed/test_symm_mem_allreduce.py
   # TODO: create a dedicated test section for multi-GPU example tests
   # when we have multiple distributed example tests
-  # OLD rlhf examples
   - pushd ../examples/offline_inference
   - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf.py
   - VLLM_ALLOW_INSECURE_SERIALIZATION=1 RAY_DEDUP_LOGS=0 python3 rlhf_colocate.py
   - popd
-  # NEW rlhf examples
-  - pushd ../examples/offline_inference/new_weight_syncing
-  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf.py
-  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_async_new_apis.py
-  - popd
 
 - label: Distributed Tests (8 GPUs) # 4min
   timeout_in_minutes: 10
@@ -514,7 +505,7 @@ steps:
     - python3 offline_inference/vision_language_multi_image.py --seed 0
     - python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
     # for pooling models
-    - python3 pooling/embed/vision_embedding_offline.py --seed 0
+    - python3 pooling/pooling/vision_language_pooling.py --seed 0
     # for features demo
     - python3 offline_inference/prefix_caching.py
     - python3 offline_inference/llm_engine_example.py
@@ -534,7 +525,6 @@ steps:
   - tests/cuda
   commands:
     - pytest -v -s cuda/test_cuda_context.py
-    - pytest -v -s cuda/test_platform_no_cuda_init.py
 
 - label: Samplers Test # 56min
   timeout_in_minutes: 75
@@ -552,7 +542,7 @@ steps:
 - label: LoRA Test %N # 20min each
   timeout_in_minutes: 30
   mirror_hardwares: [amdexperimental]
-  agent_pool: mi325_1
+  agent_pool: mi325_8
   # grade: Blocking
   source_file_dependencies:
   - vllm/lora
@@ -614,11 +604,9 @@ steps:
   - tests/compile
   commands:
   - pytest -v -s compile/fullgraph/test_full_graph.py -k 'not test_fp8_kv_scale_compile'
-    # # Limit to no custom ops to reduce running time
-    # # Wrap with quotes to escape yaml and avoid starting -k string with a -
-    # - "pytest -v -s compile/distributed/test_fusions_e2e.py -k 'TRITON and not +quant_fp8 and not Llama-4'"
-    # Old E2E tests were removed in https://github.com/vllm-project/vllm/pull/33293
-    # in favor of new tests in fusions_e2e. We avoid replicating the new jobs in this file as it's deprecated.
+    # Limit to no custom ops to reduce running time
+    # Wrap with quotes to escape yaml and avoid starting -k string with a -
+  - "pytest -v -s compile/distributed/test_fusions_e2e.py -k 'TRITON and not +quant_fp8 and not Llama-4'"
 
 - label: Cudagraph test
   timeout_in_minutes: 20
@@ -648,13 +636,12 @@ steps:
 - label: Kernels Attention Test %N # 23min
   timeout_in_minutes: 35
   mirror_hardwares: [amdexperimental, amdproduction]
-  agent_pool: mi325_1
+  agent_pool: mi325_8
   # grade: Blocking
   source_file_dependencies:
   - csrc/attention/
+  - vllm/attention
   - vllm/v1/attention
-    # TODO: remove this dependency (https://github.com/vllm-project/vllm/issues/32267)
-  - vllm/model_executor/layers/attention
   - tests/kernels/attention
   commands:
     - pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
@@ -663,7 +650,7 @@ steps:
 - label: Kernels Quantization Test %N # 64min
   timeout_in_minutes: 90
   mirror_hardwares: [amdexperimental]
-  agent_pool: mi325_1
+  agent_pool: mi325_8
   # grade: Blocking
   source_file_dependencies:
   - csrc/quantization/
@@ -676,7 +663,7 @@ steps:
 - label: Kernels MoE Test %N # 40min
   timeout_in_minutes: 60
   mirror_hardwares: [amdexperimental, amdproduction]
-  agent_pool: mi325_1
+  agent_pool: mi325_8
   # grade: Blocking
   source_file_dependencies:
   - csrc/quantization/cutlass_w8a8/moe/
@@ -754,7 +741,7 @@ steps:
 - label: Benchmarks # 11min
   timeout_in_minutes: 20
   mirror_hardwares: [amdexperimental, amdproduction]
-  agent_pool: mi325_1
+  agent_pool: mi325_8
   # grade: Blocking
   working_dir: "/vllm-workspace/.buildkite"
   source_file_dependencies:
@@ -765,7 +752,7 @@ steps:
 - label: Benchmarks CLI Test # 7min
   timeout_in_minutes: 20
   mirror_hardwares: [amdexperimental, amdproduction]
-  agent_pool: mi325_1
+  agent_pool: mi325_8
   # grade: Blocking
   source_file_dependencies:
   - vllm/
@@ -839,7 +826,7 @@ steps:
 - label: Basic Models Tests (Extra Initialization) %N
   timeout_in_minutes: 45
   mirror_hardwares: [amdexperimental, amdproduction]
-  agent_pool: mi325_1
+  agent_pool: mi325_8
   # grade: Blocking
   torch_nightly: true
   source_file_dependencies:
@@ -864,11 +851,10 @@ steps:
   torch_nightly: true
   source_file_dependencies:
   - vllm/
-  - tests/models/test_terratorch.py
   - tests/models/test_transformers.py
   - tests/models/test_registry.py
   commands:
-    - pytest -v -s models/test_terratorch.py models/test_transformers.py models/test_registry.py
+    - pytest -v -s models/test_transformers.py models/test_registry.py
 
 - label: Basic Models Test (Other CPU) # 5min
   mirror_hardwares: [amdexperimental, amdproduction]
@@ -901,7 +887,7 @@ steps:
 - label: Language Models Tests (Extra Standard) %N
   timeout_in_minutes: 45
   mirror_hardwares: [amdexperimental]
-  agent_pool: mi325_1
+  agent_pool: mi325_8
   # grade: Blocking
   torch_nightly: true
   source_file_dependencies:
@@ -922,7 +908,7 @@ steps:
 - label: Language Models Tests (Hybrid) %N
   timeout_in_minutes: 75
   mirror_hardwares: [amdexperimental]
-  agent_pool: mi325_1
+  agent_pool: mi325_8
   # grade: Blocking
   torch_nightly: true
   source_file_dependencies:
@@ -1145,7 +1131,7 @@ steps:
   - csrc/quantization/cutlass_w8a8/moe/
   - vllm/model_executor/layers/fused_moe/cutlass_moe.py
   - vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py
-  - vllm/model_executor/layers/fused_moe/flashinfer_a2a_prepare_finalize.py
+  - vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py
   - vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
   - vllm/v1/attention/backends/flashinfer.py
   - vllm/v1/attention/backends/mla/cutlass_mla.py
@@ -1191,26 +1177,44 @@ steps:
   - vllm/model_executor/layers/layernorm.py
   - vllm/model_executor/layers/activation.py
   - vllm/model_executor/layers/quantization/input_quant_fp8.py
-  - tests/compile/passes/test_fusion_attn.py
-  - tests/compile/passes/test_silu_mul_quant_fusion.py
-  - tests/compile/passes/distributed/test_fusion_all_reduce.py
+  - tests/compile/test_fusion_attn.py
+  - tests/compile/test_silu_mul_quant_fusion.py
+  - tests/compile/distributed/test_fusion_all_reduce.py
+  - tests/compile/distributed/test_fusions_e2e.py
   - tests/compile/fullgraph/test_full_graph.py
   commands:
     - nvidia-smi
-    - pytest -v -s tests/compile/passes/test_fusion_attn.py
-    - pytest -v -s tests/compile/passes/test_silu_mul_quant_fusion.py
+    - pytest -v -s tests/compile/test_fusion_attn.py
+    - pytest -v -s tests/compile/test_silu_mul_quant_fusion.py
     # this runner has 2 GPUs available even though num_gpus=2 is not set
-    - pytest -v -s tests/compile/passes/distributed/test_fusion_all_reduce.py
-
-    # # Limit to Inductor partition, no custom ops, and allreduce & attn fusion to reduce running time
-    # # Wrap with quotes to escape yaml
-    # - "pytest -v -s tests/compile/distributed/test_fusions_e2e.py::test_tp2_attn_quant_allreduce_rmsnorm -k 'True and not +quant_fp8 and not +rms_norm'"
-    # Old E2E tests were removed in https://github.com/vllm-project/vllm/pull/33293
-    # in favor of new tests in fusions_e2e. We avoid replicating the new jobs in this file as it's deprecated.
-
+    - pytest -v -s tests/compile/distributed/test_fusion_all_reduce.py
+    # Limit to Inductor partition, no custom ops, and allreduce & attn fusion to reduce running time
+    # Wrap with quotes to escape yaml
+    - "pytest -v -s tests/compile/distributed/test_fusions_e2e.py::test_tp2_attn_quant_allreduce_rmsnorm -k 'True and not +quant_fp8 and not +rms_norm'"
     # test_fp8_kv_scale_compile requires FlashAttention (not supported on default L4/L40)
     - pytest -v -s tests/compile/fullgraph/test_full_graph.py::test_fp8_kv_scale_compile
 
+- label: Blackwell Fusion E2E Tests # 30 min
+  timeout_in_minutes: 40
+  working_dir: "/vllm-workspace/"
+  gpu: b200
+  optional: true
+  num_gpus: 2
+  source_file_dependencies:
+  - csrc/quantization/fp4/
+  - vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
+  - vllm/v1/attention/backends/flashinfer.py
+  - vllm/compilation/
+  # can affect pattern matching
+  - vllm/model_executor/layers/layernorm.py
+  - vllm/model_executor/layers/activation.py
+  - vllm/model_executor/layers/quantization/input_quant_fp8.py
+  - tests/compile/distributed/test_fusions_e2e.py
+  commands:
+    - nvidia-smi
+    # Run all e2e fusion tests
+    - pytest -v -s tests/compile/distributed/test_fusions_e2e.py
+
 - label: Blackwell GPT-OSS Eval
   timeout_in_minutes: 60
   working_dir: "/vllm-workspace/"
@@ -1273,7 +1277,7 @@ steps:
 
 - label: 2 Node Tests (4 GPUs in total) # 16min
   timeout_in_minutes: 30
-  mirror_hardwares: [amdexperimental, amdmultinode]
+  mirror_hardwares: [amdexperimental]
   agent_pool: mi325_4
   # grade: Blocking
   working_dir: "/vllm-workspace/tests"
@@ -1287,15 +1291,15 @@ steps:
   - tests/distributed/
   - tests/examples/offline_inference/data_parallel.py
   commands:
-  - # the following commands are for the first node, with ip 192.168.10.10 (ray environment already set up)  | grep 'Same node test passed'   | grep 'Node count test passed'
-    - VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py 
-    - NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py 
+  - # the following commands are for the first node, with ip 192.168.10.10 (ray environment already set up)
+    - VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed'
+    - NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed'
     - python3 ../examples/offline_inference/data_parallel.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=0 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code
     - VLLM_MULTI_NODE=1 pytest -v -s distributed/test_multi_node_assignment.py
     - VLLM_MULTI_NODE=1 pytest -v -s distributed/test_pipeline_parallel.py
   - # the following commands are for the second node, with ip 192.168.10.11 (ray environment already set up)
-    - VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py 
-    - NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py 
+    - VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed'
+    - NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed'
     - python3 ../examples/offline_inference/data_parallel.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=1 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code
 
 - label: Distributed Tests (2 GPUs) # 68min
@@ -1334,7 +1338,7 @@ steps:
   - pytest -v -s ./compile/test_wrapper.py
   - VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
   - VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
-  - pytest -v -s compile/correctness_e2e/test_sequence_parallel.py
+  - pytest -v -s distributed/test_sequence_parallel.py
   - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
   - pytest -v -s v1/worker/test_worker_memory_snapshot.py
 
@@ -1504,9 +1508,6 @@ steps:
   source_file_dependencies:
   - vllm/
   commands:
-  # Work around HIP bug tracked here: https://github.com/ROCm/hip/issues/3876
-  # TODO: Remove when the bug is fixed in a future ROCm release
-  - export TORCH_NCCL_BLOCKING_WAIT=1
   # NOTE: don't test llama model here, it seems hf implementation is buggy
   # see https://github.com/vllm-project/vllm/pull/5689 for details
   - pytest -v -s distributed/test_custom_all_reduce.py
@@ -1557,15 +1558,12 @@ steps:
   working_dir: "/vllm-workspace/"
   num_gpus: 2
   commands:
-    - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/passes/distributed/test_async_tp.py
-    - pytest -v -s tests/compile/passes/distributed/test_sequence_parallelism.py
-    - pytest -v -s tests/compile/passes/distributed/test_fusion_all_reduce.py
+    - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/distributed/test_async_tp.py
+    - pytest -v -s tests/compile/distributed/test_sequence_parallelism.py
+    - pytest -v -s tests/compile/distributed/test_fusion_all_reduce.py
     #- pytest -v -s tests/compile/distributed/test_fusions_e2e.py::test_tp2_attn_quant_allreduce_rmsnorm
-    # - "VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/distributed/test_fusions_e2e.py -k 'not Llama-4'"
-    # Old E2E tests were removed in https://github.com/vllm-project/vllm/pull/33293
-    # in favor of new tests in fusions_e2e. We avoid replicating the new jobs in this file as it's deprecated.
-
-    - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py
+    - "VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/distributed/test_fusions_e2e.py -k 'not Llama-4'"
+    - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/distributed/test_sequence_parallel.py
     - pytest -v -s tests/distributed/test_context_parallel.py
     - HIP_VISIBLE_DEVICES=0,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=allgather_reducescatter --disable-nccl-for-dp-synchronization
     - pytest -v -s tests/v1/distributed/test_dbo.py

.buildkite/test-pipeline.yaml

@@ -63,7 +63,6 @@ steps:
   - vllm/
   - tests/test_inputs.py
   - tests/test_outputs.py
-  - tests/test_pooling_params.py
   - tests/multimodal
   - tests/renderers
   - tests/standalone_tests/lazy_imports.py
@@ -76,7 +75,6 @@ steps:
   - python3 standalone_tests/lazy_imports.py
   - pytest -v -s test_inputs.py
   - pytest -v -s test_outputs.py
-  - pytest -v -s test_pooling_params.py
   - pytest -v -s -m 'cpu_test' multimodal
   - pytest -v -s renderers
   - pytest -v -s tokenizers_
@@ -118,7 +116,7 @@ steps:
   - tests/entrypoints/
   commands:
   - pytest -v -s entrypoints/openai/tool_parsers
-  - pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/rpc --ignore=entrypoints/instrumentator --ignore=entrypoints/openai --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py  --ignore=entrypoints/pooling
+  - pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/rpc --ignore=entrypoints/sleep --ignore=entrypoints/instrumentator --ignore=entrypoints/openai --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py  --ignore=entrypoints/pooling
 
 - label: Entrypoints Integration Test (LLM) # 30min
   timeout_in_minutes: 40
@@ -148,7 +146,7 @@ steps:
   - tests/entrypoints/test_chat_utils
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/test_chat_with_tool_reasoning.py --ignore=entrypoints/instrumentator --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/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
 
 - label: Entrypoints Integration Test (API Server 2)
@@ -159,13 +157,13 @@ steps:
   torch_nightly: true
   source_file_dependencies:
   - vllm/
+  - tests/entrypoints/sleep
   - tests/entrypoints/rpc
-  - tests/entrypoints/instrumentator
   - tests/tool_use
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/instrumentator
-  - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc
+  - pytest -v -s entrypoints/sleep
+  - PYTHONPATH=/vllm-workspace  pytest -v -s entrypoints/rpc
   - pytest -v -s tool_use
 
 - label: Entrypoints Integration Test (Pooling)
@@ -206,7 +204,6 @@ steps:
   - tests/compile/fullgraph/test_basic_correctness.py
   - examples/offline_inference/rlhf.py
   - examples/offline_inference/rlhf_colocate.py
-  - examples/offline_inference/new_weight_syncing/
   - tests/examples/offline_inference/data_parallel.py
   - tests/v1/distributed
   - tests/v1/engine/test_engine_core_client.py
@@ -241,16 +238,10 @@ steps:
   - pytest -v -s distributed/test_symm_mem_allreduce.py
   # TODO: create a dedicated test section for multi-GPU example tests
   # when we have multiple distributed example tests
-  # OLD rlhf examples
   - pushd ../examples/offline_inference
   - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf.py
   - VLLM_ALLOW_INSECURE_SERIALIZATION=1 RAY_DEDUP_LOGS=0 python3 rlhf_colocate.py
   - popd
-  # NEW rlhf examples
-  - pushd ../examples/offline_inference/new_weight_syncing
-  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf.py
-  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_async_new_apis.py
-  - popd
 
 - label: Distributed Tests (8 GPUs) # 4min
   timeout_in_minutes: 10
@@ -371,7 +362,7 @@ steps:
     - pytest -v -s v1/sample
     - pytest -v -s v1/logits_processors
     - pytest -v -s v1/worker
-    - pytest -v -s -m 'not slow_test' v1/spec_decode
+    - pytest -v -s v1/spec_decode
     - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
     - pytest -v -s -m 'not cpu_test' v1/metrics
     - pytest -v -s v1/test_oracle.py
@@ -453,7 +444,7 @@ steps:
     - python3 offline_inference/vision_language_multi_image.py --seed 0
     - python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
     # for pooling models
-    - python3 pooling/embed/vision_embedding_offline.py --seed 0
+    - python3 pooling/pooling/vision_language_pooling.py --seed 0
     # for features demo
     - python3 offline_inference/prefix_caching.py
     - python3 offline_inference/llm_engine_example.py
@@ -519,7 +510,6 @@ steps:
   # However, find does not normally propagate error codes, so we combine it with xargs
   # (using -0 for proper path handling)
   - "find compile/ -maxdepth 1 -name 'test_*.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
-  - pytest -s -v compile/passes --ignore compile/passes/distributed
 
 - label: PyTorch Fullgraph Smoke Test # 15min
   timeout_in_minutes: 30
@@ -547,11 +537,9 @@ steps:
   commands:
     # fp8 kv scales not supported on sm89, tested on Blackwell instead
   - pytest -v -s compile/fullgraph/test_full_graph.py -k 'not test_fp8_kv_scale_compile'
-    # # Limit to no custom ops to reduce running time
-    # # Wrap with quotes to escape yaml and avoid starting -k string with a -
-    # - "pytest -v -s compile/distributed/test_fusions_e2e.py -k 'TRITON and not +quant_fp8 and not Llama-4'"
-    # Old E2E tests were removed in https://github.com/vllm-project/vllm/pull/33293
-    # in favor of new tests in fusions_e2e. We avoid replicating the new jobs in this file as it's deprecated.
+    # Limit to no custom ops to reduce running time
+    # Wrap with quotes to escape yaml and avoid starting -k string with a -
+  - "pytest -v -s compile/distributed/test_fusions_e2e.py -k 'TRITON and not +quant_fp8 and not Llama-4'"
 
 - label: Cudagraph test
   timeout_in_minutes: 20
@@ -580,9 +568,8 @@ steps:
   mirror_hardwares: [amdexperimental]
   source_file_dependencies:
   - csrc/attention/
+  - vllm/attention
   - vllm/v1/attention
-    # TODO: remove this dependency (https://github.com/vllm-project/vllm/issues/32267)
-  - vllm/model_executor/layers/attention
   - tests/kernels/attention
   commands:
     - pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
@@ -805,11 +792,10 @@ steps:
   torch_nightly: true
   source_file_dependencies:
   - vllm/
-  - tests/models/test_terratorch.py
   - tests/models/test_transformers.py
   - tests/models/test_registry.py
   commands:
-    - pytest -v -s models/test_terratorch.py models/test_transformers.py models/test_registry.py
+    - pytest -v -s models/test_transformers.py models/test_registry.py
 
 - label: Basic Models Test (Other CPU) # 5min
   timeout_in_minutes: 10
@@ -862,7 +848,7 @@ steps:
   commands:
     # Install fast path packages for testing against transformers
     # Note: also needed to run plamo2 model in vLLM
-    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
+    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.2.5'
     - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.2'
     # Shard hybrid language model tests
     - pytest -v -s models/language/generation \
@@ -881,7 +867,7 @@ steps:
   commands:
     # Install fast path packages for testing against transformers
     # Note: also needed to run plamo2 model in vLLM
-    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
+    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.2.5'
     - 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)'
 
@@ -1031,7 +1017,7 @@ steps:
   - csrc/quantization/cutlass_w8a8/moe/
   - vllm/model_executor/layers/fused_moe/cutlass_moe.py
   - vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py
-  - vllm/model_executor/layers/fused_moe/flashinfer_a2a_prepare_finalize.py
+  - vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py
   - vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
   - vllm/v1/attention/backends/flashinfer.py
   - vllm/v1/attention/backends/mla/cutlass_mla.py
@@ -1081,23 +1067,84 @@ steps:
   - vllm/model_executor/layers/quantization/input_quant_fp8.py
   - tests/compile/test_fusion_attn.py
   - tests/compile/test_silu_mul_quant_fusion.py
-  - tests/compile/passes/distributed/test_fusion_all_reduce.py
+  - tests/compile/distributed/test_fusion_all_reduce.py
+  - tests/compile/distributed/test_fusions_e2e.py
   - tests/compile/fullgraph/test_full_graph.py
   commands:
     - nvidia-smi
     - pytest -v -s tests/compile/test_fusion_attn.py
     - pytest -v -s tests/compile/test_silu_mul_quant_fusion.py
     # this runner has 2 GPUs available even though num_gpus=2 is not set
-    - pytest -v -s tests/compile/passes/distributed/test_fusion_all_reduce.py
-    #  # Limit to Inductor partition, no custom ops, and allreduce & attn fusion to reduce running time
-    #  # Wrap with quotes to escape yaml
-    #  - "pytest -v -s tests/compile/distributed/test_fusions_e2e.py::test_tp2_attn_quant_allreduce_rmsnorm -k 'True and not +quant_fp8 and not +rms_norm'"
-    # Old E2E tests were removed in https://github.com/vllm-project/vllm/pull/33293
-    # in favor of new tests in fusions_e2e. We avoid replicating the new jobs in this file as it's deprecated.
-
+    - pytest -v -s tests/compile/distributed/test_fusion_all_reduce.py
+    # Limit to Inductor partition, no custom ops, and allreduce & attn fusion to reduce running time
+    # Wrap with quotes to escape yaml
+    - "pytest -v -s tests/compile/distributed/test_fusions_e2e.py::test_tp2_attn_quant_allreduce_rmsnorm -k 'True and not +quant_fp8 and not +rms_norm'"
     # test_fp8_kv_scale_compile requires FlashAttention (not supported on default L4/L40)
     - pytest -v -s tests/compile/fullgraph/test_full_graph.py::test_fp8_kv_scale_compile
 
+- label: Blackwell Fusion E2E Tests # 30 min
+  timeout_in_minutes: 40
+  working_dir: "/vllm-workspace/"
+  gpu: b200
+  optional: true
+  num_gpus: 2
+  source_file_dependencies:
+  - csrc/quantization/fp4/
+  - vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
+  - vllm/v1/attention/backends/flashinfer.py
+  - vllm/compilation/
+  # can affect pattern matching
+  - vllm/model_executor/layers/layernorm.py
+  - vllm/model_executor/layers/activation.py
+  - vllm/model_executor/layers/quantization/input_quant_fp8.py
+  - tests/compile/distributed/test_fusions_e2e.py
+  commands:
+    - nvidia-smi
+    # Run all e2e fusion tests
+    - pytest -v -s tests/compile/distributed/test_fusions_e2e.py
+
+- label: Hopper Fusion E2E Tests (H100) # 10min
+  timeout_in_minutes: 70
+  working_dir: "/vllm-workspace/"
+  gpu: h100
+  optional: true
+  source_file_dependencies:
+  - csrc/quantization/fp4/
+  - vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
+  - vllm/v1/attention/backends/flashinfer.py
+  - vllm/compilation/
+  # can affect pattern matching
+  - vllm/model_executor/layers/layernorm.py
+  - vllm/model_executor/layers/activation.py
+  - vllm/model_executor/layers/quantization/input_quant_fp8.py
+  - tests/compile/test_fusion_attn.py
+  commands:
+    - export VLLM_TEST_CLEAN_GPU_MEMORY=1
+    # skip Llama-4 since it does not fit on this device
+    - pytest -v -s tests/compile/test_fusion_attn.py -k 'not Llama-4'
+
+- label: Hopper Fusion Distributed E2E Tests (2xH100)  # 70min
+  timeout_in_minutes: 70
+  working_dir: "/vllm-workspace/"
+  gpu: h100
+  optional: true
+  num_gpus: 2
+  source_file_dependencies:
+  - csrc/quantization/fp4/
+  - vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
+  - vllm/v1/attention/backends/flashinfer.py
+  - vllm/compilation/
+  # can affect pattern matching
+  - vllm/model_executor/layers/layernorm.py
+  - vllm/model_executor/layers/activation.py
+  - vllm/model_executor/layers/quantization/input_quant_fp8.py
+  - tests/compile/distributed/test_fusions_e2e.py
+  commands:
+    - export VLLM_TEST_CLEAN_GPU_MEMORY=1
+    # Run all e2e fusion tests
+    - pytest -v -s tests/compile/distributed/test_fusions_e2e.py -k 'not Llama-4'
+    - pytest -v -s tests/compile/distributed/test_fusion_all_reduce.py
+
 - label: Blackwell GPT-OSS Eval
   timeout_in_minutes: 60
   working_dir: "/vllm-workspace/"
@@ -1155,8 +1202,6 @@ steps:
   - pytest -v -s distributed/test_shm_broadcast.py
   - pytest -v -s distributed/test_shm_buffer.py
   - pytest -v -s distributed/test_shm_storage.py
-  - pytest -v -s distributed/test_packed_tensor.py
-  - pytest -v -s distributed/test_weight_transfer.py
 
 - label: 2 Node Tests (4 GPUs in total) # 16min
   timeout_in_minutes: 30
@@ -1271,7 +1316,7 @@ steps:
   - pytest -v -s distributed/test_distributed_oot.py
   - 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 lora resolver plugins
+  - pytest -v -s plugins/lora_resolvers # unit tests for in-tree lora resolver plugins
 
 - label: Pipeline + Context Parallelism Test # 45min
   timeout_in_minutes: 60
@@ -1374,20 +1419,6 @@ steps:
   - TARGET_TEST_SUITE=A100 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
   - pytest -v -s -x lora/test_mixtral.py
 
-- label: Acceptance Length Test (Large Models) # optional
-  timeout_in_minutes: 120
-  gpu: h100
-  optional: true
-  num_gpus: 1
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/v1/spec_decode/
-  - vllm/model_executor/models/mlp_speculator.py
-  - tests/v1/spec_decode/test_acceptance_length.py
-  commands:
-    - export VLLM_ALLOW_INSECURE_SERIALIZATION=1
-    - pytest -v -s v1/spec_decode/test_acceptance_length.py -m slow_test
-
 - label: LM Eval Large Models # optional
   gpu: a100
   optional: true
@@ -1422,8 +1453,8 @@ steps:
   commands:
     - export VLLM_TEST_CLEAN_GPU_MEMORY=1
     # Run sequence parallel tests
-    - pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py
-    - pytest -v -s tests/compile/passes/distributed/test_sequence_parallelism.py
+    - pytest -v -s tests/distributed/test_sequence_parallel.py
+    - pytest -v -s tests/compile/distributed/test_sequence_parallelism.py
 
 - label: Distributed Tests (H100) # optional
   gpu: h100
@@ -1431,7 +1462,7 @@ steps:
   working_dir: "/vllm-workspace/"
   num_gpus: 2
   commands:
-    - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/passes/distributed/test_async_tp.py
+    - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/distributed/test_async_tp.py
     - pytest -v -s tests/distributed/test_context_parallel.py
     - 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

.buildkite/test_areas/attention.yaml

@@ -4,7 +4,7 @@ depends_on:
 steps:
 - label: V1 attention (H100)
   timeout_in_minutes: 30
-  device: h100
+  gpu: h100
   source_file_dependencies:
     - vllm/config/attention.py
     - vllm/model_executor/layers/attention
@@ -15,7 +15,7 @@ steps:
 
 - label: V1 attention (B200)
   timeout_in_minutes: 30
-  device: b200
+  gpu: b200
   source_file_dependencies:
     - vllm/config/attention.py
     - vllm/model_executor/layers/attention

.buildkite/test_areas/basic_correctness.yaml

@@ -14,8 +14,3 @@ steps:
   - pytest -v -s basic_correctness/test_cumem.py
   - pytest -v -s basic_correctness/test_basic_correctness.py
   - pytest -v -s basic_correctness/test_cpu_offload.py
-  mirror:
-    amd:
-      device: mi325_1
-      depends_on:
-      - image-build-amd

.buildkite/test_areas/benchmarks.yaml

@@ -17,15 +17,3 @@ steps:
   - tests/benchmarks/
   commands:
   - pytest -v -s benchmarks/
-
-- label: Attention Benchmarks Smoke Test (B200)
-  device: b200
-  num_gpus: 2
-  optional: true
-  working_dir: "/vllm-workspace/"
-  timeout_in_minutes: 10
-  source_file_dependencies:
-  - benchmarks/attention_benchmarks/
-  - vllm/v1/attention/
-  commands:
-  - python3 benchmarks/attention_benchmarks/benchmark.py --backends flash flashinfer --batch-specs "8q1s1k" --repeats 1 --warmup-iters 1

.buildkite/test_areas/compile.yaml

@@ -2,200 +2,56 @@ group: Compile
 depends_on: 
   - image-build
 steps:
-- label: Sequence Parallel Correctness Tests (2 GPUs)
-  timeout_in_minutes: 50
+- label: Fusion and Compile Tests (B200)
+  timeout_in_minutes: 40
   working_dir: "/vllm-workspace/"
-  num_devices: 2
-  source_file_dependencies:
-  - vllm/model_executor/layers/
-  - vllm/compilation/
-  - vllm/v1/worker/
-  - vllm/v1/cudagraph_dispatcher.py
-  - tests/compile/correctness_e2e/test_sequence_parallel.py
-  commands:
-  - export VLLM_TEST_CLEAN_GPU_MEMORY=1
-  - pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py
-
-- label: Sequence Parallel Correctness Tests (2xH100)
-  timeout_in_minutes: 50
-  working_dir: "/vllm-workspace/"
-  device: h100
-  optional: true
-  num_devices: 2
-  commands:
-  - export VLLM_TEST_CLEAN_GPU_MEMORY=1
-  - pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py
-
-- label: AsyncTP Correctness Tests (2xH100)
-  timeout_in_minutes: 50
-  working_dir: "/vllm-workspace/"
-  device: h100
-  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/"
-  device: h100
-  num_devices: 2
-  source_file_dependencies:
-  - vllm/compilation/
-  - vllm/model_executor/layers
-  - tests/compile/passes/distributed/
-  commands:
-  - export VLLM_TEST_CLEAN_GPU_MEMORY=1
-  - pytest -s -v tests/compile/passes/distributed
-
-- label: Fusion and Compile Unit Tests (B200)
-  timeout_in_minutes: 20
-  working_dir: "/vllm-workspace/"
-  device: b200
+  gpu: b200
   source_file_dependencies:
   - csrc/quantization/fp4/
-  - vllm/model_executor/layers/quantization/
+  - vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
+  - vllm/v1/attention/backends/flashinfer.py
+  - vllm/v1/worker/
+  - vllm/v1/cudagraph_dispatcher.py
+  - vllm/compilation/
+  # can affect pattern matching
   - vllm/model_executor/layers/layernorm.py
   - vllm/model_executor/layers/activation.py
-  - vllm/model_executor/layers/attention/attention.py
-  - vllm/v1/attention/backends/flashinfer.py
-  - vllm/compilation/ # TODO(luka) limit to vllm/compilation/passes
-  - tests/compile/passes/test_fusion_attn.py
-  - tests/compile/passes/test_silu_mul_quant_fusion.py
-  - tests/compile/passes/distributed/test_fusion_all_reduce.py
+  - vllm/model_executor/layers/quantization/input_quant_fp8.py
+  - tests/compile/test_fusion_attn.py
+  - tests/compile/test_silu_mul_quant_fusion.py
+  - tests/compile/distributed/test_fusion_all_reduce.py
+  - tests/compile/distributed/test_fusions_e2e.py
   - tests/compile/fullgraph/test_full_graph.py
   commands:
-    # b200 runners are limited, so we limit the tests to the minimum set only supported on Blackwell
     - nvidia-smi
-    - pytest -v -s tests/compile/passes/test_fusion_attn.py -k FLASHINFER
-    - pytest -v -s tests/compile/passes/test_silu_mul_quant_fusion.py
-    # this runner has 2 GPUs available even though num_devices=2 is not set
-    - pytest -v -s tests/compile/passes/distributed/test_fusion_all_reduce.py
+    - pytest -v -s tests/compile/test_fusion_attn.py
+    - pytest -v -s tests/compile/test_silu_mul_quant_fusion.py
+    # this runner has 2 GPUs available even though num_gpus=2 is not set
+    - pytest -v -s tests/compile/distributed/test_fusion_all_reduce.py
+    # Limit to Inductor partition, no custom ops, and allreduce & attn fusion to reduce running time
+    # Wrap with quotes to escape yaml
+    - "pytest -v -s tests/compile/distributed/test_fusions_e2e.py::test_tp2_attn_quant_allreduce_rmsnorm -k 'True and not +quant_fp8 and not +rms_norm'"
     # test_fp8_kv_scale_compile requires FlashAttention (not supported on default L4/L40)
-    # TODO(luka) move to H100 once pass tests run on H100
     - pytest -v -s tests/compile/fullgraph/test_full_graph.py::test_fp8_kv_scale_compile
 
-- label: Fusion E2E Quick (H100)
-  timeout_in_minutes: 15
+- label: Fusion E2E (2 GPUs)(B200)
+  timeout_in_minutes: 40
   working_dir: "/vllm-workspace/"
-  device: h100
-  num_devices: 1
-  source_file_dependencies:
-    - csrc/quantization/
-    - vllm/model_executor/
-    - vllm/v1/attention/
-    - vllm/compilation/
-    - tests/compile/fusions_e2e/
-  commands:
-    - nvidia-smi
-    # Run all models and attn backends but only Inductor partition and native custom ops
-    - pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "inductor_partition and not +rms_norm and not +quant_fp8"
-    # Qwen requires +quant_fp8 as -quant_fp8 rms+quant fusion is not supported
-    - pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "inductor_partition and not +rms_norm and +quant_fp8 and qwen3"
-
-- label: Fusion E2E Config Sweep (H100)
-  timeout_in_minutes: 30
-  working_dir: "/vllm-workspace/"
-  device: h100
-  num_devices: 1
-  source_file_dependencies:
-    - csrc/quantization/
-    - vllm/compilation/
-    # can affect pattern matching
-    - vllm/model_executor/layers/layernorm.py
-    - vllm/model_executor/layers/activation.py
-    - vllm/model_executor/layers/attention/attention.py
-    - vllm/model_executor/layers/quantization/input_quant_fp8.py
-    - tests/compile/fusions_e2e/
-  commands:
-    - nvidia-smi
-    # Run just llama3 (fp8) for all config combinations
-    - pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "llama-3"
-
-- label: Fusion E2E Config Sweep (B200)
-  timeout_in_minutes: 30
-  working_dir: "/vllm-workspace/"
-  device: b200
-  num_devices: 1
+  gpu: b200
   optional: true
-  commands:
-    - nvidia-smi
-    # Run all models but only FLASHINFER, Inductor partition and native custom ops
-    # Qwen requires +quant_fp8 as -quant_fp8 rms+quant fusion is not supported
-    # Run just llama3 (fp8 & fp4) for all config combinations (only inductor partition)
-    - pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "inductor_partition and (FLASHINFER and not +rms_norm and (not +quant_fp8 or +quant_fp8 and qwen3) or llama-3)"
-
-- label: Fusion E2E TP2 Quick (H100)
-  timeout_in_minutes: 20
-  working_dir: "/vllm-workspace/"
-  device: h100
-  num_devices: 2
+  num_gpus: 2
   source_file_dependencies:
-    - csrc/quantization/
-    - vllm/model_executor/
-    - vllm/v1/attention/
-    - vllm/compilation/
-    - tests/compile/fusions_e2e/
+  - csrc/quantization/fp4/
+  - vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
+  - vllm/v1/attention/backends/flashinfer.py
+  - vllm/compilation/
+  # can affect pattern matching
+  - vllm/model_executor/layers/layernorm.py
+  - vllm/model_executor/layers/activation.py
+  - vllm/model_executor/layers/quantization/input_quant_fp8.py
+  - tests/compile/distributed/test_fusions_e2e.py
   commands:
     - nvidia-smi
-    # Run all models and attn backends but only Inductor partition and native custom ops
-    - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k "inductor_partition and not +rms_norm and not +quant_fp8"
-    - pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "inductor_partition and not +rms_norm and not +quant_fp8"
+    # Run all e2e fusion tests
+    - pytest -v -s tests/compile/distributed/test_fusions_e2e.py
 
-- label: Fusion E2E TP2 AR-RMS Config Sweep (H100)
-  timeout_in_minutes: 40
-  working_dir: "/vllm-workspace/"
-  device: h100
-  num_devices: 2
-  source_file_dependencies:
-    - csrc/quantization/
-    - vllm/compilation/
-    # can affect pattern matching
-    - vllm/model_executor/layers/layernorm.py
-    - vllm/model_executor/layers/activation.py
-    - vllm/model_executor/layers/attention/attention.py
-    - vllm/model_executor/layers/quantization/input_quant_fp8.py
-    - tests/compile/fusions_e2e/
-  commands:
-    - nvidia-smi
-    # Run just llama3 (fp8 & bf16) for all config combinations
-    - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k "llama-3"
-
-- label: Fusion E2E TP2 AsyncTP Config Sweep (H100)
-  timeout_in_minutes: 40
-  working_dir: "/vllm-workspace/"
-  device: h100
-  num_devices: 2
-  source_file_dependencies:
-    - csrc/quantization/
-    - vllm/compilation/
-    # can affect pattern matching
-    - vllm/model_executor/layers/layernorm.py
-    - vllm/model_executor/layers/activation.py
-    - vllm/model_executor/layers/attention/attention.py
-    - vllm/model_executor/layers/quantization/input_quant_fp8.py
-    - tests/compile/fusions_e2e/
-  commands:
-    - nvidia-smi
-    # Run just llama3 (fp8 & bf16) for all config combinations
-    - pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "llama-3"
-
-- label: Fusion E2E TP2 (B200)
-  timeout_in_minutes: 20
-  working_dir: "/vllm-workspace/"
-  device: b200
-  num_devices: 2
-  source_file_dependencies:
-    - csrc/quantization/
-    - vllm/model_executor/
-    - vllm/v1/attention/
-    - vllm/compilation/
-    - tests/compile/fusions_e2e/
-  commands:
-    - nvidia-smi
-    # Run all models but only FLASHINFER, Inductor partition and native custom ops
-    # include qwen with +quant_fp8 as -quant_fp8 rms+quant fusion is not supported
-    # for ar-rms-quant-fp4, also sweep llama3
-    - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k "(FLASHINFER and inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and qwen3)) or Llama-3.1-8B-Instruct-FP4"
-    - pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "FLASHINFER and inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and qwen3)"

.buildkite/test_areas/cuda.yaml

@@ -9,7 +9,6 @@ steps:
   - tests/cuda
   commands:
     - pytest -v -s cuda/test_cuda_context.py
-    - pytest -v -s cuda/test_platform_no_cuda_init.py
 
 - label: Cudagraph
   timeout_in_minutes: 20

.buildkite/test_areas/distributed.yaml

@@ -5,7 +5,7 @@ steps:
 - label: Distributed Comm Ops
   timeout_in_minutes: 20
   working_dir: "/vllm-workspace/tests"
-  num_devices: 2
+  num_gpus: 2
   source_file_dependencies:
   - vllm/distributed
   - tests/distributed
@@ -16,9 +16,9 @@ steps:
   - pytest -v -s distributed/test_shm_storage.py
 
 - label: Distributed (2 GPUs)
-  timeout_in_minutes: 60
+  timeout_in_minutes: 90
   working_dir: "/vllm-workspace/tests"
-  num_devices: 2
+  num_gpus: 2
   source_file_dependencies:
   - vllm/compilation/
   - vllm/distributed/
@@ -47,13 +47,14 @@ steps:
   - pytest -v -s ./compile/test_wrapper.py
   - VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
   - VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
+  - pytest -v -s distributed/test_sequence_parallel.py
   - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
   - pytest -v -s v1/worker/test_worker_memory_snapshot.py
 
 - label: Distributed Tests (4 GPUs)
   timeout_in_minutes: 50
   working_dir: "/vllm-workspace/tests"
-  num_devices: 4
+  num_gpus: 4
   source_file_dependencies:
   - vllm/distributed/
   - tests/distributed/test_utils
@@ -62,7 +63,6 @@ steps:
   - tests/compile/fullgraph/test_basic_correctness.py
   - examples/offline_inference/rlhf.py
   - examples/offline_inference/rlhf_colocate.py
-  - examples/offline_inference/new_weight_syncing/
   - tests/examples/offline_inference/data_parallel.py
   - tests/v1/distributed
   - tests/v1/engine/test_engine_core_client.py
@@ -97,19 +97,14 @@ steps:
   - pytest -v -s distributed/test_symm_mem_allreduce.py
   # TODO: create a dedicated test section for multi-GPU example tests
   # when we have multiple distributed example tests
-  # OLD rlhf examples
   - cd ../examples/offline_inference
   - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf.py
   - 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.py
-  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_async_new_apis.py
 
 - label: Distributed Tests (8 GPUs)(H100)
   timeout_in_minutes: 10
-  device: h100
-  num_devices: 8
+  gpu: h100
+  num_gpus: 8
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - examples/offline_inference/torchrun_dp_example.py
@@ -125,9 +120,9 @@ steps:
   - torchrun --nproc-per-node=8 ../examples/offline_inference/torchrun_dp_example.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep
 
 - label: Distributed Tests (4 GPUs)(A100)
-  device: a100
+  gpu: a100
   optional: true
-  num_devices: 4
+  num_gpus: 4
   source_file_dependencies:
   - vllm/
   commands:
@@ -138,22 +133,26 @@ steps:
   - TARGET_TEST_SUITE=A100 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
   - pytest -v -s -x lora/test_mixtral.py
 
-- label: Distributed Tests (2 GPUs)(H100)
-  timeout_in_minutes: 15
-  device: h100
+- label: Distributed Tests (2 GPUs)(H200)
+  gpu: h200
   optional: true
   working_dir: "/vllm-workspace/"
-  num_devices: 2
+  num_gpus: 2
   commands:
+    - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/distributed/test_async_tp.py
+    - pytest -v -s tests/compile/distributed/test_sequence_parallelism.py
+    - pytest -v -s tests/compile/distributed/test_fusion_all_reduce.py
+    - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/distributed/test_fusions_e2e.py -k 'not Llama-4'
+    - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/distributed/test_sequence_parallel.py
     - pytest -v -s tests/distributed/test_context_parallel.py
-    - 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
+    - CUDA_VISIBLE_DEVICES=1,2 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
 
 - label: Distributed Tests (2 GPUs)(B200)
-  device: b200
+  gpu: b200
   optional: true
   working_dir: "/vllm-workspace/"
-  num_devices: 2
+  num_gpus: 2
   commands:
     - pytest -v -s tests/distributed/test_context_parallel.py
     - pytest -v -s tests/distributed/test_nccl_symm_mem_allreduce.py
@@ -162,9 +161,8 @@ steps:
 - label: 2 Node Test (4 GPUs)
   timeout_in_minutes: 30
   working_dir: "/vllm-workspace/tests"
-  num_devices: 2
+  num_gpus: 2
   num_nodes: 2
-  no_plugin: true
   source_file_dependencies:
   - vllm/distributed/
   - vllm/engine/
@@ -173,12 +171,12 @@ steps:
   - tests/distributed/
   - tests/examples/offline_inference/data_parallel.py
   commands:
-    - ./.buildkite/scripts/run-multi-node-test.sh /vllm-workspace/tests 2 2 $IMAGE_TAG "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/offline_inference/data_parallel.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=0 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_multi_node_assignment.py && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_pipeline_parallel.py" "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/offline_inference/data_parallel.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=1 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code"
+    - ./.buildkite/scripts/run-multi-node-test.sh /vllm-workspace/tests 2 2 public.ecr.aws/q9t5s3a7/vllm-ci-postmerge-repo:0bec63fa317e1fbd62e19b0fc31c43c81bf89077 "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/offline_inference/data_parallel.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=0 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_multi_node_assignment.py && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_pipeline_parallel.py" "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/offline_inference/data_parallel.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=1 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code"
 
 - label: Distributed NixlConnector PD accuracy (4 GPUs)
   timeout_in_minutes: 30
   working_dir: "/vllm-workspace/tests"
-  num_devices: 4
+  num_gpus: 4
   source_file_dependencies:
     - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
     - tests/v1/kv_connector/nixl_integration/
@@ -186,21 +184,10 @@ steps:
     - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
     - bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
 
-- label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs)
-  timeout_in_minutes: 30
-  working_dir: "/vllm-workspace/tests"
-  num_devices: 4
-  source_file_dependencies:
-    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
-    - tests/v1/kv_connector/nixl_integration/
-  commands:
-    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
-    - DP_EP=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
-
 - label: Pipeline + Context Parallelism (4 GPUs))
   timeout_in_minutes: 60
   working_dir: "/vllm-workspace/tests"
-  num_devices: 4
+  num_gpus: 4
   source_file_dependencies:
   - vllm/distributed/
   - vllm/engine/
@@ -209,4 +196,4 @@ steps:
   - tests/distributed/
   commands:
   - pytest -v -s distributed/test_pp_cudagraph.py
-  - pytest -v -s distributed/test_pipeline_parallel.py
+  - pytest -v -s distributed/test_pipeline_parallel.py

.buildkite/test_areas/e2e_integration.yaml

@@ -4,27 +4,27 @@ depends_on:
 steps:
 - label: DeepSeek V2-Lite Accuracy
   timeout_in_minutes: 60
-  device: h100
+  gpu: h100
   optional: true
-  num_devices: 4
+  num_gpus: 4
   working_dir: "/vllm-workspace"
   commands:
   - bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh 0.25 200 8010
 
 - label: Qwen3-30B-A3B-FP8-block Accuracy
   timeout_in_minutes: 60
-  device: h100
+  gpu: h100
   optional: true
-  num_devices: 4
+  num_gpus: 4
   working_dir: "/vllm-workspace"
   commands:
   - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020
 
 - label: Qwen3-30B-A3B-FP8-block Accuracy (B200)
   timeout_in_minutes: 60
-  device: b200
+  gpu: b200
   optional: true
-  num_devices: 2
+  num_gpus: 2
   working_dir: "/vllm-workspace"
   commands:
   - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020 2 1
@@ -33,11 +33,10 @@ steps:
   timeout_in_minutes: 30
   optional: true
   soft_fail: true
-  num_devices: 2
+  num_gpus: 2
   working_dir: "/vllm-workspace"
   source_file_dependencies:
   - vllm/
   - .buildkite/scripts/run-prime-rl-test.sh
   commands:
-    - nvidia-smi
     - bash .buildkite/scripts/run-prime-rl-test.sh

.buildkite/test_areas/engine.yaml

@@ -23,8 +23,4 @@ steps:
     # TODO: accuracy does not match, whether setting
     # VLLM_USE_FLASHINFER_SAMPLER or not on H100.
     - pytest -v -s v1/e2e
-    # Run this test standalone for now;
-    # need to untangle use (implicit) use of spawn/fork across the tests.
-    - 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
+    - pytest -v -s v1/engine

.buildkite/test_areas/entrypoints.yaml

@@ -24,11 +24,6 @@ steps:
   - pytest -v -s entrypoints/llm --ignore=entrypoints/llm/test_generate.py --ignore=entrypoints/llm/test_collective_rpc.py
   - pytest -v -s entrypoints/llm/test_generate.py # it needs a clean process
   - pytest -v -s entrypoints/offline_mode # Needs to avoid interference with other tests
-  mirror:
-    amd:
-      device: mi325_1
-      depends_on:
-      - image-build-amd
 
 - label: Entrypoints Integration (API Server 1)
   timeout_in_minutes: 130
@@ -47,13 +42,15 @@ steps:
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - vllm/
-  - tests/entrypoints/rpc
-  - tests/entrypoints/instrumentator
   - tests/tool_use
+  - tests/entrypoints/sleep
+  - tests/entrypoints/instrumentator
+  - tests/entrypoints/rpc
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/instrumentator
   - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc
+  - pytest -v -s entrypoints/instrumentator
+  - pytest -v -s entrypoints/sleep
   - pytest -v -s tool_use
 
 - label: Entrypoints Integration (Pooling)

.buildkite/test_areas/expert_parallelism.yaml

@@ -14,7 +14,7 @@ steps:
 - label: EPLB Execution
   timeout_in_minutes: 20
   working_dir: "/vllm-workspace/tests"
-  num_devices: 4
+  num_gpus: 4
   source_file_dependencies:
   - vllm/distributed/eplb
   - tests/distributed/test_eplb_execute.py

.buildkite/test_areas/kernels.yaml

@@ -15,9 +15,8 @@ steps:
   timeout_in_minutes: 35
   source_file_dependencies:
   - csrc/attention/
+  - vllm/attention
   - vllm/v1/attention
-    # TODO: remove this dependency (https://github.com/vllm-project/vllm/issues/32267)
-  - vllm/model_executor/layers/attention
   - tests/kernels/attention
   commands:
     - pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
@@ -58,8 +57,8 @@ steps:
 
 - label: Kernels DeepGEMM Test (H100)
   timeout_in_minutes: 45
-  device: h100
-  num_devices: 1
+  gpu: h100
+  num_gpus: 1
   source_file_dependencies:
   - tools/install_deepgemm.sh
   - vllm/utils/deep_gemm.py
@@ -78,7 +77,7 @@ steps:
 - label: Kernels (B200)
   timeout_in_minutes: 30
   working_dir: "/vllm-workspace/"
-  device: b200
+  gpu: b200
   # optional: true
   source_file_dependencies:
   - csrc/quantization/fp4/
@@ -86,7 +85,7 @@ steps:
   - csrc/quantization/cutlass_w8a8/moe/
   - vllm/model_executor/layers/fused_moe/cutlass_moe.py
   - vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py
-  - vllm/model_executor/layers/fused_moe/flashinfer_a2a_prepare_finalize.py
+  - vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py
   - vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
   - vllm/v1/attention/backends/flashinfer.py
   - vllm/v1/attention/backends/mla/cutlass_mla.py
@@ -115,55 +114,4 @@ 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_cutedsl_moe.py
-    # e2e
-    - pytest -v -s tests/models/quantization/test_nvfp4.py
-
-- label: Kernels Helion Test
-  timeout_in_minutes: 30
-  device: h100
-  source_file_dependencies:
-  - vllm/utils/import_utils.py
-  - tests/kernels/helion/
-  commands:
-    - pip install helion
-    - pytest -v -s kernels/helion/
-
- 
-- label: Kernels FP8 MoE Test (1 H100)
-  timeout_in_minutes: 90
-  device: h100
-  num_devices: 1
-  optional: true
-  commands:
-    - pytest -v -s kernels/moe/test_cutlass_moe.py
-    - pytest -v -s kernels/moe/test_flashinfer.py
-    - pytest -v -s kernels/moe/test_gpt_oss_triton_kernels.py
-    - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py
-    - pytest -v -s kernels/moe/test_moe.py
-    # - pytest -v -s kernels/moe/test_block_fp8.py - failing on main
-    - pytest -v -s kernels/moe/test_block_int8.py
-    - pytest -v -s kernels/moe/test_triton_moe_no_act_mul.py
-    - pytest -v -s kernels/moe/test_triton_moe_ptpc_fp8.py
-
-- label: Kernels FP8 MoE Test (2 H100s)
-  timeout_in_minutes: 90
-  device: h100
-  num_devices: 2
-  optional: true
-  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
-  num_devices: 1
-  optional: true
-  commands:
-    - pytest -v -s kernels/moe/test_cutedsl_moe.py
-    - pytest -v -s kernels/moe/test_flashinfer_moe.py
-    - pytest -v -s kernels/moe/test_nvfp4_moe.py
-    - pytest -v -s kernels/moe/test_ocp_mx_moe.py
+    - pytest -v -s tests/kernels/moe/test_cutedsl_moe.py

.buildkite/test_areas/lm_eval.yaml

@@ -12,9 +12,9 @@ steps:
   - 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
+  gpu: a100
   optional: true
-  num_devices: 4
+  num_gpus: 4
   working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
   source_file_dependencies:
   - csrc/
@@ -24,9 +24,9 @@ steps:
   - 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
+  gpu: h100
   optional: true
-  num_devices: 4
+  num_gpus: 4
   working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
   source_file_dependencies:
   - csrc/
@@ -37,39 +37,10 @@ steps:
 
 - label: LM Eval Small Models (B200)
   timeout_in_minutes: 120
-  device: b200
+  gpu: b200
   optional: true
   source_file_dependencies:
   - csrc/
   - vllm/model_executor/layers/quantization
   commands:
   - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-blackwell.txt
-
-- label: LM Eval Large Models (H200)
-  timeout_in_minutes: 60
-  device: h200
-  optional: true
-  num_devices: 8
-  commands:
-    - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-h200.txt
-
-- label: MoE Refactor Integration Test (H100 - TEMPORARY)
-  device: h100
-  optional: true
-  num_devices: 2
-  commands:
-    - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor/config-h100.txt
-  
-- label: MoE Refactor Integration Test (B200 - TEMPORARY)
-  device: b200
-  optional: true
-  num_devices: 2
-  commands:
-    - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor/config-b200.txt
-
-- label: MoE Refactor Integration Test (B200 DP - TEMPORARY)
-  device: b200
-  optional: true
-  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

.buildkite/test_areas/lora.yaml

@@ -14,7 +14,7 @@ steps:
 
 - label: LoRA TP (Distributed)
   timeout_in_minutes: 30
-  num_devices: 4
+  num_gpus: 4
   source_file_dependencies:
   - vllm/lora
   - tests/lora

.buildkite/test_areas/misc.yaml

@@ -16,7 +16,7 @@ steps:
     - pytest -v -s v1/sample
     - pytest -v -s v1/logits_processors
     - pytest -v -s v1/worker
-    - pytest -v -s -m 'not slow_test' v1/spec_decode
+    - pytest -v -s v1/spec_decode
     - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
     - pytest -v -s -m 'not cpu_test' v1/metrics
     - pytest -v -s v1/test_oracle.py
@@ -27,12 +27,11 @@ steps:
     - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
 
 - label: V1 Others (CPU)
-  depends_on:
-    - image-build-cpu
+  depends_on: ~
   source_file_dependencies:
     - vllm/
     - tests/v1
-  device: cpu
+  no_gpu: true
   commands:
     # split the test to avoid interference
     - pytest -v -s -m 'cpu_test' v1/core
@@ -72,7 +71,7 @@ steps:
     - python3 offline_inference/vision_language_multi_image.py --seed 0
     - python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
      # for pooling models
-    - python3 pooling/embed/vision_embedding_offline.py --seed 0
+    - python3 pooling/pooling/vision_language_pooling.py --seed 0
     # for features demo
     - python3 offline_inference/prefix_caching.py
     - python3 offline_inference/llm_engine_example.py
@@ -83,7 +82,7 @@ steps:
 
 - label: Metrics, Tracing (2 GPUs)
   timeout_in_minutes: 20
-  num_devices: 2
+  num_gpus: 2
   source_file_dependencies:
   - vllm/
   - tests/v1/tracing
@@ -115,14 +114,12 @@ steps:
   - pytest -v -s utils_
 
 - label: Async Engine, Inputs, Utils, Worker, Config (CPU)
-  depends_on: 
-  - image-build-cpu
+  depends_on: ~
   timeout_in_minutes: 30
   source_file_dependencies:
   - vllm/
   - tests/test_inputs.py
   - tests/test_outputs.py
-  - tests/test_pooling_params.py
   - tests/multimodal
   - tests/renderers
   - tests/standalone_tests/lazy_imports.py
@@ -130,12 +127,11 @@ steps:
   - tests/tool_parsers
   - tests/transformers_utils
   - tests/config
-  device: cpu
+  no_gpu: true
   commands:
   - python3 standalone_tests/lazy_imports.py
   - pytest -v -s test_inputs.py
   - pytest -v -s test_outputs.py
-  - pytest -v -s test_pooling_params.py
   - pytest -v -s -m 'cpu_test' multimodal
   - pytest -v -s renderers
   - pytest -v -s tokenizers_
@@ -146,7 +142,7 @@ steps:
 - label: GPT-OSS Eval (B200)
   timeout_in_minutes: 60
   working_dir: "/vllm-workspace/"
-  device: b200
+  gpu: b200
   optional: true
   source_file_dependencies:
   - tests/evals/gpt_oss
@@ -159,7 +155,7 @@ steps:
 
 - label: Batch Invariance (H100)
   timeout_in_minutes: 25
-  device: h100
+  gpu: h100
   source_file_dependencies:
     - vllm/v1/attention
     - vllm/model_executor/layers
@@ -168,18 +164,4 @@ steps:
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
     - pip install pytest-timeout pytest-forked
     - pytest -v -s v1/determinism/test_batch_invariance.py
-    - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
-  
-- label: Acceptance Length Test (Large Models) # optional
-  timeout_in_minutes: 25
-  gpu: h100
-  optional: true
-  num_gpus: 1
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/v1/spec_decode/
-  - vllm/model_executor/models/mlp_speculator.py
-  - tests/v1/spec_decode/test_acceptance_length.py
-  commands:
-    - export VLLM_ALLOW_INSECURE_SERIALIZATION=1
-    - pytest -v -s v1/spec_decode/test_acceptance_length.py -m slow_test
+    - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py

.buildkite/test_areas/models_basic.yaml

@@ -4,6 +4,7 @@ depends_on:
 steps:
 - label: Basic Models Tests (Initialization)
   timeout_in_minutes: 45
+  mirror_hardwares: [amdexperimental]
   torch_nightly: true
   source_file_dependencies:
   - vllm/
@@ -15,6 +16,7 @@ steps:
 
 - label: Basic Models Tests (Extra Initialization) %N
   timeout_in_minutes: 45
+  mirror_hardwares: [amdexperimental]
   torch_nightly: true
   source_file_dependencies:
   - vllm/model_executor/models/
@@ -31,27 +33,18 @@ steps:
   timeout_in_minutes: 45
   source_file_dependencies:
   - vllm/
-  - tests/models/test_terratorch.py
   - tests/models/test_transformers.py
   - tests/models/test_registry.py
   commands:
-    - pytest -v -s models/test_terratorch.py models/test_transformers.py models/test_registry.py
-  mirror:
-    amd:
-      device: mi325_1
-      depends_on:
-      - image-build-amd
-    
+    - pytest -v -s models/test_transformers.py models/test_registry.py
 
 - label: Basic Models Test (Other CPU) # 5min
-  depends_on: 
-  - image-build-cpu
   timeout_in_minutes: 10
   source_file_dependencies:
   - vllm/
   - tests/models/test_utils.py
   - tests/models/test_vision.py
-  device: cpu
+  no_gpu: true
   commands:
     - pytest -v -s models/test_utils.py models/test_vision.py
 

.buildkite/test_areas/models_distributed.yaml

@@ -5,7 +5,7 @@ steps:
 - label: Distributed Model Tests (2 GPUs)
   timeout_in_minutes: 50
   working_dir: "/vllm-workspace/tests"
-  num_devices: 2
+  num_gpus: 2
   source_file_dependencies:
   - vllm/model_executor/model_loader/sharded_state_loader.py
   - vllm/model_executor/models/

.buildkite/test_areas/models_language.yaml

@@ -4,6 +4,7 @@ depends_on:
 steps:
 - label: Language Models Tests (Standard)
   timeout_in_minutes: 25
+  mirror_hardwares: [amdexperimental]
   torch_nightly: true
   source_file_dependencies:
   - vllm/
@@ -15,6 +16,7 @@ steps:
 
 - label: Language Models Tests (Extra Standard) %N
   timeout_in_minutes: 45
+  mirror_hardwares: [amdexperimental]
   torch_nightly: true
   source_file_dependencies:
   - vllm/model_executor/models/
@@ -30,6 +32,7 @@ steps:
 
 - label: Language Models Tests (Hybrid) %N
   timeout_in_minutes: 75
+  mirror_hardwares: [amdexperimental]
   torch_nightly: true
   source_file_dependencies:
   - vllm/
@@ -37,7 +40,7 @@ steps:
   commands:
     # Install fast path packages for testing against transformers
     # Note: also needed to run plamo2 model in vLLM
-    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
+    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.2.5'
     - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.2'
     # Shard hybrid language model tests
     - pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
@@ -45,6 +48,7 @@ steps:
 
 - label: Language Models Test (Extended Generation) # 80min
   timeout_in_minutes: 110
+  mirror_hardwares: [amdexperimental]
   optional: true
   source_file_dependencies:
   - vllm/
@@ -52,12 +56,13 @@ steps:
   commands:
     # Install fast path packages for testing against transformers
     # Note: also needed to run plamo2 model in vLLM
-    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
+    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.2.5'
     - 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
+  mirror_hardwares: [amdexperimental]
   optional: true
   source_file_dependencies:
   - vllm/
@@ -67,6 +72,7 @@ steps:
 
 - label: Language Models Test (Extended Pooling)  # 36min
   timeout_in_minutes: 50
+  mirror_hardwares: [amdexperimental]
   optional: true
   source_file_dependencies:
   - vllm/
@@ -76,6 +82,7 @@ steps:
 
 - label: Language Models Test (MTEB)
   timeout_in_minutes: 110
+  mirror_hardwares: [amdexperimental]
   optional: true
   source_file_dependencies:
   - vllm/

.buildkite/test_areas/models_multimodal.yaml

@@ -14,13 +14,11 @@ steps:
     - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model  # Otherwise, mp_method="spawn" doesn't work
 
 - label: Multi-Modal Processor Test (CPU)
-  depends_on: 
-  - image-build-cpu
   timeout_in_minutes: 60
   source_file_dependencies:
   - vllm/
   - tests/models/multimodal
-  device: cpu
+  no_gpu: true
   commands:
     - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
     - pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py

.buildkite/test_areas/plugins.yaml

@@ -5,7 +5,7 @@ steps:
 - label: Plugin Tests (2 GPUs)
   timeout_in_minutes: 60
   working_dir: "/vllm-workspace/tests"
-  num_devices: 2
+  num_gpus: 2
   source_file_dependencies:
   - vllm/plugins/
   - tests/plugins/

.buildkite/test_areas/pytorch.yaml

@@ -3,7 +3,7 @@ depends_on:
   - image-build
 steps:
 - label: PyTorch Compilation Unit Tests
-  timeout_in_minutes: 10
+  timeout_in_minutes: 30
   source_file_dependencies:
     - vllm/
     - tests/compile
@@ -17,16 +17,8 @@ steps:
   # (using -0 for proper path handling)
   - "find compile/ -maxdepth 1 -name 'test_*.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
 
-- label: PyTorch Compilation Passes Unit Tests
-  timeout_in_minutes: 20
-  source_file_dependencies:
-    - vllm/
-    - tests/compile/passes
-  commands:
-  - pytest -s -v compile/passes --ignore compile/passes/distributed
-
 - label: PyTorch Fullgraph Smoke Test
-  timeout_in_minutes: 35
+  timeout_in_minutes: 30
   source_file_dependencies:
   - vllm/
   - tests/compile
@@ -38,13 +30,16 @@ steps:
   - "find compile/fullgraph/ -name 'test_*.py' -not -name 'test_full_graph.py' -exec pytest -s -v {} \\;"
 
 - label: PyTorch Fullgraph
-  timeout_in_minutes: 30
+  timeout_in_minutes: 40
   source_file_dependencies:
   - vllm/
   - tests/compile
   commands:
     # fp8 kv scales not supported on sm89, tested on Blackwell instead
   - pytest -v -s compile/fullgraph/test_full_graph.py -k 'not test_fp8_kv_scale_compile'
+    # Limit to no custom ops to reduce running time
+    # Wrap with quotes to escape yaml and avoid starting -k string with a -
+  - "pytest -v -s compile/distributed/test_fusions_e2e.py -k 'TRITON and not +quant_fp8 and not Llama-4'"
 
 - label: Pytorch Nightly Dependency Override Check # 2min
   # if this test fails, it means the nightly torch version is not compatible with some

.buildkite/test_areas/quantization.yaml

@@ -16,14 +16,14 @@ steps:
   # https://github.com/pytorch/ao/issues/2919, we'll have to skip new torchao tests for now
   # we can only upgrade after this is resolved
   # TODO(jerryzh168): resolve the above comment
-  - uv pip install --system torchao==0.14.1 --index-url https://download.pytorch.org/whl/cu129
+  - uv pip install --system torchao==0.13.0 --index-url https://download.pytorch.org/whl/cu129
   - uv pip install --system conch-triton-kernels
   - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py
 
 - label: Quantized MoE Test (B200)
   timeout_in_minutes: 60
   working_dir: "/vllm-workspace/"
-  device: b200
+  gpu: b200
   source_file_dependencies:
   - tests/quantization/test_blackwell_moe.py
   - vllm/model_executor/models/deepseek_v2.py

.buildkite/test_areas/samplers.yaml

@@ -12,10 +12,3 @@ steps:
   commands:
     - pytest -v -s samplers
     - VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers
-  mirror:
-    amd:
-      device: mi325_1
-      depends_on:
-      - image-build-amd
-      commands:
-      - pytest -v -s -m 'not skip_v1' samplers

.buildkite/test_areas/weight_loading.yaml

@@ -5,7 +5,7 @@ steps:
 - label: Weight Loading Multiple GPU  # 33min
   timeout_in_minutes: 45
   working_dir: "/vllm-workspace/tests"
-  num_devices: 2
+  num_gpus: 2
   optional: true
   source_file_dependencies:
   - vllm/
@@ -15,8 +15,8 @@ steps:
 
 - label: Weight Loading Multiple GPU - Large Models # optional
   working_dir: "/vllm-workspace/tests"
-  num_devices: 2
-  device: a100
+  num_gpus: 2
+  gpu: a100
   optional: true
   source_file_dependencies:
   - vllm/

.github/CODEOWNERS

@@ -2,10 +2,8 @@
 # for more info about CODEOWNERS file
 
 # This lists cover the "core" components of vLLM that require careful review
-/vllm/compilation @zou3519 @youkaichao @ProExpertProg
-/vllm/distributed/kv_transfer @NickLucche @ApostaC @orozery
-/vllm/lora @jeejeelee
-/vllm/model_executor/layers/attention @LucasWilkinson
+/vllm/attention @LucasWilkinson
+/vllm/executor/executor_base.py @zhuohan123 @youkaichao @alexm-redhat @njhill @22quinn
 /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
@@ -13,34 +11,18 @@
 /vllm/model_executor/layers/batch_invariant.py @yewentao256 
 /vllm/multimodal @DarkLight1337 @ywang96 @NickLucche @tjtanaa
 /vllm/vllm_flash_attn @LucasWilkinson
+/vllm/lora @jeejeelee
+/vllm/reasoning @aarnphm @chaunceyjiang
+/vllm/entrypoints @aarnphm @chaunceyjiang
+/vllm/tool_parsers @aarnphm @chaunceyjiang
+/vllm/compilation @zou3519 @youkaichao @ProExpertProg
+/vllm/distributed/kv_transfer @NickLucche @ApostaC
 CMakeLists.txt @tlrmchlsmth @LucasWilkinson
 
 # Any change to the VllmConfig changes can have a large user-facing impact,
 # so spam a lot of people
 /vllm/config @WoosukKwon @youkaichao @robertgshaw2-redhat @mgoin @tlrmchlsmth @houseroad @hmellor @yewentao256 @ProExpertProg
-/vllm/config/cache.py @heheda12345
-
-# Entrypoints
-/vllm/entrypoints/anthropic @mgoin @DarkLight1337
-/vllm/entrypoints/cli @hmellor @mgoin @DarkLight1337 @russellb
-/vllm/entrypoints/mcp @heheda12345
-/vllm/entrypoints/openai @aarnphm @chaunceyjiang @DarkLight1337 @russellb
-/vllm/entrypoints/openai/realtime @njhill
-/vllm/entrypoints/openai/speech_to_text @NickLucche
-/vllm/entrypoints/pooling @noooop
-/vllm/entrypoints/sagemaker @DarkLight1337
-/vllm/entrypoints/serve @njhill
-/vllm/entrypoints/*.py @njhill
-/vllm/entrypoints/chat_utils.py @DarkLight1337
-/vllm/entrypoints/llm.py @DarkLight1337
-
-# Input/Output Processing
-/vllm/sampling_params.py @njhill @NickLucche
-/vllm/pooling_params.py @noooop @DarkLight1337
-/vllm/tokenizers @DarkLight1337 @njhill
-/vllm/renderers @DarkLight1337 @njhill
-/vllm/reasoning @aarnphm @chaunceyjiang
-/vllm/tool_parsers @aarnphm @chaunceyjiang
+/vllm/config/cache.py @WoosukKwon @youkaichao @robertgshaw2-redhat @mgoin @tlrmchlsmth @houseroad @hmellor @yewentao256 @ProExpertProg @heheda12345
 
 # vLLM V1
 /vllm/v1/attention @LucasWilkinson
@@ -48,14 +30,12 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
 /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/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC
 /vllm/v1/sample @22quinn @houseroad @njhill
 /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/worker/gpu/kv_connector.py @orozery
-/vllm/v1/worker/kv_connector_model_runner_mixin.py @orozery
+/vllm/v1/offloading @ApostaC
 
 # Model runner V2
 /vllm/v1/worker/gpu @WoosukKwon
@@ -74,13 +54,13 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
 /tests/test_inputs.py @DarkLight1337 @ywang96
 /tests/v1/entrypoints/llm/test_struct_output_generate.py @mgoin @russellb @aarnphm
 /tests/v1/structured_output @mgoin @russellb @aarnphm
-/tests/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC @orozery
+/tests/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC
 /tests/weight_loading @mgoin @youkaichao @yewentao256
 /tests/lora @jeejeelee
 /tests/models/language/generation/test_hybrid.py @tdoublep
 /tests/v1/kv_connector/nixl_integration @NickLucche
-/tests/v1/kv_connector @ApostaC @orozery
-/tests/v1/kv_offload @ApostaC @orozery
+/tests/v1/kv_connector @ApostaC
+/tests/v1/offloading @ApostaC
 /tests/v1/determinism @yewentao256 
 
 # Transformers modeling backend
@@ -133,8 +113,8 @@ mkdocs.yaml @hmellor
 /vllm/model_executor/models/mixtral*.py @patrickvonplaten
 /vllm/model_executor/models/voxtral*.py @patrickvonplaten
 /vllm/model_executor/models/pixtral*.py @patrickvonplaten
-/vllm/tokenizers/mistral.py @patrickvonplaten
 /vllm/transformers_utils/configs/mistral.py @patrickvonplaten
+/vllm/transformers_utils/tokenizers/mistral.py @patrickvonplaten
 
 # Kernels
 /vllm/v1/attention/ops/chunked_prefill_paged_decode.py @tdoublep
@@ -170,7 +150,9 @@ mkdocs.yaml @hmellor
 /examples/pooling @noooop
 /tests/models/*/pooling* @noooop
 /tests/entrypoints/pooling @noooop
+/vllm/entrypoints/pooling @noooop
 /vllm/config/pooler.py @noooop
+/vllm/pooling_params.py @noooop
 /vllm/model_executor/layers/pooler @noooop
 
 # Security guide and policies

.github/workflows/cleanup_pr_body.yml

@@ -19,7 +19,6 @@ jobs:
         uses: actions/setup-python@83679a892e2d95755f2dac6acb0bfd1e9ac5d548 # v6.1.0
         with:
           python-version: '3.12'
-          cache: 'pip'
 
       - name: Install Python dependencies
         run: |

.gitignore

@@ -238,6 +238,3 @@ ep_kernels_workspace/
 vllm/grpc/vllm_engine_pb2.py
 vllm/grpc/vllm_engine_pb2_grpc.py
 vllm/grpc/vllm_engine_pb2.pyi
-
-# Ignore generated cpu headers 
-csrc/cpu/cpu_attn_dispatch_generated.h

.pre-commit-config.yaml

@@ -121,9 +121,24 @@ repos:
     name: Update Dockerfile dependency graph
     entry: tools/pre_commit/update-dockerfile-graph.sh
     language: script
-  - id: check-forbidden-imports
-    name: Check for forbidden imports
-    entry: python tools/pre_commit/check_forbidden_imports.py
+  - id: enforce-import-regex-instead-of-re
+    name: Enforce import regex as re
+    entry: python tools/pre_commit/enforce_regex_import.py
+    language: python
+    types: [python]
+    pass_filenames: false
+    additional_dependencies: [regex]
+  # forbid directly import triton
+  - id: forbid-direct-triton-import
+    name: "Forbid direct 'import triton'"
+    entry: python tools/pre_commit/check_triton_import.py
+    language: python
+    types: [python]
+    pass_filenames: false
+    additional_dependencies: [regex]
+  - id: check-pickle-imports
+    name: Prevent new pickle/cloudpickle imports
+    entry: python tools/pre_commit/check_pickle_imports.py
     language: python
     types: [python]
     additional_dependencies: [regex]
@@ -139,15 +154,6 @@ repos:
     files: ^docker/(Dockerfile|versions\.json)$
     pass_filenames: false
     additional_dependencies: [dockerfile-parse]
-  - id: attention-backend-docs
-    name: Check attention backend documentation is up to date
-    entry: python tools/pre_commit/generate_attention_backend_docs.py --check
-    language: python
-  - id: check-boolean-context-manager
-    name: Check for boolean ops in with-statements
-    entry: python tools/pre_commit/check_boolean_context_manager.py
-    language: python
-    types: [python]
   # Keep `suggestion` last
   - id: suggestion
     name: Suggestion

.readthedocs.yaml

@@ -9,14 +9,13 @@ build:
     python: "3.12"
   jobs:
     post_checkout:
-      - git fetch origin main --unshallow --no-tags --filter=blob:none || true
-    pre_create_environment:
-      - pip install uv
-    create_environment:
-      - uv venv $READTHEDOCS_VIRTUALENV_PATH
-    install:
-      - uv pip install --python $READTHEDOCS_VIRTUALENV_PATH/bin/python --no-cache-dir -r requirements/docs.txt 
+      - git fetch --unshallow || true
 
 mkdocs:
   configuration: mkdocs.yaml
   fail_on_warning: true
+
+# Optionally declare the Python requirements required to build your docs
+python:
+  install:
+    - requirements: requirements/docs.txt

CMakeLists.txt

@@ -56,8 +56,8 @@ endif()
 # requirements.txt files and should be kept consistent.  The ROCm torch
 # versions are derived from docker/Dockerfile.rocm
 #
-set(TORCH_SUPPORTED_VERSION_CUDA "2.10.0")
-set(TORCH_SUPPORTED_VERSION_ROCM "2.10.0")
+set(TORCH_SUPPORTED_VERSION_CUDA "2.9.1")
+set(TORCH_SUPPORTED_VERSION_ROCM "2.9.1")
 
 #
 # Try to find python package with an executable that exactly matches
@@ -293,7 +293,6 @@ set(VLLM_EXT_SRC
   "csrc/fused_qknorm_rope_kernel.cu"
   "csrc/layernorm_quant_kernels.cu"
   "csrc/sampler.cu"
-  "csrc/topk.cu"
   "csrc/cuda_view.cu"
   "csrc/quantization/gptq/q_gemm.cu"
   "csrc/quantization/w8a8/int8/scaled_quant.cu"
@@ -434,7 +433,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
       list(APPEND VLLM_EXT_SRC ${MARLIN_TEMPLATE_BF16_KERNEL_SRC})
     endif()
 
-    if (MARLIN_SM75_ARCHS)
+    if (MARLIN_SM75_ARCHS) 
       file(GLOB MARLIN_TEMPLATE_SM75_KERNEL_SRC "csrc/quantization/marlin/sm75_kernel_*.cu")
       set_gencode_flags_for_srcs(
         SRCS "${MARLIN_TEMPLATE_SM75_KERNEL_SRC}"
@@ -446,7 +445,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
       list(APPEND VLLM_EXT_SRC ${MARLIN_TEMPLATE_SM75_KERNEL_SRC})
     endif()
 
-    if (MARLIN_FP8_ARCHS)
+    if (MARLIN_FP8_ARCHS) 
       file(GLOB MARLIN_TEMPLATE_FP8_KERNEL_SRC "csrc/quantization/marlin/sm89_kernel_*.cu")
       set_gencode_flags_for_srcs(
         SRCS "${MARLIN_TEMPLATE_FP8_KERNEL_SRC}"
@@ -459,6 +458,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
     endif()
 
     set(MARLIN_SRCS
+       "csrc/quantization/marlin/sparse/marlin_24_cuda_kernel.cu"
        "csrc/quantization/marlin/marlin.cu"
        "csrc/quantization/marlin/marlin_int4_fp8_preprocess.cu"
        "csrc/quantization/marlin/gptq_marlin_repack.cu"
@@ -1043,7 +1043,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
       list(APPEND VLLM_MOE_EXT_SRC ${MARLIN_MOE_SRC})
     endif()
 
-    if (MARLIN_MOE_SM75_ARCHS)
+    if (MARLIN_MOE_SM75_ARCHS) 
       file(GLOB MARLIN_MOE_SM75_SRC "csrc/moe/marlin_moe_wna16/sm75_kernel_*.cu")
       set_gencode_flags_for_srcs(
         SRCS "${MARLIN_MOE_SM75_SRC}"

benchmarks/README.md

@@ -11,7 +11,7 @@ This directory used to contain vLLM's benchmark scripts and utilities for perfor
 
 ## Usage
 
-For detailed usage instructions, examples, and dataset information, see the [Benchmark CLI documentation](https://docs.vllm.ai/en/latest/benchmarking/cli/#benchmark-cli).
+For detailed usage instructions, examples, and dataset information, see the [Benchmark CLI documentation](https://docs.vllm.ai/en/latest/contributing/benchmarks.html#benchmark-cli).
 
 For full CLI reference see:
 

benchmarks/attention_benchmarks/README.md

@@ -1,266 +0,0 @@
-# vLLM Attention Benchmarking Suite
-
-Fast, flexible benchmarking for vLLM attention and MLA backends with an extended batch specification grammar.
-
-## Quick Start
-
-```bash
-cd benchmarks/attention_benchmarks
-
-# Run a pre-configured benchmark
-python benchmark.py --config configs/mla_decode.yaml
-python benchmark.py --config configs/mla_mixed_batch.yaml
-python benchmark.py --config configs/speculative_decode.yaml
-python benchmark.py --config configs/standard_attention.yaml
-python benchmark.py --config configs/reorder_threshold.yaml
-
-# Or run custom benchmarks
-python benchmark.py \
-    --backends flash flashinfer \
-    --batch-specs "q2k" "8q1s1k" "2q2k_32q1s1k" \
-    --output-csv results.csv
-```
-
-## Simplified Batch Specification Grammar
-
-Express workloads concisely using query length and sequence length:
-
-```python
-"q2k"              # 2048-token prefill (q_len=2048, seq_len=2048)
-"q1s1k"            # Decode: 1 token with 1K sequence
-"8q1s1k"           # 8 decode requests
-"q4s1k"            # 4-token extend (e.g., spec decode)
-"2q2k_32q1s1k"     # Mixed: 2 prefills + 32 decodes
-"16q4s1k"          # 16 spec decode (4 tokens each)
-```
-
-### Grammar Rule
-
-```text
-Format: (<count>?) q<q_len>(k?) (s<seq_len>(k?))?
-
-- count:   Number of identical requests (optional, default=1)
-- q_len:   Query length (number of new tokens)
-- seq_len: Total sequence length (optional, defaults to q_len for prefill)
-- 'k':     Multiplies value by 1024
-
-Mixed batches: Use _ to combine (e.g., "2q2k_32q1s1k")
-```
-
-**Note**: Decode, prefill, and spec decode are just different query lengths - no special syntax needed!
-
-## Pre-configured Benchmarks
-
-The suite includes several pre-configured YAML benchmark configurations:
-
-### MLA Decode Benchmark
-
-Tests pure decode performance across MLA backends with varying batch sizes and sequence lengths.
-
-```bash
-python benchmark.py --config configs/mla_decode.yaml
-```
-
-### MLA Mixed Batch Benchmark
-
-Tests chunked prefill performance with mixed prefill + decode batches.
-
-```bash
-python benchmark.py --config configs/mla_mixed_batch.yaml
-```
-
-### Speculative Decoding Benchmark
-
-Tests speculative decode scenarios (K-token verification) and reorder_batch_threshold optimization.
-
-```bash
-python benchmark.py --config configs/speculative_decode.yaml
-```
-
-### Standard Attention Benchmark
-
-Tests standard attention backends (Flash/Triton/FlashInfer) with pure prefill, decode, and mixed batches.
-
-```bash
-python benchmark.py --config configs/standard_attention.yaml
-```
-
-### Reorder Threshold Study
-
-**Question:** At what query length does the prefill pipeline become faster than the decode pipeline?
-
-Tests query lengths from 1-1024 across 9 batch sizes to find the crossover point. Uses `decode_vs_prefill` mode to compare both pipelines for each query length.
-
-```bash
-python benchmark.py --config configs/reorder_threshold.yaml
-```
-
----
-
-## Universal Benchmark
-
-The `benchmark.py` script handles **all** backends - both standard attention and MLA.
-
-### Standard Attention (Flash/Triton/FlashInfer)
-
-```bash
-python benchmark.py \
-    --backends flash triton flashinfer \
-    --batch-specs "q2k" "8q1s1k" "2q2k_32q1s1k" \
-    --num-layers 10 \
-    --repeats 5 \
-    --output-csv results.csv
-```
-
-### MLA Backends
-
-```bash
-# Compare all MLA backends
-python benchmark.py \
-    --backends cutlass_mla flashinfer_mla flashattn_mla flashmla \
-    --batch-specs "64q1s1k" "64q1s4k" \
-    --output-csv mla_results.csv
-```
-
-### Parameter Sweeps
-
-Use `--sweep-param` and `--sweep-values` to run parameter sweeps from the CLI:
-
-#### CUTLASS MLA num-splits Optimization
-
-**Question:** What is the optimal `num_kv_splits` for CUTLASS MLA?
-
-```bash
-python benchmark.py \
-    --backend cutlass_mla \
-    --batch-specs "64q1s1k" "64q1s4k" "64q1s16k" \
-    --sweep-param num_kv_splits \
-    --sweep-values 1 2 4 8 16 \
-    --output-json optimal_splits.json
-```
-
-#### Reorder Batch Threshold Optimization
-
-**Question:** What's the optimal `reorder_batch_threshold` for speculative decoding?
-
-```bash
-python benchmark.py \
-    --backend flashmla \
-    --batch-specs "q4s1k" "q8s2k" \
-    --sweep-param reorder_batch_threshold \
-    --sweep-values 1 4 16 64 256 512 \
-    --output-csv threshold_sweep.csv
-```
-
-### All Command-Line Options
-
-```text
---config CONFIG                     # Path to YAML config file (overrides other args)
---backends BACKEND [BACKEND ...]    # flash, triton, flashinfer, cutlass_mla,
-                                    # flashinfer_mla, flashattn_mla, flashmla
---backend BACKEND                   # Single backend (alternative to --backends)
---batch-specs SPEC [SPEC ...]       # Batch specifications using extended grammar
-
-# Model configuration
---num-layers N                      # Number of layers
---head-dim N                        # Head dimension
---num-q-heads N                     # Query heads
---num-kv-heads N                    # KV heads
---block-size N                      # Block size
-
-# Benchmark settings
---device DEVICE                     # Device (default: cuda:0)
---repeats N                         # Repetitions
---warmup-iters N                    # Warmup iterations
---profile-memory                    # Profile memory usage
-
-# Parameter sweeps
---sweep-param PARAM                 # Parameter name to sweep (e.g., num_kv_splits,
-                                    # reorder_batch_threshold)
---sweep-values N [N ...]            # Values to sweep for the parameter
-
-# Output
---output-csv FILE                   # Save to CSV
---output-json FILE                  # Save to JSON
-```
-
-## Hardware Requirements
-
-| Backend | Hardware |
-|---------|----------|
-| Flash/Triton/FlashInfer | Any CUDA GPU |
-| CUTLASS MLA | Blackwell (SM100+) |
-| FlashAttn MLA | Hopper (SM90+) |
-| FlashMLA | Hopper (SM90+) |
-| FlashInfer-MLA | Any CUDA GPU |
-
-## Using MLA Runner Directly
-
-All MLA backends are available through `mla_runner.run_mla_benchmark()`:
-
-```python
-from mla_runner import run_mla_benchmark
-from common import BenchmarkConfig
-
-config = BenchmarkConfig(
-    backend="cutlass_mla",
-    batch_spec="64q1s4k",
-    num_layers=10,
-    head_dim=576,
-    num_q_heads=128,
-    num_kv_heads=1,
-    block_size=128,
-    device="cuda:0",
-    repeats=5,
-    warmup_iters=3,
-)
-
-# CUTLASS MLA with specific num_kv_splits
-result = run_mla_benchmark("cutlass_mla", config, num_kv_splits=4)
-print(f"Time: {result.mean_time:.6f}s")
-
-# FlashInfer-MLA
-result = run_mla_benchmark("flashinfer_mla", config)
-
-# FlashAttn MLA (Hopper SM90+)
-result = run_mla_benchmark("flashattn_mla", config, reorder_batch_threshold=64)
-
-# FlashMLA (Hopper SM90+)
-result = run_mla_benchmark("flashmla", config, reorder_batch_threshold=64)
-```
-
-## Python API
-
-```python
-from batch_spec import parse_batch_spec, format_batch_spec, get_batch_stats
-from common import BenchmarkConfig, BenchmarkResult, ResultsFormatter
-
-# Parse batch specs
-requests = parse_batch_spec("2q2k_q4s1k_32q1s1k")
-print(format_batch_spec(requests))
-# "2 prefill (2x2k), 1 extend (1xq4kv1k), 32 decode (32x1k)"
-
-# Get batch statistics
-stats = get_batch_stats(requests)
-print(f"Total tokens: {stats['total_tokens']}")
-print(f"Num decode: {stats['num_decode']}, Num prefill: {stats['num_prefill']}")
-
-# Format results
-formatter = ResultsFormatter()
-formatter.save_csv(results, "output.csv")
-formatter.save_json(results, "output.json")
-```
-
-## Tips
-
-**1. Warmup matters** - Use `--warmup-iters 10` for stable results
-
-**2. Multiple repeats** - Use `--repeats 20` for low variance
-
-**3. Save results** - Always use `--output-csv` or `--output-json`
-
-**4. Test incrementally** - Start with `--num-layers 1 --repeats 1`
-
-**5. Extended grammar** - Leverage spec decode, chunked prefill patterns
-
-**6. Parameter sweeps** - Use `--sweep-param` and `--sweep-values` to find optimal values

benchmarks/attention_benchmarks/__init__.py

@@ -1,44 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""vLLM Attention Benchmarking Suite."""
-
-from .batch_spec import (
-    BatchRequest,
-    format_batch_spec,
-    get_batch_stats,
-    parse_batch_spec,
-    reorder_for_flashinfer,
-    split_by_type,
-)
-from .common import (
-    BenchmarkConfig,
-    BenchmarkResult,
-    MockLayer,
-    MockModelConfig,
-    ResultsFormatter,
-    get_attention_scale,
-    is_mla_backend,
-    setup_mla_dims,
-)
-
-__all__ = [
-    # Batch specification
-    "BatchRequest",
-    "parse_batch_spec",
-    "format_batch_spec",
-    "reorder_for_flashinfer",
-    "split_by_type",
-    "get_batch_stats",
-    # Benchmarking infrastructure
-    "BenchmarkConfig",
-    "BenchmarkResult",
-    "ResultsFormatter",
-    # Mock objects
-    "MockLayer",
-    "MockModelConfig",
-    # Utilities
-    "setup_mla_dims",
-    "get_attention_scale",
-    "is_mla_backend",
-]

benchmarks/attention_benchmarks/batch_spec.py

@@ -1,268 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""
-Simplified batch specification grammar for attention benchmarks.
-
-Grammar (underscore-separated segments):
-  Format: (<count>?) q<q_len>(k?) (s<seq_len>(k?))?
-
-  - count: Number of identical requests (optional, default=1)
-  - q_len: Query length (number of new tokens)
-  - seq_len: Total sequence length (optional, defaults to q_len for prefill)
-  - 'k' suffix: Multiplies value by 1024
-
-Common patterns:
-  - Prefill:  q_len == seq_len  (e.g., "q2k" → 2048 new tokens, 2048 seq)
-  - Decode:   q_len == 1        (e.g., "q1s1k" → 1 token, 1024 seq length)
-  - Extend:   q_len < seq_len   (e.g., "q4s1k" → 4 tokens, 1024 seq length)
-
-Examples:
-  q2k              -> [(2048, 2048)]           # Prefill: 2048 tokens
-  q1s1k            -> [(1, 1024)]              # Decode: 1 token, 1K sequence
-  8q1s1k           -> [(1, 1024)] * 8          # 8 decode requests
-  q4s1k            -> [(4, 1024)]              # 4-token extend (spec decode)
-  2q1k_32q1s1k     -> [(1024, 1024)] * 2 + [(1, 1024)] * 32  # Mixed batch
-  16q4s1k          -> [(4, 1024)] * 16         # 16 spec decode requests
-"""
-
-from collections import Counter
-from dataclasses import dataclass
-
-import regex as re
-
-
-@dataclass
-class BatchRequest:
-    """Represents a single request in a batch."""
-
-    q_len: int  # Query length (number of new tokens)
-    kv_len: int  # Total KV cache length
-
-    @property
-    def is_decode(self) -> bool:
-        """True if this is a decode request (q_len == 1)."""
-        return self.q_len == 1
-
-    @property
-    def is_prefill(self) -> bool:
-        """True if this is a pure prefill (q_len == kv_len)."""
-        return self.q_len == self.kv_len
-
-    @property
-    def is_extend(self) -> bool:
-        """True if this is context extension (q_len > 1, kv_len > q_len)."""
-        return self.q_len > 1 and self.kv_len > self.q_len
-
-    @property
-    def context_len(self) -> int:
-        """Context length (KV cache - query)."""
-        return self.kv_len - self.q_len
-
-    def as_tuple(self) -> tuple[int, int]:
-        """Return as (q_len, kv_len) tuple for compatibility."""
-        return (self.q_len, self.kv_len)
-
-
-def _parse_size(size_str: str, k_suffix: str) -> int:
-    """Parse size string with optional 'k' suffix."""
-    size = int(size_str)
-    return size * 1024 if k_suffix == "k" else size
-
-
-def parse_batch_spec(spec: str) -> list[BatchRequest]:
-    """
-    Parse batch specification string into list of BatchRequest objects.
-
-    Grammar: (<count>?) q<q_len>(k?) (s<seq_len>(k?))?
-
-    Args:
-        spec: Batch specification string (see module docstring for grammar)
-
-    Returns:
-        List of BatchRequest objects
-
-    Raises:
-        ValueError: If spec format is invalid
-    """
-    requests = []
-
-    for seg in spec.split("_"):
-        # Unified pattern: (<count>?) q<q_len>(k?) (s<seq_len>(k?))?
-        m = re.match(r"^(?:(\d+))?q(\d+)(k?)(?:s(\d+)(k?))?$", seg)
-        if m:
-            cnt = int(m.group(1)) if m.group(1) else 1
-            q_len = _parse_size(m.group(2), m.group(3))
-            kv_len = _parse_size(m.group(4), m.group(5)) if m.group(4) else q_len
-            requests.extend([BatchRequest(q_len=q_len, kv_len=kv_len)] * cnt)
-            continue
-
-        raise ValueError(f"Invalid batch spec segment: '{seg}'")
-
-    return requests
-
-
-def format_batch_spec(requests: list[BatchRequest]) -> str:
-    """
-    Format list of BatchRequest into human-readable string.
-
-    Groups requests by type and provides counts and sizes.
-
-    Args:
-        requests: List of BatchRequest objects
-
-    Returns:
-        Formatted string describing the batch
-    """
-    kinds = {
-        "prefill": [],
-        "extend": [],
-        "decode": [],
-    }
-
-    for req in requests:
-        tup = (req.q_len, req.kv_len)
-        if req.is_prefill:
-            kinds["prefill"].append(tup)
-        elif req.is_extend:
-            kinds["extend"].append(tup)
-        elif req.is_decode:
-            kinds["decode"].append(tup)
-
-    parts = []
-    for kind in ["prefill", "extend", "decode"]:
-        lst = kinds[kind]
-        if not lst:
-            continue
-
-        cnt_total = len(lst)
-        ctr = Counter(lst)
-        inner = []
-
-        for (q, kv), cnt in ctr.items():
-            if kind == "prefill":
-                size = f"{q // 1024}k" if q % 1024 == 0 else str(q)
-                inner.append(f"{cnt}x{size}")
-            elif kind == "decode":
-                size = f"{kv // 1024}k" if kv % 1024 == 0 else str(kv)
-                inner.append(f"{cnt}x{size}")
-            else:  # extend
-                qstr = f"{q // 1024}k" if q % 1024 == 0 else str(q)
-                kstr = f"{kv // 1024}k" if kv % 1024 == 0 else str(kv)
-                inner.append(f"{cnt}xq{qstr}kv{kstr}")
-
-        parts.append(f"{cnt_total} {kind} ({', '.join(inner)})")
-
-    return ", ".join(parts)
-
-
-def reorder_for_flashinfer(requests: list[BatchRequest]) -> list[BatchRequest]:
-    """
-    Reorder requests for FlashInfer: decode first, then prefill.
-
-    FlashInfer expects decode requests before prefill requests for
-    optimal performance.
-
-    Args:
-        requests: Original list of BatchRequest
-
-    Returns:
-        Reordered list with decode requests first
-    """
-    decodes = [r for r in requests if r.is_decode]
-    non_decodes = [r for r in requests if not r.is_decode]
-    return decodes + non_decodes
-
-
-def split_by_type(
-    requests: list[BatchRequest],
-) -> dict[str, list[BatchRequest]]:
-    """
-    Split requests by type for analysis.
-
-    Args:
-        requests: List of BatchRequest
-
-    Returns:
-        Dict with keys: 'decode', 'prefill', 'extend'
-    """
-    result = {
-        "decode": [],
-        "prefill": [],
-        "extend": [],
-    }
-
-    for req in requests:
-        if req.is_decode:
-            result["decode"].append(req)
-        elif req.is_prefill:
-            result["prefill"].append(req)
-        elif req.is_extend:
-            result["extend"].append(req)
-
-    return result
-
-
-def get_batch_stats(requests: list[BatchRequest]) -> dict:
-    """
-    Compute statistics about a batch.
-
-    Args:
-        requests: List of BatchRequest
-
-    Returns:
-        Dict with batch statistics
-    """
-    by_type = split_by_type(requests)
-
-    return {
-        "total_requests": len(requests),
-        "num_decode": len(by_type["decode"]),
-        "num_prefill": len(by_type["prefill"]),
-        "num_extend": len(by_type["extend"]),
-        "total_tokens": sum(r.q_len for r in requests),
-        "total_kv_cache": sum(r.kv_len for r in requests),
-        "max_q_len": max((r.q_len for r in requests), default=0),
-        "max_kv_len": max((r.kv_len for r in requests), default=0),
-        "avg_q_len": sum(r.q_len for r in requests) / len(requests) if requests else 0,
-        "avg_kv_len": (
-            sum(r.kv_len for r in requests) / len(requests) if requests else 0
-        ),
-    }
-
-
-def get_batch_type(batch_spec: str, spec_decode_threshold: int = 8) -> str:
-    """
-    Classify a batch spec into a type string.
-
-    Args:
-        batch_spec: Batch specification string (e.g., "q2k", "8q1s1k", "2q2k_8q1s1k")
-        spec_decode_threshold: Max q_len to be considered spec-decode vs extend
-
-    Returns:
-        Type string: "prefill", "decode", "spec-decode", "extend", or "mixed (types...)"
-    """
-    requests = parse_batch_spec(batch_spec)
-
-    # Classify each request
-    types_present = set()
-    for req in requests:
-        if req.is_decode:
-            types_present.add("decode")
-        elif req.is_prefill:
-            types_present.add("prefill")
-        elif req.is_extend:
-            # Distinguish spec-decode (small q_len) from extend (chunked prefill)
-            if req.q_len <= spec_decode_threshold:
-                types_present.add("spec-decode")
-            else:
-                types_present.add("extend")
-
-    if len(types_present) == 1:
-        return types_present.pop()
-    elif len(types_present) > 1:
-        # Sort for consistent output
-        sorted_types = sorted(types_present)
-        return f"mixed ({'+'.join(sorted_types)})"
-    else:
-        return "unknown"

benchmarks/attention_benchmarks/benchmark.py

@@ -1,895 +0,0 @@
-#!/usr/bin/env python3
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""
-Universal vLLM Attention Benchmark
-
-Benchmark any attention backend with the extended grammar.
-Supports standard attention (Flash/Triton/FlashInfer) and MLA backends.
-
-Examples:
-    # Standard attention
-    python benchmark.py --backends flash flashinfer --batch-specs "q2k" "8q1s1k"
-
-    # MLA backends
-    python benchmark.py --backends cutlass_mla flashinfer_mla --batch-specs "64q1s1k"
-
-    # Parameter sweep (CLI)
-    python benchmark.py --backend cutlass_mla \
-                        --batch-specs "64q1s1k" \
-                        --sweep-param num_kv_splits \
-                        --sweep-values 1 4 8 16
-
-    # Parameter sweep (YAML config - recommended)
-    python benchmark.py --config configs/cutlass_numsplits.yaml
-"""
-
-import argparse
-import sys
-from dataclasses import replace
-from pathlib import Path
-
-import yaml
-from rich.console import Console
-from tqdm import tqdm
-
-sys.path.insert(0, str(Path(__file__).parent.parent.parent))
-
-from batch_spec import parse_batch_spec
-from common import (
-    BenchmarkConfig,
-    BenchmarkResult,
-    ModelParameterSweep,
-    ParameterSweep,
-    ResultsFormatter,
-    batch_spec_sort_key,
-    is_mla_backend,
-)
-
-
-def run_standard_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
-    """Run standard attention benchmark (Flash/Triton/FlashInfer)."""
-    from runner import run_attention_benchmark
-
-    return run_attention_benchmark(config)
-
-
-def run_mla_benchmark(config: BenchmarkConfig, **kwargs) -> BenchmarkResult:
-    """Run MLA benchmark with appropriate backend."""
-    from mla_runner import run_mla_benchmark as run_mla
-
-    return run_mla(config.backend, config, **kwargs)
-
-
-def run_benchmark(config: BenchmarkConfig, **kwargs) -> BenchmarkResult:
-    """
-    Run a single benchmark with proper backend selection.
-
-    Args:
-        config: BenchmarkConfig with backend, batch_spec, and model params
-        **kwargs: Additional arguments passed to MLA benchmarks
-
-    Returns:
-        BenchmarkResult (may have error field set on failure)
-    """
-    try:
-        if is_mla_backend(config.backend):
-            return run_mla_benchmark(config, **kwargs)
-        else:
-            return run_standard_attention_benchmark(config)
-    except Exception as e:
-        return BenchmarkResult(
-            config=config,
-            mean_time=float("inf"),
-            std_time=0,
-            min_time=float("inf"),
-            max_time=float("inf"),
-            error=str(e),
-        )
-
-
-def run_model_parameter_sweep(
-    backends: list[str],
-    batch_specs: list[str],
-    base_config_args: dict,
-    sweep: ModelParameterSweep,
-    console: Console,
-) -> list[BenchmarkResult]:
-    """
-    Run model parameter sweep for given backends and batch specs.
-
-    Args:
-        backends: List of backend names
-        batch_specs: List of batch specifications
-        base_config_args: Base configuration arguments (num_layers, head_dim, etc.)
-        sweep: ModelParameterSweep configuration
-        console: Rich console for output
-
-    Returns:
-        List of BenchmarkResult objects
-    """
-    all_results = []
-
-    console.print(
-        f"[yellow]Model sweep mode: testing {sweep.param_name} = {sweep.values}[/]"
-    )
-
-    total = len(backends) * len(batch_specs) * len(sweep.values)
-
-    with tqdm(total=total, desc="Benchmarking") as pbar:
-        for backend in backends:
-            for spec in batch_specs:
-                for value in sweep.values:
-                    # Create config with modified model parameter
-                    config_args = base_config_args.copy()
-                    config_args[sweep.param_name] = value
-
-                    # Create config with original backend for running
-                    clean_config = BenchmarkConfig(
-                        backend=backend, batch_spec=spec, **config_args
-                    )
-
-                    # Run benchmark
-                    result = run_benchmark(clean_config)
-
-                    # Replace backend with labeled version for display
-                    backend_label = sweep.get_label(backend, value)
-                    labeled_config = replace(result.config, backend=backend_label)
-                    result = replace(result, config=labeled_config)
-                    all_results.append(result)
-
-                    if not result.success:
-                        console.print(
-                            f"[red]Error {backend} {spec} {sweep.param_name}="
-                            f"{value}: {result.error}[/]"
-                        )
-
-                    pbar.update(1)
-
-    # Display sweep results - create separate table for each parameter value
-    console.print("\n[bold green]Model Parameter Sweep Results:[/]")
-    formatter = ResultsFormatter(console)
-
-    # Group results by parameter value and extract backend mapping
-    by_param_value = {}
-    backend_mapping = {}  # Maps labeled backend -> original backend
-
-    for r in all_results:
-        # Extract original backend and param value from labeled backend
-        # The label format is: {backend}_{param_name}_{value}
-        # We need to reverse engineer this
-        labeled_backend = r.config.backend
-
-        # Try each backend to find which one this result belongs to
-        for backend in backends:
-            for value in sweep.values:
-                expected_label = sweep.get_label(backend, value)
-                if labeled_backend == expected_label:
-                    backend_mapping[labeled_backend] = backend
-                    param_value = str(value)
-
-                    if param_value not in by_param_value:
-                        by_param_value[param_value] = []
-                    by_param_value[param_value].append(r)
-                    break
-
-    # Create a table for each parameter value
-    sorted_param_values = sorted(
-        by_param_value.keys(), key=lambda x: int(x) if x.isdigit() else x
-    )
-
-    for param_value in sorted_param_values:
-        console.print(f"\n[bold cyan]{sweep.param_name} = {param_value}[/]")
-        param_results = by_param_value[param_value]
-
-        # Create modified results with original backend names
-        modified_results = []
-        for r in param_results:
-            # Get the original backend name from our mapping
-            original_backend = backend_mapping[r.config.backend]
-            modified_config = replace(r.config, backend=original_backend)
-            modified_result = replace(r, config=modified_config)
-            modified_results.append(modified_result)
-
-        # Print table with original backend names
-        formatter.print_table(modified_results, backends, compare_to_fastest=True)
-
-    # Show optimal backend for each (param_value, batch_spec) combination
-    console.print(
-        f"\n[bold cyan]Optimal backend for each ({sweep.param_name}, batch_spec):[/]"
-    )
-
-    # Group by (param_value, batch_spec)
-    by_param_and_spec = {}
-    for r in all_results:
-        if r.success:
-            # Find which (backend, value) this result corresponds to
-            labeled_backend = r.config.backend
-            for backend in backends:
-                for value in sweep.values:
-                    expected_label = sweep.get_label(backend, value)
-                    if labeled_backend == expected_label:
-                        param_value = str(value)
-                        spec = r.config.batch_spec
-                        key = (param_value, spec)
-
-                        if key not in by_param_and_spec:
-                            by_param_and_spec[key] = []
-                        by_param_and_spec[key].append(r)
-                        break
-
-    # Sort by param value then spec (batch_size, q_len, kv_len)
-    sorted_keys = sorted(
-        by_param_and_spec.keys(),
-        key=lambda x: (
-            int(x[0]) if x[0].isdigit() else x[0],
-            batch_spec_sort_key(x[1]),
-        ),
-    )
-
-    current_param_value = None
-    for param_value, spec in sorted_keys:
-        # Print header when param value changes
-        if param_value != current_param_value:
-            console.print(f"\n  [bold]{sweep.param_name}={param_value}:[/]")
-            current_param_value = param_value
-
-        results = by_param_and_spec[(param_value, spec)]
-        best = min(results, key=lambda r: r.mean_time)
-
-        # Extract original backend name using the mapping
-        backend_name = backend_mapping[best.config.backend]
-
-        # Show all backends' times for comparison
-        times_str = " | ".join(
-            [
-                f"{backend_mapping[r.config.backend]}: {r.mean_time:.6f}s"
-                for r in sorted(results, key=lambda r: r.mean_time)
-            ]
-        )
-
-        console.print(
-            f"    {spec:12s} -> [bold green]{backend_name:15s}[/] ({times_str})"
-        )
-
-    return all_results
-
-
-def run_parameter_sweep(
-    backends: list[str],
-    batch_specs: list[str],
-    base_config_args: dict,
-    sweep: ParameterSweep,
-    console: Console,
-) -> list[BenchmarkResult]:
-    """
-    Run parameter sweep for given backends and batch specs.
-
-    Args:
-        backends: List of backend names
-        batch_specs: List of batch specifications
-        base_config_args: Base configuration arguments (num_layers, head_dim, etc.)
-        sweep: ParameterSweep configuration
-        console: Rich console for output
-
-    Returns:
-        List of BenchmarkResult objects
-    """
-    all_results = []
-
-    # Build list of values to sweep (including auto if requested)
-    sweep_values = list(sweep.values)
-    if sweep.include_auto:
-        sweep_values.append("auto")
-
-    console.print(f"[yellow]Sweep mode: testing {sweep.param_name} = {sweep_values}[/]")
-
-    total = len(backends) * len(batch_specs) * len(sweep_values)
-
-    with tqdm(total=total, desc="Benchmarking") as pbar:
-        for backend in backends:
-            for spec in batch_specs:
-                for value in sweep_values:
-                    # Create config with original backend for running
-                    config = BenchmarkConfig(
-                        backend=backend, batch_spec=spec, **base_config_args
-                    )
-
-                    # Prepare kwargs for benchmark runner
-                    kwargs = {}
-                    if value != "auto":
-                        kwargs[sweep.param_name] = value
-
-                    # Run benchmark
-                    result = run_benchmark(config, **kwargs)
-
-                    # Replace backend with labeled version for display
-                    backend_label = sweep.get_label(backend, value)
-                    labeled_config = replace(result.config, backend=backend_label)
-                    result = replace(result, config=labeled_config)
-                    all_results.append(result)
-
-                    if not result.success:
-                        console.print(
-                            f"[red]Error {backend} {spec} {sweep.param_name}="
-                            f"{value}: {result.error}[/]"
-                        )
-
-                    pbar.update(1)
-
-    # Display sweep results
-    console.print("\n[bold green]Sweep Results:[/]")
-    backend_labels = [sweep.get_label(b, v) for b in backends for v in sweep_values]
-    formatter = ResultsFormatter(console)
-    formatter.print_table(all_results, backend_labels)
-
-    # Show optimal values
-    console.print(f"\n[bold cyan]Optimal {sweep.param_name} per batch spec:[/]")
-    by_spec = {}
-    for r in all_results:
-        if r.success:
-            spec = r.config.batch_spec
-            if spec not in by_spec:
-                by_spec[spec] = []
-            by_spec[spec].append(r)
-
-    for spec in sorted(by_spec.keys(), key=batch_spec_sort_key):
-        results = by_spec[spec]
-        best = min(results, key=lambda r: r.mean_time)
-        console.print(
-            f"  {spec}: [bold green]{best.config.backend}[/] ({best.mean_time:.6f}s)"
-        )
-
-    return all_results
-
-
-def load_config_from_yaml(config_path: str) -> dict:
-    """Load configuration from YAML file."""
-    with open(config_path) as f:
-        return yaml.safe_load(f)
-
-
-def generate_batch_specs_from_ranges(ranges: list[dict]) -> list[str]:
-    """
-    Generate batch specs from range specifications.
-
-    Args:
-        ranges: List of range specifications, each containing:
-            - template: Batch spec template (e.g., "q{q_len}kv1k")
-            - q_len: Dict with start, stop, step, end_inclusive (optional)
-            - Other parameters can also be ranges
-
-    Returns:
-        List of generated batch spec strings
-
-    Example:
-        ranges = [
-            {
-                "template": "q{q_len}kv1k",
-                "q_len": {
-                    "start": 1,
-                    "stop": 16,
-                    "step": 1,
-                    "end_inclusive": true  # Optional, defaults to true
-                }
-            }
-        ]
-        Returns: ["q1kv1k", "q2kv1k", ..., "q16kv1k"]
-    """
-    all_specs = []
-
-    for range_spec in ranges:
-        template = range_spec.get("template")
-        if not template:
-            raise ValueError("Range specification must include 'template'")
-
-        # Extract all range parameters from the spec
-        range_params = {}
-        for key, value in range_spec.items():
-            if key == "template":
-                continue
-            if isinstance(value, dict) and "start" in value:
-                # This is a range specification
-                start = value["start"]
-                stop = value["stop"]
-                step = value.get("step", 1)
-                # Check if end should be inclusive (default: True)
-                end_inclusive = value.get("end_inclusive", True)
-
-                # Adjust stop based on end_inclusive
-                if end_inclusive:
-                    range_params[key] = list(range(start, stop + 1, step))
-                else:
-                    range_params[key] = list(range(start, stop, step))
-            else:
-                # This is a fixed value
-                range_params[key] = [value]
-
-        # Generate all combinations (Cartesian product)
-        if range_params:
-            import itertools
-
-            param_names = list(range_params.keys())
-            param_values = [range_params[name] for name in param_names]
-
-            for values in itertools.product(*param_values):
-                params = dict(zip(param_names, values))
-                spec = template.format(**params)
-                all_specs.append(spec)
-        else:
-            # No parameters, just use template as-is
-            all_specs.append(template)
-
-    return all_specs
-
-
-def main():
-    parser = argparse.ArgumentParser(
-        description="Universal vLLM attention benchmark",
-        formatter_class=argparse.RawDescriptionHelpFormatter,
-        epilog=__doc__,
-    )
-
-    # Config file
-    parser.add_argument(
-        "--config",
-        help="Path to YAML config file (overrides other args)",
-    )
-
-    # Backend selection
-    parser.add_argument(
-        "--backends",
-        nargs="+",
-        help="Backends to benchmark (flash, triton, flashinfer, cutlass_mla, "
-        "flashinfer_mla, flashattn_mla, flashmla)",
-    )
-    parser.add_argument(
-        "--backend",
-        help="Single backend (alternative to --backends)",
-    )
-
-    # Batch specifications
-    parser.add_argument(
-        "--batch-specs",
-        nargs="+",
-        default=["q2k", "8q1s1k"],
-        help="Batch specifications using extended grammar",
-    )
-
-    # Model config
-    parser.add_argument("--num-layers", type=int, default=10, help="Number of layers")
-    parser.add_argument("--head-dim", type=int, default=128, help="Head dimension")
-    parser.add_argument("--num-q-heads", type=int, default=32, help="Query heads")
-    parser.add_argument("--num-kv-heads", type=int, default=8, help="KV heads")
-    parser.add_argument("--block-size", type=int, default=16, help="Block size")
-
-    # Benchmark settings
-    parser.add_argument("--device", default="cuda:0", help="Device")
-    parser.add_argument("--repeats", type=int, default=1, help="Repetitions")
-    parser.add_argument("--warmup-iters", type=int, default=3, help="Warmup iterations")
-    parser.add_argument("--profile-memory", action="store_true", help="Profile memory")
-
-    # Parameter sweep (use YAML config for advanced sweeps)
-    parser.add_argument(
-        "--sweep-param",
-        help="Parameter name to sweep (e.g., num_kv_splits, reorder_batch_threshold)",
-    )
-    parser.add_argument(
-        "--sweep-values",
-        type=int,
-        nargs="+",
-        help="Values to sweep for the parameter",
-    )
-
-    # Output
-    parser.add_argument("--output-csv", help="Save to CSV")
-    parser.add_argument("--output-json", help="Save to JSON")
-
-    args = parser.parse_args()
-
-    console = Console()
-    console.print("[bold cyan]vLLM Attention Benchmark[/]")
-
-    # Load config from YAML if provided
-    if args.config:
-        console.print(f"[yellow]Loading config from: {args.config}[/]")
-        yaml_config = load_config_from_yaml(args.config)
-
-        # Show description if available
-        if "description" in yaml_config:
-            console.print(f"[dim]{yaml_config['description']}[/]")
-
-        # Override args with YAML values, but CLI args take precedence
-        # Check if CLI provided backends (they would be non-None and not default)
-        cli_backends_provided = args.backends is not None or args.backend is not None
-
-        # Backend(s) - only use YAML if CLI didn't specify
-        if not cli_backends_provided:
-            if "backend" in yaml_config:
-                args.backend = yaml_config["backend"]
-                args.backends = None
-            elif "backends" in yaml_config:
-                args.backends = yaml_config["backends"]
-                args.backend = None
-
-        # Check for special modes
-        if "mode" in yaml_config:
-            args.mode = yaml_config["mode"]
-        else:
-            args.mode = None
-
-        # Batch specs and sizes
-        # Support both explicit batch_specs and generated batch_spec_ranges
-        if "batch_spec_ranges" in yaml_config:
-            # Generate batch specs from ranges
-            generated_specs = generate_batch_specs_from_ranges(
-                yaml_config["batch_spec_ranges"]
-            )
-            # Combine with any explicit batch_specs
-            if "batch_specs" in yaml_config:
-                args.batch_specs = yaml_config["batch_specs"] + generated_specs
-            else:
-                args.batch_specs = generated_specs
-            console.print(
-                f"[dim]Generated {len(generated_specs)} batch specs from ranges[/]"
-            )
-        elif "batch_specs" in yaml_config:
-            args.batch_specs = yaml_config["batch_specs"]
-
-        if "batch_sizes" in yaml_config:
-            args.batch_sizes = yaml_config["batch_sizes"]
-        else:
-            args.batch_sizes = None
-
-        # Model config
-        if "model" in yaml_config:
-            model = yaml_config["model"]
-            args.num_layers = model.get("num_layers", args.num_layers)
-            args.head_dim = model.get("head_dim", args.head_dim)
-            args.num_q_heads = model.get("num_q_heads", args.num_q_heads)
-            args.num_kv_heads = model.get("num_kv_heads", args.num_kv_heads)
-            args.block_size = model.get("block_size", args.block_size)
-
-        # Benchmark settings (top-level keys)
-        if "device" in yaml_config:
-            args.device = yaml_config["device"]
-        if "repeats" in yaml_config:
-            args.repeats = yaml_config["repeats"]
-        if "warmup_iters" in yaml_config:
-            args.warmup_iters = yaml_config["warmup_iters"]
-        if "profile_memory" in yaml_config:
-            args.profile_memory = yaml_config["profile_memory"]
-
-        # Parameter sweep configuration
-        if "parameter_sweep" in yaml_config:
-            sweep_config = yaml_config["parameter_sweep"]
-            args.parameter_sweep = ParameterSweep(
-                param_name=sweep_config["param_name"],
-                values=sweep_config["values"],
-                include_auto=sweep_config.get("include_auto", False),
-                label_format=sweep_config.get(
-                    "label_format", "{backend}_{param_name}_{value}"
-                ),
-            )
-        else:
-            args.parameter_sweep = None
-
-        # Model parameter sweep configuration
-        if "model_parameter_sweep" in yaml_config:
-            sweep_config = yaml_config["model_parameter_sweep"]
-            args.model_parameter_sweep = ModelParameterSweep(
-                param_name=sweep_config["param_name"],
-                values=sweep_config["values"],
-                label_format=sweep_config.get(
-                    "label_format", "{backend}_{param_name}_{value}"
-                ),
-            )
-        else:
-            args.model_parameter_sweep = None
-
-        # Output
-        if "output" in yaml_config:
-            output = yaml_config["output"]
-            if "csv" in output and not args.output_csv:
-                args.output_csv = output["csv"]
-            if "json" in output and not args.output_json:
-                args.output_json = output["json"]
-
-        console.print()
-
-    # Handle CLI-based parameter sweep (if not from YAML)
-    if (
-        (not hasattr(args, "parameter_sweep") or args.parameter_sweep is None)
-        and args.sweep_param
-        and args.sweep_values
-    ):
-        args.parameter_sweep = ParameterSweep(
-            param_name=args.sweep_param,
-            values=args.sweep_values,
-            include_auto=False,
-            label_format="{backend}_{param_name}_{value}",
-        )
-
-    # Determine backends
-    backends = args.backends or ([args.backend] if args.backend else ["flash"])
-    console.print(f"Backends: {', '.join(backends)}")
-    console.print(f"Batch specs: {', '.join(args.batch_specs)}")
-    console.print()
-
-    # Run benchmarks
-    all_results = []
-
-    # Handle special mode: decode_vs_prefill comparison
-    if hasattr(args, "mode") and args.mode == "decode_vs_prefill":
-        console.print("[yellow]Mode: Decode vs Prefill pipeline comparison[/]")
-        console.print(
-            "[dim]For each query length, testing both decode and prefill pipelines[/]"
-        )
-        console.print("[dim]Using batched execution for optimal performance[/]")
-
-        # Extract batch sizes from config
-        batch_sizes = getattr(args, "batch_sizes", [1])
-        backend = backends[0]  # Use first backend (should only be one)
-
-        # Calculate total benchmarks
-        total = len(batch_sizes)
-
-        with tqdm(total=total, desc="Benchmarking") as pbar:
-            for batch_size in batch_sizes:
-                # Prepare all configs for this batch size
-                configs_with_thresholds = []
-
-                for spec in args.batch_specs:
-                    # Parse the batch spec to get query length
-                    requests = parse_batch_spec(spec)
-                    if not requests:
-                        console.print(
-                            f"[red]Error: Could not parse batch spec '{spec}'[/]"
-                        )
-                        continue
-
-                    # Get query length from first request
-                    query_length = requests[0].q_len
-
-                    # Create batch spec for this batch size
-                    # For batch_size > 1, we need to prepend the count
-                    batch_spec = f"{batch_size}{spec}" if batch_size > 1 else spec
-
-                    # Create base config (without backend name)
-                    base_config = BenchmarkConfig(
-                        backend=backend,  # Will be overridden later
-                        batch_spec=batch_spec,
-                        num_layers=args.num_layers,
-                        head_dim=args.head_dim,
-                        num_q_heads=args.num_q_heads,
-                        num_kv_heads=args.num_kv_heads,
-                        block_size=args.block_size,
-                        device=args.device,
-                        repeats=args.repeats,
-                        warmup_iters=args.warmup_iters,
-                        profile_memory=args.profile_memory,
-                    )
-
-                    # Add decode pipeline config
-                    decode_threshold = query_length
-                    config_decode = replace(
-                        base_config,
-                        backend=f"{backend}_decode_qlen{query_length}_bs{batch_size}",
-                    )
-                    configs_with_thresholds.append((config_decode, decode_threshold))
-
-                    # Add prefill pipeline config if query_length > 1
-                    if query_length > 1:
-                        prefill_threshold = query_length - 1
-                        config_prefill = replace(
-                            base_config,
-                            backend=f"{backend}_prefill_qlen{query_length}"
-                            f"_bs{batch_size}",
-                        )
-                        configs_with_thresholds.append(
-                            (config_prefill, prefill_threshold)
-                        )
-
-                # Run all benchmarks for this batch size in one go (batched mode)
-                try:
-                    from mla_runner import run_mla_benchmark as run_mla
-
-                    # Use batched API: pass list of (config, threshold) tuples
-                    timing_results = run_mla(backend, configs_with_thresholds)
-
-                    # Create BenchmarkResult objects from timing results
-                    for (config, _), timing in zip(
-                        configs_with_thresholds, timing_results
-                    ):
-                        result = BenchmarkResult(
-                            config=config,
-                            mean_time=timing["mean"],
-                            std_time=timing["std"],
-                            min_time=timing["min"],
-                            max_time=timing["max"],
-                            throughput_tokens_per_sec=timing.get("throughput", None),
-                        )
-                        all_results.append(result)
-
-                except Exception as e:
-                    import traceback
-
-                    console.print(
-                        f"[red]Error running batched benchmarks for "
-                        f"batch_size={batch_size}: {e}[/]"
-                    )
-                    console.print("[red]Traceback:[/]")
-                    traceback.print_exc()
-                    # Add error results for all configs
-                    for config, _ in configs_with_thresholds:
-                        result = BenchmarkResult(
-                            config=config,
-                            mean_time=float("inf"),
-                            std_time=0,
-                            min_time=float("inf"),
-                            max_time=float("inf"),
-                            error=str(e),
-                        )
-                        all_results.append(result)
-
-                pbar.update(1)
-
-        # Display decode vs prefill results
-        console.print("\n[bold green]Decode vs Prefill Results:[/]")
-
-        # Group by batch size
-        by_batch_size = {}
-        for r in all_results:
-            if r.success:
-                # Extract batch size from backend name
-                parts = r.config.backend.split("_")
-                bs_part = [p for p in parts if p.startswith("bs")]
-                if bs_part:
-                    bs = int(bs_part[0][2:])
-                    if bs not in by_batch_size:
-                        by_batch_size[bs] = []
-                    by_batch_size[bs].append(r)
-
-        # For each batch size, analyze crossover point
-        for bs in sorted(by_batch_size.keys()):
-            console.print(f"\n[bold cyan]Batch size: {bs}[/]")
-            results = by_batch_size[bs]
-
-            # Group by query length
-            by_qlen = {}
-            for r in results:
-                parts = r.config.backend.split("_")
-                qlen_part = [p for p in parts if p.startswith("qlen")]
-                if qlen_part:
-                    qlen = int(qlen_part[0][4:])
-                    if qlen not in by_qlen:
-                        by_qlen[qlen] = {}
-
-                    pipeline = "decode" if "decode" in r.config.backend else "prefill"
-                    by_qlen[qlen][pipeline] = r
-
-            # Find crossover point
-            last_decode_faster = None
-            for qlen in sorted(by_qlen.keys()):
-                pipelines = by_qlen[qlen]
-                if "decode" in pipelines and "prefill" in pipelines:
-                    decode_time = pipelines["decode"].mean_time
-                    prefill_time = pipelines["prefill"].mean_time
-                    faster = "decode" if decode_time < prefill_time else "prefill"
-
-                    speedup = (
-                        prefill_time / decode_time
-                        if decode_time < prefill_time
-                        else decode_time / prefill_time
-                    )
-
-                    console.print(
-                        f"  qlen={qlen:3d}: decode={decode_time:.6f}s, "
-                        f"prefill={prefill_time:.6f}s -> "
-                        f"[bold]{faster}[/] ({speedup:.2f}x)"
-                    )
-
-                    if faster == "decode":
-                        last_decode_faster = qlen
-
-            if last_decode_faster is not None:
-                optimal_threshold = last_decode_faster
-                console.print(
-                    f"\n  [bold green]Optimal threshold for batch_size={bs}: "
-                    f"{optimal_threshold}[/]"
-                )
-                console.print(
-                    f"  [dim](Use decode pipeline for query_length <= "
-                    f"{optimal_threshold})[/]"
-                )
-            else:
-                console.print(
-                    f"\n  [yellow]Prefill always faster for batch_size={bs}[/]"
-                )
-
-    # Handle model parameter sweep mode
-    elif hasattr(args, "model_parameter_sweep") and args.model_parameter_sweep:
-        # Model parameter sweep
-        base_config_args = {
-            "num_layers": args.num_layers,
-            "head_dim": args.head_dim,
-            "num_q_heads": args.num_q_heads,
-            "num_kv_heads": args.num_kv_heads,
-            "block_size": args.block_size,
-            "device": args.device,
-            "repeats": args.repeats,
-            "warmup_iters": args.warmup_iters,
-            "profile_memory": args.profile_memory,
-        }
-        all_results = run_model_parameter_sweep(
-            backends,
-            args.batch_specs,
-            base_config_args,
-            args.model_parameter_sweep,
-            console,
-        )
-
-    # Handle parameter sweep mode (unified)
-    elif hasattr(args, "parameter_sweep") and args.parameter_sweep:
-        # Unified parameter sweep
-        base_config_args = {
-            "num_layers": args.num_layers,
-            "head_dim": args.head_dim,
-            "num_q_heads": args.num_q_heads,
-            "num_kv_heads": args.num_kv_heads,
-            "block_size": args.block_size,
-            "device": args.device,
-            "repeats": args.repeats,
-            "warmup_iters": args.warmup_iters,
-            "profile_memory": args.profile_memory,
-        }
-        all_results = run_parameter_sweep(
-            backends, args.batch_specs, base_config_args, args.parameter_sweep, console
-        )
-
-    else:
-        # Normal mode: compare backends
-        total = len(backends) * len(args.batch_specs)
-
-        with tqdm(total=total, desc="Benchmarking") as pbar:
-            for spec in args.batch_specs:
-                for backend in backends:
-                    config = BenchmarkConfig(
-                        backend=backend,
-                        batch_spec=spec,
-                        num_layers=args.num_layers,
-                        head_dim=args.head_dim,
-                        num_q_heads=args.num_q_heads,
-                        num_kv_heads=args.num_kv_heads,
-                        block_size=args.block_size,
-                        device=args.device,
-                        repeats=args.repeats,
-                        warmup_iters=args.warmup_iters,
-                        profile_memory=args.profile_memory,
-                    )
-
-                    result = run_benchmark(config)
-                    all_results.append(result)
-
-                    if not result.success:
-                        console.print(f"[red]Error {backend} {spec}: {result.error}[/]")
-
-                    pbar.update(1)
-
-        # Display results
-        console.print("\n[bold green]Results:[/]")
-        formatter = ResultsFormatter(console)
-        formatter.print_table(all_results, backends)
-
-    # Save results
-    if all_results:
-        formatter = ResultsFormatter(console)
-        if args.output_csv:
-            formatter.save_csv(all_results, args.output_csv)
-        if args.output_json:
-            formatter.save_json(all_results, args.output_json)
-
-
-if __name__ == "__main__":
-    main()

benchmarks/attention_benchmarks/common.py

@@ -1,568 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""Common utilities for attention benchmarking."""
-
-import csv
-import json
-import math
-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
-from rich.table import Table
-
-
-def batch_spec_sort_key(spec: str) -> tuple[int, int, int]:
-    """
-    Extract sorting key from batch spec: (batch_size, max_q_len, max_kv_len).
-
-    This ensures results are sorted by batch size first, then query length,
-    then sequence length, rather than alphabetically.
-    """
-    try:
-        requests = parse_batch_spec(spec)
-        batch_size = len(requests)
-        max_q_len = max(r.q_len for r in requests) if requests else 0
-        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 unparseable specs
-        return (0, 0, 0)
-
-
-# Mock classes for vLLM attention infrastructure
-
-
-class MockHfConfig:
-    """Mock HuggingFace config that satisfies vLLM's requirements."""
-
-    def __init__(self, mla_dims: dict, index_topk: int | None = None):
-        self.num_attention_heads = mla_dims["num_q_heads"]
-        self.num_key_value_heads = mla_dims["num_kv_heads"]
-        self.hidden_size = mla_dims["head_dim"] * mla_dims["num_q_heads"]
-        self.model_type = "deepseek_v2"
-        self.is_encoder_decoder = False
-        self.kv_lora_rank = mla_dims["kv_lora_rank"]
-        self.qk_nope_head_dim = mla_dims["qk_nope_head_dim"]
-        self.qk_rope_head_dim = mla_dims["qk_rope_head_dim"]
-        self.v_head_dim = mla_dims["v_head_dim"]
-        self.qk_head_dim = mla_dims["qk_nope_head_dim"] + mla_dims["qk_rope_head_dim"]
-        if index_topk is not None:
-            self.index_topk = index_topk
-
-    def get_text_config(self):
-        return self
-
-
-# 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
-
-
-class MockKVBProj:
-    """Mock KV projection layer for MLA prefill mode.
-
-    Mimics ColumnParallelLinear behavior for kv_b_proj in MLA backends.
-    Projects kv_c_normed to [qk_nope_head_dim + v_head_dim] per head.
-    """
-
-    def __init__(self, num_heads: int, qk_nope_head_dim: int, v_head_dim: int):
-        self.num_heads = num_heads
-        self.qk_nope_head_dim = qk_nope_head_dim
-        self.v_head_dim = v_head_dim
-        self.out_dim = qk_nope_head_dim + v_head_dim
-
-    def __call__(self, x: torch.Tensor) -> tuple[torch.Tensor]:
-        """
-        Project kv_c_normed to output space.
-
-        Args:
-            x: Input tensor [num_tokens, kv_lora_rank]
-
-        Returns:
-            Tuple containing output tensor
-                [num_tokens, num_heads, qk_nope_head_dim + v_head_dim]
-        """
-        num_tokens = x.shape[0]
-        result = torch.randn(
-            num_tokens,
-            self.num_heads,
-            self.out_dim,
-            device=x.device,
-            dtype=x.dtype,
-        )
-        return (result,)  # Return as tuple to match ColumnParallelLinear API
-
-
-class MockIndexer:
-    """Mock Indexer for sparse MLA backends.
-
-    Provides topk_indices_buffer that sparse MLA backends use to determine
-    which KV cache slots to attend to for each token.
-    """
-
-    def __init__(
-        self,
-        max_num_tokens: int,
-        topk_tokens: int,
-        device: torch.device,
-    ):
-        self.topk_tokens = topk_tokens
-        self.topk_indices_buffer = torch.zeros(
-            (max_num_tokens, topk_tokens),
-            dtype=torch.int32,
-            device=device,
-        )
-
-    def fill_random_indices(self, num_tokens: int, max_kv_len: int):
-        """Fill topk_indices_buffer with random valid indices for benchmarking."""
-        indices = torch.randint(
-            0,
-            max_kv_len,
-            (num_tokens, self.topk_tokens),
-            dtype=torch.int32,
-            device=self.topk_indices_buffer.device,
-        )
-        self.topk_indices_buffer[:num_tokens] = indices
-
-
-class MockLayer(AttentionLayerBase):
-    """Mock attention layer with scale parameters and impl.
-
-    Inherits from AttentionLayerBase so it passes isinstance checks
-    in get_layers_from_vllm_config when FlashInfer prefill is enabled.
-    """
-
-    def __init__(self, device: torch.device, impl=None, kv_cache_spec=None):
-        # Don't call super().__init__() as AttentionLayerBase doesn't have __init__
-        self._k_scale = torch.tensor(1.0, device=device)
-        self._v_scale = torch.tensor(1.0, device=device)
-        self._q_scale = torch.tensor(1.0, device=device)
-        # Scalar floats for kernels that need them
-        self._k_scale_float = float(self._k_scale.item())
-        self._v_scale_float = float(self._v_scale.item())
-        self._q_scale_float = float(self._q_scale.item())
-        # AttentionImpl for metadata builders to query
-        self.impl = impl
-        # KV cache spec for get_kv_cache_spec
-        self._kv_cache_spec = kv_cache_spec
-
-    def get_attn_backend(self):
-        """Get the attention backend class (required by AttentionLayerBase)."""
-        # Return None as this is just a mock layer for benchmarking
-        return None
-
-    def get_kv_cache_spec(self):
-        """Get the KV cache spec (required by 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."""
-
-    param_name: str  # Name of the backend parameter to sweep
-    values: list[Any]  # List of values to test
-    include_auto: bool = False  # Also test with param unset (auto mode)
-    label_format: str = "{backend}_{param_name}_{value}"  # Result label template
-
-    def get_label(self, backend: str, value: Any) -> str:
-        """Generate a label for a specific parameter value."""
-        return self.label_format.format(
-            backend=backend, param_name=self.param_name, value=value
-        )
-
-
-@dataclass
-class ModelParameterSweep:
-    """Configuration for sweeping a model configuration parameter."""
-
-    param_name: str  # Name of the model config parameter to sweep (e.g., "num_q_heads")
-    values: list[Any]  # List of values to test
-    label_format: str = "{backend}_{param_name}_{value}"  # Result label template
-
-    def get_label(self, backend: str, value: Any) -> str:
-        """Generate a label for a specific parameter value."""
-        return self.label_format.format(
-            backend=backend, param_name=self.param_name, value=value
-        )
-
-
-@dataclass
-class BenchmarkConfig:
-    """Configuration for a single benchmark run."""
-
-    backend: str
-    batch_spec: str
-    num_layers: int
-    head_dim: int
-    num_q_heads: int
-    num_kv_heads: int
-    block_size: int
-    device: str
-    dtype: torch.dtype = torch.float16
-    repeats: int = 1
-    warmup_iters: int = 3
-    profile_memory: bool = False
-    use_cuda_graphs: bool = False
-
-    # MLA-specific
-    kv_lora_rank: int | None = None
-    qk_nope_head_dim: int | None = None
-    qk_rope_head_dim: int | None = None
-    v_head_dim: int | None = None
-
-    # Backend-specific tuning
-    num_kv_splits: int | None = None  # CUTLASS MLA
-    reorder_batch_threshold: int | None = None  # FlashAttn MLA, FlashMLA
-
-
-@dataclass
-class BenchmarkResult:
-    """Results from a single benchmark run."""
-
-    config: BenchmarkConfig
-    mean_time: float  # seconds
-    std_time: float  # seconds
-    min_time: float  # seconds
-    max_time: float  # seconds
-    throughput_tokens_per_sec: float | None = None
-    memory_allocated_mb: float | None = None
-    memory_reserved_mb: float | None = None
-    error: str | None = None
-
-    @property
-    def success(self) -> bool:
-        """Whether benchmark completed successfully."""
-        return self.error is None
-
-    def to_dict(self) -> dict[str, Any]:
-        """Convert to dictionary for serialization."""
-        return {
-            "config": asdict(self.config),
-            "mean_time": self.mean_time,
-            "std_time": self.std_time,
-            "min_time": self.min_time,
-            "max_time": self.max_time,
-            "throughput_tokens_per_sec": self.throughput_tokens_per_sec,
-            "memory_allocated_mb": self.memory_allocated_mb,
-            "memory_reserved_mb": self.memory_reserved_mb,
-            "error": self.error,
-        }
-
-
-class ResultsFormatter:
-    """Format and display benchmark results."""
-
-    def __init__(self, console: Console | None = None):
-        self.console = console or Console()
-
-    def print_table(
-        self,
-        results: list[BenchmarkResult],
-        backends: list[str],
-        compare_to_fastest: bool = True,
-    ):
-        """
-        Print results as a rich table.
-
-        Args:
-            results: List of BenchmarkResult
-            backends: List of backend names being compared
-            compare_to_fastest: Show percentage comparison to fastest
-        """
-        # Group by batch spec, preserving first-occurrence order
-        by_spec = {}
-        specs_order = []
-        for r in results:
-            spec = r.config.batch_spec
-            if spec not in by_spec:
-                by_spec[spec] = {}
-                specs_order.append(spec)
-            by_spec[spec][r.config.backend] = r
-
-        # Sort specs by (batch_size, q_len, kv_len) instead of alphabetically
-        specs_order = sorted(by_spec.keys(), key=batch_spec_sort_key)
-
-        # Create shortened backend names for display
-        def shorten_backend_name(name: str) -> str:
-            """Shorten long backend names for table display."""
-            # Remove common prefixes
-            name = name.replace("flashattn_mla", "famla")
-            name = name.replace("flashinfer_mla", "fimla")
-            name = name.replace("flashmla", "fmla")
-            name = name.replace("cutlass_mla", "cmla")
-            name = name.replace("numsplits", "ns")
-            return name
-
-        table = Table(title="Attention Benchmark Results")
-        table.add_column("Batch\nSpec", no_wrap=True)
-        table.add_column("Type", no_wrap=True)
-        table.add_column("Batch\nSize", justify="right", no_wrap=True)
-
-        multi = len(backends) > 1
-        for backend in backends:
-            short_name = shorten_backend_name(backend)
-            # Time column
-            col_time = f"{short_name}\nTime (s)"
-            table.add_column(col_time, justify="right", no_wrap=False)
-            if multi and compare_to_fastest:
-                # Relative performance column
-                col_rel = f"{short_name}\nvs Best"
-                table.add_column(col_rel, justify="right", no_wrap=False)
-
-        # Add rows
-        for spec in specs_order:
-            spec_results = by_spec[spec]
-            times = {b: r.mean_time for b, r in spec_results.items() if r.success}
-            best_time = min(times.values()) if times else 0.0
-
-            batch_type = get_batch_type(spec)
-            batch_size = len(parse_batch_spec(spec))
-            row = [spec, batch_type, str(batch_size)]
-            for backend in backends:
-                if backend in spec_results:
-                    r = spec_results[backend]
-                    if r.success:
-                        row.append(f"{r.mean_time:.6f}")
-                        if multi and compare_to_fastest:
-                            pct = (
-                                (r.mean_time / best_time * 100) if best_time > 0 else 0
-                            )
-                            pct_str = f"{pct:.1f}%"
-                            if r.mean_time == best_time:
-                                pct_str = f"[bold green]{pct_str}[/]"
-                            row.append(pct_str)
-                    else:
-                        row.append("[red]ERROR[/]")
-                        if multi and compare_to_fastest:
-                            row.append("-")
-                else:
-                    row.append("-")
-                    if multi and compare_to_fastest:
-                        row.append("-")
-
-            table.add_row(*row)
-
-        self.console.print(table)
-
-    def save_csv(self, results: list[BenchmarkResult], path: str):
-        """Save results to CSV file."""
-        if not results:
-            return
-
-        path_obj = Path(path)
-        path_obj.parent.mkdir(parents=True, exist_ok=True)
-
-        with open(path, "w", newline="") as f:
-            writer = csv.DictWriter(
-                f,
-                fieldnames=[
-                    "backend",
-                    "batch_spec",
-                    "num_layers",
-                    "mean_time",
-                    "std_time",
-                    "throughput",
-                    "memory_mb",
-                ],
-            )
-            writer.writeheader()
-            for r in results:
-                writer.writerow(
-                    {
-                        "backend": r.config.backend,
-                        "batch_spec": r.config.batch_spec,
-                        "num_layers": r.config.num_layers,
-                        "mean_time": r.mean_time,
-                        "std_time": r.std_time,
-                        "throughput": r.throughput_tokens_per_sec or 0,
-                        "memory_mb": r.memory_allocated_mb or 0,
-                    }
-                )
-
-        self.console.print(f"[green]Saved CSV results to {path}[/]")
-
-    def save_json(self, results: list[BenchmarkResult], path: str):
-        """Save results to JSON file."""
-        path_obj = Path(path)
-        path_obj.parent.mkdir(parents=True, exist_ok=True)
-
-        data = [r.to_dict() for r in results]
-        with open(path, "w") as f:
-            json.dump(data, f, indent=2, default=str)
-
-        self.console.print(f"[green]Saved JSON results to {path}[/]")
-
-
-def setup_mla_dims(model_name: str = "deepseek-v3") -> dict:
-    """
-    Get MLA dimensions for known models.
-
-    Args:
-        model_name: Model identifier
-
-    Returns:
-        Dict with MLA dimension configuration
-    """
-    configs = {
-        "deepseek-v2": {
-            "kv_lora_rank": 512,
-            "qk_nope_head_dim": 128,
-            "qk_rope_head_dim": 64,
-            "v_head_dim": 128,
-            "num_q_heads": 128,
-            "num_kv_heads": 1,
-            "head_dim": 576,
-        },
-        "deepseek-v3": {
-            "kv_lora_rank": 512,
-            "qk_nope_head_dim": 128,
-            "qk_rope_head_dim": 64,
-            "v_head_dim": 128,
-            "num_q_heads": 128,
-            "num_kv_heads": 1,
-            "head_dim": 576,
-        },
-        "deepseek-v2-lite": {
-            "kv_lora_rank": 512,
-            "qk_nope_head_dim": 128,
-            "qk_rope_head_dim": 64,
-            "v_head_dim": 128,
-            "num_q_heads": 16,
-            "num_kv_heads": 1,
-            "head_dim": 576,
-        },
-    }
-
-    if model_name not in configs:
-        raise ValueError(
-            f"Unknown model '{model_name}'. Known models: {list(configs.keys())}"
-        )
-
-    return configs[model_name]
-
-
-def get_attention_scale(head_dim: int) -> float:
-    """Compute attention scale factor (1/sqrt(d))."""
-    return 1.0 / math.sqrt(head_dim)
-
-
-def is_mla_backend(backend: str) -> bool:
-    """
-    Check if backend is an MLA backend using the AttentionBackendEnum.
-
-    Args:
-        backend: Backend name matching AttentionBackendEnum exactly
-        (e.g., "FLASHMLA_SPARSE")
-
-    Returns:
-        True if the backend is an MLA backend, False otherwise
-    """
-    from vllm.v1.attention.backends.registry import AttentionBackendEnum
-
-    try:
-        backend_enum = AttentionBackendEnum[backend]
-        backend_class = backend_enum.get_class()
-        return backend_class.is_mla()
-    except (KeyError, ValueError, ImportError, AttributeError):
-        return False

benchmarks/attention_benchmarks/configs/mla_decode.yaml

@@ -1,70 +0,0 @@
-# MLA decode-only benchmark configuration
-
-model:
-  name: "deepseek-v3"
-  num_layers: 60
-  num_q_heads: 128  # Base value, can be swept for TP simulation
-  num_kv_heads: 1  # MLA uses single latent KV
-  head_dim: 576
-  kv_lora_rank: 512
-  qk_nope_head_dim: 128
-  qk_rope_head_dim: 64
-  v_head_dim: 128
-  block_size: 128  # CUTLASS MLA and FlashAttn MLA use 128
-
-# Model parameter sweep: simulate tensor parallelism by varying num_q_heads
-# TP=1: 128 heads, TP=2: 64 heads, TP=4: 32 heads, TP=8: 16 heads
-model_parameter_sweep:
-  param_name: "num_q_heads"
-  values: [128, 64, 32, 16]
-  label_format: "{backend}_{value}h"
-
-batch_specs:
-  # Small batches, varying sequence lengths
-  - "16q1s512"     # 16 requests, 512 KV cache
-  - "16q1s1k"      # 16 requests, 1k KV cache
-  - "16q1s2k"      # 16 requests, 2k KV cache
-  - "16q1s4k"      # 16 requests, 4k KV cache
-
-  # Medium batches
-  - "32q1s1k"      # 32 requests, 1k KV cache
-  - "32q1s2k"      # 32 requests, 2k KV cache
-  - "32q1s4k"      # 32 requests, 4k KV cache
-  - "32q1s8k"      # 32 requests, 8k KV cache
-
-  # Large batches
-  - "64q1s1k"      # 64 requests, 1k KV cache
-  - "64q1s2k"      # 64 requests, 2k KV cache
-  - "64q1s4k"      # 64 requests, 4k KV cache
-  - "64q1s8k"      # 64 requests, 8k KV cache
-
-  # Very large batches
-  - "128q1s1k"     # 128 requests, 1k KV cache
-  - "128q1s2k"     # 128 requests, 2k KV cache
-  - "128q1s4k"     # 128 requests, 4k KV cache
-  - "128q1s8k"     # 128 requests, 8k KV cache
-
-  # Long context
-  - "32q1s16k"     # 32 requests, 16k KV cache
-  - "32q1s32k"     # 32 requests, 32k KV cache
-
-backends:
-  - CUTLASS_MLA
-  - FLASHINFER_MLA
-  - FLASH_ATTN_MLA  # Hopper only
-  - FLASHMLA        # Hopper only
-
-device: "cuda:0"
-repeats: 100
-warmup_iters: 10
-profile_memory: true
-
-# Backend-specific tuning
-CUTLASS_MLA:
-  num_kv_splits: auto  # or specific value like 4, 8, 16
-
-FLASH_ATTN_MLA:
-  reorder_batch_threshold: 512
-
-FLASHMLA:
-  reorder_batch_threshold: 1

benchmarks/attention_benchmarks/configs/mla_mixed_batch.yaml

@@ -1,60 +0,0 @@
-# MLA mixed batch benchmark (prefill + decode)
-# Tests chunked prefill performance
-
-model:
-  name: "deepseek-v3"
-  num_layers: 60
-  num_q_heads: 128
-  num_kv_heads: 1
-  head_dim: 576
-  kv_lora_rank: 512
-  qk_nope_head_dim: 128
-  qk_rope_head_dim: 64
-  v_head_dim: 128
-  block_size: 128
-
-batch_specs:
-  # Small prefill + decode
-  - "1q1k_8q1s1k"           # 1 prefill + 8 decode
-  - "2q2k_16q1s1k"          # 2 prefill + 16 decode
-  - "4q1k_32q1s2k"          # 4 prefill + 32 decode
-
-  # Medium prefill + decode
-  - "2q4k_32q1s2k"          # 2 medium prefill + 32 decode
-  - "4q4k_64q1s2k"          # 4 medium prefill + 64 decode
-  - "8q2k_64q1s4k"          # 8 prefill + 64 decode
-
-  # Large prefill + decode (chunked prefill stress test)
-  - "2q8k_32q1s1k"          # 2 large prefill + 32 decode
-  - "1q16k_16q1s2k"         # 1 very large prefill + 16 decode
-  - "2q16k_32q1s4k"         # 2 very large prefill + 32 decode
-
-  # Context extension + decode
-  - "2q1kkv2k_16q1s1k"       # 2 extend + 16 decode
-  - "4q2kkv4k_32q1s2k"       # 4 extend + 32 decode
-  - "2q1kkv8k_32q1s2k"       # 2 large extend + 32 decode
-
-  # Explicitly chunked prefill
-  - "q8k"           # 8k prefill with chunking hint
-  - "q16k"          # 16k prefill with chunking hint
-  - "2q8k_32q1s2k"    # 2 chunked prefill + 32 decode
-
-  # High decode ratio (realistic serving)
-  - "1q2k_63q1s1k"          # 1 prefill + 63 decode
-  - "2q2k_62q1s2k"          # 2 prefill + 62 decode
-  - "4q4k_60q1s4k"          # 4 prefill + 60 decode
-
-backends:
-  - CUTLASS_MLA
-  - FLASHINFER_MLA
-  - FLASH_ATTN_MLA   # Hopper only
-  - FLASHMLA         # Hopper only
-
-device: "cuda:0"
-repeats: 5
-warmup_iters: 3
-profile_memory: true
-
-# Analyze chunked prefill workspace size impact
-chunked_prefill:
-  test_workspace_sizes: [4096, 8192, 16384, 32768, 65536]

benchmarks/attention_benchmarks/configs/mla_prefill.yaml

@@ -1,62 +0,0 @@
-# MLA prefill-only benchmark configuration for sparse backends
-
-model:
-  name: "deepseek-v3"
-  num_layers: 60
-  num_q_heads: 128
-  num_kv_heads: 1
-  head_dim: 576
-  kv_lora_rank: 512
-  qk_nope_head_dim: 128
-  qk_rope_head_dim: 64
-  v_head_dim: 128
-  block_size: 128
-
-# Model parameter sweep: simulate tensor parallelism by varying num_q_heads
-# TP=1: 128 heads, TP=2: 64 heads, TP=4: 32 heads, TP=8: 16 heads
-model_parameter_sweep:
-  param_name: "num_q_heads"
-  values: [128, 64, 32, 16]
-  label_format: "{backend}_{value}h"
-
-batch_specs:
-  # Pure prefill
-  - "1q512"
-  - "1q1k"
-  - "1q2k"
-  - "1q4k"
-  - "1q8k"
-
-  # Batched pure prefill
-  - "2q512"
-  - "2q1k"
-  - "2q2k"
-  - "2q4k"
-  - "2q8k"
-  - "4q512"
-  - "4q1k"
-  - "4q2k"
-  - "4q4k"
-  - "4q8k"
-  - "8q512"
-  - "8q1k"
-  - "8q2k"
-  - "8q4k"
-  - "8q8k"
-
-  # Extend
-  - "1q512s4k"
-  - "1q512s8k"
-  - "1q1ks8k"
-  - "1q2ks8k"
-  - "1q2ks16k"
-  - "1q4ks16k"
-
-backends:
-  - FLASHMLA_SPARSE
-  - FLASHINFER_MLA_SPARSE
-
-device: "cuda:0"
-repeats: 10
-warmup_iters: 3
-profile_memory: true

benchmarks/attention_benchmarks/configs/reorder_threshold.yaml

@@ -1,87 +0,0 @@
-# Study 4: What is optimal reorder_batch_threshold for MLA backends supporting query length > 1?
-# Question: At what query length does prefill pipeline become faster than decode pipeline?
-# Methodology: For each query length, compare decode vs prefill performance to find crossover point
-# Applies to: FlashAttn MLA, FlashMLA
-
-description: "Decode vs Prefill pipeline crossover analysis"
-
-# Test FlashAttn MLA
-backend: FLASH_ATTN_MLA
-
-# Mode: decode_vs_prefill comparison (special sweep mode)
-# For each batch spec, we'll test both decode and prefill pipelines
-mode: "decode_vs_prefill"
-
-# Query lengths to test (from old benchmark_mla_threshold.py methodology)
-# Each query length will be tested with BOTH decode and prefill pipelines:
-#   - decode: threshold >= query_length (forces decode pipeline)
-#   - prefill: threshold < query_length (forces prefill pipeline)
-#
-# We use q<N>s1k format which creates q_len=N, seq_len=1024 requests
-# This tests different query lengths with fixed sequence length context
-#
-# Using batch_spec_ranges for automatic generation:
-batch_spec_ranges:
-  - template: "q{q_len}s1k"
-    q_len:
-      start: 1
-      stop: 16
-      step: 1
-      end_inclusive: false
-  - template: "q{q_len}s1k"
-    q_len:
-      start: 16
-      stop: 64
-      step: 2
-      end_inclusive: false
-  - template: "q{q_len}s1k"
-    q_len:
-      start: 64
-      stop: 1024
-      step: 4
-      end_inclusive: true
-
-# Batch sizes to test (from old script)
-batch_sizes:
-  - 1
-  - 2
-  - 4
-  - 8
-  - 16
-  - 32
-  - 64
-  - 128
-  - 256
-
-# Model configuration (DeepSeek V2/V3 defaults)
-model:
-  num_layers: 10
-  head_dim: 576
-  num_q_heads: 128
-  num_kv_heads: 1
-  block_size: 128
-
-# Benchmark settings
-device: "cuda:0"
-repeats: 15          # More repeats for spec decode variance
-warmup_iters: 5
-profile_memory: false
-
-# Output
-output:
-  csv: "reorder_threshold_results.csv"
-  json: "reorder_threshold_results.json"
-
-# Expected outcome (reproduces old benchmark_mla_threshold.py study):
-# - For each batch size, find the crossover point where prefill becomes faster than decode
-# - Show decode vs prefill performance across all query lengths
-# - Determine optimal reorder_batch_threshold based on last query length where decode is faster
-# - Understand how crossover point varies with batch size
-# - Provide data-driven guidance for default threshold value
-#
-# Methodology (from old script):
-# - Each query length tested with BOTH pipelines:
-#     * decode: threshold >= query_length (forces decode pipeline)
-#     * prefill: threshold < query_length (forces prefill pipeline)
-# - Compare which is faster to find crossover point
-#

benchmarks/attention_benchmarks/configs/speculative_decode.yaml

@@ -1,61 +0,0 @@
-# Speculative decoding benchmark configuration
-# Tests reorder_batch_threshold optimization
-
-model:
-  name: "deepseek-v3"
-  num_layers: 60
-  num_q_heads: 128
-  num_kv_heads: 1
-  head_dim: 576
-  kv_lora_rank: 512
-  qk_nope_head_dim: 128
-  qk_rope_head_dim: 64
-  v_head_dim: 128
-
-batch_specs:
-  # Pure speculative decode (K-token verification)
-  - "q2s1k"      # 2-token spec, 1k KV
-  - "q4s1k"      # 4-token spec, 1k KV
-  - "q8s1k"      # 8-token spec, 1k KV
-  - "q16s1k"     # 16-token spec, 1k KV
-
-  # Speculative with different context lengths
-  - "q4s2k"      # 4-token spec, 2k KV
-  - "q4s4k"      # 4-token spec, 4k KV
-  - "q8s2k"      # 8-token spec, 2k KV
-  - "q8s4k"      # 8-token spec, 4k KV
-
-  # Mixed: speculative + regular decode
-  - "32q4s1k"                    # 32 spec requests
-  - "16q4s1k_16q1s1k"              # 16 spec + 16 regular
-  - "8q8s2k_24q1s2k"               # 8 spec (8-tok) + 24 regular
-
-  # Mixed: speculative + prefill + decode
-  - "2q1k_16q4s1k_16q1s1k"         # 2 prefill + 16 spec + 16 decode
-  - "4q2k_32q4s2k_32q1s2k"         # 4 prefill + 32 spec + 32 decode
-
-  # Large batches with speculation
-  - "64q4s1k"                    # 64 spec requests
-  - "32q8s2k"                    # 32 spec (8-token)
-  - "16q16s4k"                   # 16 spec (16-token)
-
-# Backends that support query length > 1
-backends:
-  - FLASH_ATTN_MLA    # reorder_batch_threshold = 512
-  - FLASHMLA          # reorder_batch_threshold = 1 (tunable)
-
-# FlashInfer-MLA also supports uniform spec-as-decode but with different mechanism
-# - FLASHINFER_MLA
-
-# Benchmark settings
-device: "cuda:0"
-repeats: 10  # More repeats for statistical significance
-warmup_iters: 5
-profile_memory: false
-
-# Test these threshold values for optimization
-parameter_sweep:
-  param_name: "reorder_batch_threshold"
-  values: [1, 2, 4, 8, 16, 32, 64, 128, 256, 512]
-  include_auto: false
-  label_format: "{backend}_threshold_{value}"

benchmarks/attention_benchmarks/configs/standard_attention.yaml

@@ -1,48 +0,0 @@
-# Standard attention backend benchmark configuration
-
-model:
-  num_layers: 32
-  num_q_heads: 32
-  num_kv_heads: 8  # GQA with 4:1 ratio
-  head_dim: 128
-  block_size: 16
-
-batch_specs:
-  # Pure prefill
-  - "q512"      # Small prefill (512 tokens)
-  - "q2k"       # Medium prefill (2048 tokens)
-  - "q4k"       # Large prefill (4096 tokens)
-  - "q8k"       # Very large prefill (8192 tokens)
-
-  # Pure decode
-  - "8q1s1k"      # 8 requests, 1k KV cache each
-  - "16q1s2k"     # 16 requests, 2k KV cache each
-  - "32q1s1k"     # 32 requests, 1k KV cache each
-  - "64q1s4k"     # 64 requests, 4k KV cache each
-
-  # Mixed prefill/decode
-  - "2q2k_8q1s1k"      # 2 prefill + 8 decode
-  - "4q1k_16q1s2k"     # 4 prefill + 16 decode
-  - "2q4k_32q1s1k"     # 2 large prefill + 32 decode
-
-  # Speculative decode (q <= 8)
-  - "16q2s1k"         # 16 requests, 2 spec tokens, 1k KV cache
-  - "16q4s1k"         # 16 requests, 4 spec tokens, 1k KV cache
-  - "16q8s1k"         # 16 requests, 8 spec tokens, 1k KV cache
-  - "32q4s2k"         # 32 requests, 4 spec tokens, 2k KV cache
-  - "8q8s4k"          # 8 requests, 8 spec tokens, 4k KV cache
-
-  # Context extension (chunked prefill)
-  - "q1ks2k"          # 1k query, 2k sequence
-  - "2q1ks4k"         # 2 requests: 1k query, 4k sequence
-
-# Available backends: FLASH_ATTN, TRITON_ATTN, FLASHINFER
-backends:
-  - FLASH_ATTN
-  - TRITON_ATTN
-  - FLASHINFER
-
-device: "cuda:0"
-repeats: 5
-warmup_iters: 3
-profile_memory: false

benchmarks/attention_benchmarks/mla_runner.py

@@ -1,891 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""
-MLA benchmark runner - shared utilities for MLA benchmarks.
-
-This module provides helpers for running MLA backends without
-needing full VllmConfig integration.
-"""
-
-import numpy as np
-import torch
-from batch_spec import parse_batch_spec
-from common import (
-    BenchmarkResult,
-    MockHfConfig,
-    MockIndexer,
-    MockKVBProj,
-    MockLayer,
-    setup_mla_dims,
-)
-
-from vllm.config import (
-    CacheConfig,
-    CompilationConfig,
-    ModelConfig,
-    ParallelConfig,
-    SchedulerConfig,
-    VllmConfig,
-    set_current_vllm_config,
-)
-
-# ============================================================================
-# VllmConfig Creation
-# ============================================================================
-
-
-def _add_mock_methods_to_model_config(model_config: ModelConfig) -> None:
-    """
-    Add mock methods for layer-specific queries to ModelConfig.
-
-    These methods are needed by metadata builders but aren't normally
-    present on ModelConfig when used in benchmark contexts.
-    """
-    import types
-
-    model_config.get_num_layers = types.MethodType(lambda self: 1, model_config)
-    model_config.get_sliding_window_for_layer = types.MethodType(
-        lambda self, _i: None, model_config
-    )
-    model_config.get_logits_soft_cap_for_layer = types.MethodType(
-        lambda self, _i: None, model_config
-    )
-    model_config.get_sm_scale_for_layer = types.MethodType(
-        lambda self, _i: 1.0 / model_config.get_head_size() ** 0.5, model_config
-    )
-
-
-def create_minimal_vllm_config(
-    model_name: str = "deepseek-v3",
-    block_size: int = 128,
-    max_num_seqs: int = 256,
-    mla_dims: dict | None = None,
-    index_topk: int | None = None,
-) -> VllmConfig:
-    """
-    Create minimal VllmConfig for MLA benchmarks.
-
-    Args:
-        model_name: Model name (deepseek-v2, deepseek-v3, etc.) - used if mla_dims not
-                    provided
-        block_size: KV cache block size
-        max_num_seqs: Maximum number of sequences
-        mla_dims: Optional custom MLA dimensions dict. If not provided, uses
-                  setup_mla_dims(model_name)
-        index_topk: Optional topk value for sparse MLA backends. If provided,
-                    the config will include index_topk for sparse attention.
-
-    Returns:
-        VllmConfig for benchmarking
-    """
-    # Get MLA dimensions - use provided or load from model name
-    if mla_dims is None:
-        mla_dims = setup_mla_dims(model_name)
-
-    # Create mock HF config first (avoids downloading from HuggingFace)
-    mock_hf_config = MockHfConfig(mla_dims, index_topk=index_topk)
-
-    # Create a temporary minimal config.json to avoid HF downloads
-    # This ensures consistent ModelConfig construction without network access
-    import json
-    import os
-    import shutil
-    import tempfile
-
-    minimal_config = {
-        "architectures": ["DeepseekV2ForCausalLM"],
-        "model_type": "deepseek_v2",
-        "num_attention_heads": mla_dims["num_q_heads"],
-        "num_key_value_heads": mla_dims["num_kv_heads"],
-        "hidden_size": mla_dims["head_dim"] * mla_dims["num_q_heads"],
-        "torch_dtype": "bfloat16",
-        "max_position_embeddings": 163840,  # DeepSeek V3 default
-        "rope_theta": 10000.0,
-        "vocab_size": 128256,
-    }
-
-    # Create temporary directory with config.json
-    temp_dir = tempfile.mkdtemp(prefix="vllm_bench_")
-    config_path = os.path.join(temp_dir, "config.json")
-    with open(config_path, "w") as f:
-        json.dump(minimal_config, f)
-
-    try:
-        # Create model config using local path - no HF downloads
-        model_config = ModelConfig(
-            model=temp_dir,  # Use local temp directory
-            tokenizer=None,
-            tokenizer_mode="auto",
-            trust_remote_code=True,
-            dtype="bfloat16",
-            seed=0,
-            max_model_len=32768,
-            quantization=None,
-            enforce_eager=False,
-            max_logprobs=20,
-            disable_sliding_window=False,
-            skip_tokenizer_init=True,
-            served_model_name=None,
-            limit_mm_per_prompt=None,
-            config_format="auto",
-        )
-    finally:
-        # Clean up temporary directory
-        shutil.rmtree(temp_dir, ignore_errors=True)
-
-    # Override with our mock config
-    model_config.hf_config = mock_hf_config
-    model_config.hf_text_config = mock_hf_config
-
-    # Add mock methods for layer-specific queries
-    _add_mock_methods_to_model_config(model_config)
-
-    # Create sub-configs
-    cache_config = CacheConfig(
-        block_size=block_size,
-        gpu_memory_utilization=0.9,
-        swap_space=0,
-        cache_dtype="auto",
-        enable_prefix_caching=False,
-    )
-
-    scheduler_config = SchedulerConfig(
-        max_num_seqs=max_num_seqs,
-        max_num_batched_tokens=8192,
-        max_model_len=32768,
-        is_encoder_decoder=False,
-        enable_chunked_prefill=True,
-    )
-
-    parallel_config = ParallelConfig(
-        tensor_parallel_size=1,
-    )
-
-    compilation_config = CompilationConfig()
-
-    return VllmConfig(
-        model_config=model_config,
-        cache_config=cache_config,
-        parallel_config=parallel_config,
-        scheduler_config=scheduler_config,
-        compilation_config=compilation_config,
-    )
-
-
-# ============================================================================
-# Backend Configuration
-# ============================================================================
-
-
-# Backend-specific properties that can't be inferred from the backend class
-# Keys are AttentionBackendEnum names (uppercase)
-_BACKEND_PROPERTIES = {
-    "FLASHMLA": {
-        "query_format": "concat",  # Single concatenated tensor (vs tuple)
-    },
-    "FLASHMLA_SPARSE": {
-        "query_format": "concat",  # Single concatenated tensor (vs tuple)
-    },
-}
-
-
-def _get_backend_config(backend: str) -> dict:
-    """
-    Get backend configuration from AttentionBackendEnum.
-
-    Uses the registry to get the backend class and extract configuration
-    from its methods (get_impl_cls, get_builder_cls, is_sparse, etc.).
-
-    Args:
-        backend: Backend name matching AttentionBackendEnum exactly
-        (e.g., "FLASHMLA_SPARSE")
-
-    Returns:
-        Dict with backend configuration
-    """
-    from vllm.v1.attention.backends.registry import AttentionBackendEnum
-
-    try:
-        backend_enum = AttentionBackendEnum[backend]
-        backend_class = backend_enum.get_class()
-    except (KeyError, ValueError) as e:
-        valid_backends = [e.name for e in AttentionBackendEnum if e.name != "CUSTOM"]
-        raise ValueError(
-            f"Unknown backend: {backend}. "
-            f"Valid MLA backends: {[b for b in valid_backends if 'MLA' in b]}"
-        ) from e
-
-    # Get block size from backend class
-    block_sizes = backend_class.get_supported_kernel_block_sizes()
-    # Use first supported block size (backends typically support one for MLA)
-    block_size = block_sizes[0] if block_sizes else None
-    if hasattr(block_size, "value"):
-        # Handle MultipleOf enum
-        block_size = None
-
-    # Check if sparse via class method if available
-    is_sparse = getattr(backend_class, "is_sparse", lambda: False)()
-
-    # Get properties that can't be inferred
-    props = _BACKEND_PROPERTIES.get(backend, {})
-
-    return {
-        "backend_class": backend_class,
-        "impl_class": backend_class.get_impl_cls(),
-        "builder_class": backend_class.get_builder_cls(),
-        "query_format": props.get("query_format", "tuple"),
-        "block_size": block_size,
-        "is_sparse": is_sparse,
-    }
-
-
-# ============================================================================
-# Metadata Building Helpers
-# ============================================================================
-
-
-def _build_attention_metadata(
-    requests: list,
-    block_size: int,
-    device: torch.device,
-    builder_instance,
-) -> tuple:
-    """
-    Build attention metadata from batch requests.
-
-    Args:
-        requests: List of BatchRequest objects
-        block_size: KV cache block size
-        device: Target device
-        builder_instance: Metadata builder instance
-
-    Returns:
-        Tuple of (metadata, kv_cache_num_blocks)
-    """
-    q_lens = [r.q_len for r in requests]
-    kv_lens = [r.kv_len for r in requests]
-    total_q = sum(q_lens)
-    max_kv = max(kv_lens)
-
-    # Build query start locations
-    q_start_cpu = torch.tensor(
-        [0] + [sum(q_lens[: i + 1]) for i in range(len(q_lens))],
-        dtype=torch.int32,
-    )
-    q_start_gpu = q_start_cpu.to(device)
-
-    # Build sequence lengths
-    seq_lens_cpu = torch.tensor(kv_lens, dtype=torch.int32)
-    seq_lens_gpu = seq_lens_cpu.to(device)
-
-    # Build num_computed_tokens (context length for each request)
-    context_lens = [kv_len - q_len for q_len, kv_len in zip(q_lens, kv_lens)]
-    num_computed_tokens_cpu = torch.tensor(context_lens, dtype=torch.int32)
-
-    # Build block table
-    num_blocks_per_req = [(kv + block_size - 1) // block_size for kv in kv_lens]
-    max_num_blocks = max(num_blocks_per_req)
-
-    block_table_cpu = np.zeros((len(requests), max_num_blocks), dtype=np.int32)
-    current_block = 0
-    for i, num_blocks in enumerate(num_blocks_per_req):
-        for j in range(num_blocks):
-            block_table_cpu[i, j] = current_block
-            current_block += 1
-
-    block_table_gpu = torch.from_numpy(block_table_cpu).to(device)
-
-    # Build slot mapping
-    slot_mapping_list = []
-    for i, (q_len, kv_len, num_blocks) in enumerate(
-        zip(q_lens, kv_lens, num_blocks_per_req)
-    ):
-        context_len = kv_len - q_len
-        for j in range(q_len):
-            token_kv_idx = context_len + j
-            block_idx = token_kv_idx // block_size
-            offset_in_block = token_kv_idx % block_size
-            global_block_id = block_table_cpu[i, block_idx]
-            slot_id = global_block_id * block_size + offset_in_block
-            slot_mapping_list.append(slot_id)
-
-    slot_mapping = torch.tensor(slot_mapping_list, dtype=torch.int64, device=device)
-
-    # Create CommonAttentionMetadata
-    from vllm.v1.attention.backends.utils import CommonAttentionMetadata
-
-    common_attn_metadata = CommonAttentionMetadata(
-        num_reqs=len(requests),
-        max_query_len=max(q_lens),
-        max_seq_len=max_kv,
-        num_actual_tokens=total_q,
-        query_start_loc=q_start_gpu,
-        query_start_loc_cpu=q_start_cpu,
-        seq_lens=seq_lens_gpu,
-        _seq_lens_cpu=seq_lens_cpu,
-        _num_computed_tokens_cpu=num_computed_tokens_cpu,
-        slot_mapping=slot_mapping,
-        block_table_tensor=block_table_gpu,
-        dcp_local_seq_lens=None,
-    )
-
-    # Use the production build() method
-    metadata = builder_instance.build(
-        common_prefix_len=0,
-        common_attn_metadata=common_attn_metadata,
-        fast_build=False,
-    )
-
-    return metadata, current_block
-
-
-def _create_input_tensors(
-    total_q: int,
-    mla_dims: dict,
-    query_format: str,
-    device: torch.device,
-    dtype: torch.dtype,
-):
-    """
-    Create input tensors for both decode and prefill modes.
-
-    MLA requires different tensor formats for decode vs prefill:
-    - Decode: Uses kv_lora_rank (512) dimension
-    - Prefill: Uses qk_nope_head_dim (128) to stay under FlashAttention's 256 limit
-
-    Args:
-        total_q: Total number of query tokens
-        mla_dims: MLA dimension configuration
-        query_format: Either "tuple" or "concat"
-        device: Target device
-        dtype: Tensor dtype
-
-    Returns:
-        Tuple of (decode_inputs, prefill_inputs)
-        - decode_inputs: Query tensor(s) for decode mode
-        - prefill_inputs: Dict with 'q', 'k_c_normed', 'k_pe', 'k_scale' for prefill
-    """
-    if query_format == "tuple":
-        # Decode mode format: (q_nope, q_pe) where q_nope has kv_lora_rank dim
-        q_nope_decode = torch.randn(
-            total_q,
-            mla_dims["num_q_heads"],
-            mla_dims["kv_lora_rank"],
-            device=device,
-            dtype=dtype,
-        )
-        q_pe = torch.randn(
-            total_q,
-            mla_dims["num_q_heads"],
-            mla_dims["qk_rope_head_dim"],
-            device=device,
-            dtype=dtype,
-        )
-        decode_inputs = (q_nope_decode, q_pe)
-
-        # For prefill, we need q with qk_nope_head_dim instead of kv_lora_rank
-        q_nope_prefill = torch.randn(
-            total_q,
-            mla_dims["num_q_heads"],
-            mla_dims["qk_nope_head_dim"],
-            device=device,
-            dtype=dtype,
-        )
-        prefill_q = torch.cat([q_nope_prefill, q_pe], dim=-1)
-    else:  # concat
-        decode_inputs = torch.randn(
-            total_q,
-            mla_dims["num_q_heads"],
-            mla_dims["kv_lora_rank"] + mla_dims["qk_rope_head_dim"],
-            device=device,
-            dtype=dtype,
-        )
-        # For prefill with concat format
-        prefill_q = torch.randn(
-            total_q,
-            mla_dims["num_q_heads"],
-            mla_dims["qk_nope_head_dim"] + mla_dims["qk_rope_head_dim"],
-            device=device,
-            dtype=dtype,
-        )
-
-    # Create additional inputs needed for prefill forward
-    k_c_normed = torch.randn(
-        total_q,
-        mla_dims["kv_lora_rank"],
-        device=device,
-        dtype=dtype,
-    )
-    k_pe = torch.randn(
-        total_q,
-        1,  # Single head for MLA
-        mla_dims["qk_rope_head_dim"],
-        device=device,
-        dtype=dtype,
-    )
-    k_scale = torch.ones(1, device=device, dtype=torch.float32)
-
-    output = torch.zeros(
-        total_q,
-        mla_dims["num_q_heads"] * mla_dims["v_head_dim"],
-        device=device,
-        dtype=dtype,
-    )
-
-    prefill_inputs = {
-        "q": prefill_q,
-        "k_c_normed": k_c_normed,
-        "k_pe": k_pe,
-        "k_scale": k_scale,
-        "output": output,
-    }
-
-    return decode_inputs, prefill_inputs
-
-
-# ============================================================================
-# Backend Initialization
-# ============================================================================
-
-
-def _create_backend_impl(
-    backend_cfg: dict,
-    mla_dims: dict,
-    vllm_config: VllmConfig,
-    device: torch.device,
-    max_num_tokens: int = 8192,
-    index_topk: int | None = None,
-):
-    """
-    Create backend implementation instance.
-
-    Args:
-        backend_cfg: Backend configuration dict from _get_backend_config()
-        mla_dims: MLA dimension configuration
-        vllm_config: VllmConfig instance
-        device: Target device
-        max_num_tokens: Maximum number of tokens for sparse indexer buffer
-        index_topk: Topk value for sparse MLA backends
-
-    Returns:
-        Tuple of (impl, layer, builder_instance, indexer)
-    """
-    # Get classes from backend config (already resolved by _get_backend_config)
-    impl_class = backend_cfg["impl_class"]
-    builder_class = backend_cfg["builder_class"]
-
-    # Calculate scale
-    scale = 1.0 / np.sqrt(mla_dims["qk_nope_head_dim"] + mla_dims["qk_rope_head_dim"])
-
-    # Create mock kv_b_proj layer for prefill mode
-    mock_kv_b_proj = MockKVBProj(
-        num_heads=mla_dims["num_q_heads"],
-        qk_nope_head_dim=mla_dims["qk_nope_head_dim"],
-        v_head_dim=mla_dims["v_head_dim"],
-    )
-
-    # Create indexer for sparse backends
-    indexer = None
-    if backend_cfg.get("is_sparse", False):
-        if index_topk is None:
-            index_topk = 2048  # Default topk for sparse MLA
-        indexer = MockIndexer(
-            max_num_tokens=max_num_tokens,
-            topk_tokens=index_topk,
-            device=device,
-        )
-
-    # Build impl kwargs
-    impl_kwargs = {
-        "num_heads": mla_dims["num_q_heads"],
-        "head_size": mla_dims["head_dim"],
-        "scale": scale,
-        "num_kv_heads": mla_dims["num_kv_heads"],
-        "alibi_slopes": None,
-        "sliding_window": None,
-        "kv_cache_dtype": "auto",
-        "logits_soft_cap": None,
-        "attn_type": "decoder",
-        "kv_sharing_target_layer_name": None,
-        "q_lora_rank": None,
-        "kv_lora_rank": mla_dims["kv_lora_rank"],
-        "qk_nope_head_dim": mla_dims["qk_nope_head_dim"],
-        "qk_rope_head_dim": mla_dims["qk_rope_head_dim"],
-        "qk_head_dim": mla_dims["qk_nope_head_dim"] + mla_dims["qk_rope_head_dim"],
-        "v_head_dim": mla_dims["v_head_dim"],
-        "kv_b_proj": mock_kv_b_proj,
-    }
-
-    # Add indexer for sparse backends
-    if indexer is not None:
-        impl_kwargs["indexer"] = indexer
-
-    # Create impl
-    impl = impl_class(**impl_kwargs)
-
-    # Initialize DCP attributes
-    if not hasattr(impl, "dcp_world_size") or impl.dcp_world_size in (None, -1):
-        impl.dcp_world_size = 1
-        impl.dcp_rank = 0
-
-    # Create KV cache spec for MockLayer
-    from vllm.v1.kv_cache_interface import FullAttentionSpec
-
-    kv_cache_spec = FullAttentionSpec(
-        block_size=backend_cfg["block_size"] or vllm_config.cache_config.block_size,
-        num_kv_heads=1,  # MLA uses 1 KV head
-        head_size=576,  # MLA head dim
-        dtype=torch.bfloat16,
-    )
-
-    # Create mock layer
-    layer = MockLayer(device, impl=impl, kv_cache_spec=kv_cache_spec)
-
-    # Create builder instance if needed
-    builder_instance = None
-    if builder_class:
-        # Populate static_forward_context so builder can find the layer
-        # MockLayer inherits from AttentionLayerBase, so isinstance checks pass
-        vllm_config.compilation_config.static_forward_context = {"placeholder": layer}
-
-        builder_instance = builder_class(
-            kv_cache_spec=kv_cache_spec,
-            layer_names=["placeholder"],
-            vllm_config=vllm_config,
-            device=device,
-        )
-
-    return impl, layer, builder_instance, indexer
-
-
-# ============================================================================
-# Config Helpers
-# ============================================================================
-
-
-def _extract_mla_dims_from_config(config) -> dict | None:
-    """
-    Extract MLA dimensions from BenchmarkConfig if all required fields are present.
-
-    Args:
-        config: BenchmarkConfig instance
-
-    Returns:
-        Dict with MLA dimensions if all fields are provided, None otherwise
-    """
-    # Check if all MLA-specific fields are provided
-    if all(
-        [
-            config.kv_lora_rank is not None,
-            config.qk_nope_head_dim is not None,
-            config.qk_rope_head_dim is not None,
-            config.v_head_dim is not None,
-        ]
-    ):
-        return {
-            "kv_lora_rank": config.kv_lora_rank,
-            "qk_nope_head_dim": config.qk_nope_head_dim,
-            "qk_rope_head_dim": config.qk_rope_head_dim,
-            "v_head_dim": config.v_head_dim,
-            "num_q_heads": config.num_q_heads,
-            "num_kv_heads": config.num_kv_heads,
-            "head_dim": config.head_dim,
-        }
-    # Fallback: if MLA fields not fully specified, try to construct from basic fields
-    elif config.head_dim == 576:
-        # This looks like a DeepSeek MLA config, use standard dimensions with custom
-        # head count
-        return {
-            "kv_lora_rank": 512,
-            "qk_nope_head_dim": 128,
-            "qk_rope_head_dim": 64,
-            "v_head_dim": 128,
-            "num_q_heads": config.num_q_heads,
-            "num_kv_heads": config.num_kv_heads,
-            "head_dim": config.head_dim,
-        }
-    return None
-
-
-# ============================================================================
-# Benchmark Execution
-# ============================================================================
-
-
-def _run_single_benchmark(
-    config,
-    impl,
-    layer,
-    builder_instance,
-    backend_cfg: dict,
-    mla_dims: dict,
-    device: torch.device,
-    indexer=None,
-) -> BenchmarkResult:
-    """
-    Run a single benchmark iteration.
-
-    Args:
-        config: BenchmarkConfig instance
-        impl: Backend implementation instance
-        layer: MockLayer instance
-        builder_instance: Metadata builder instance
-        backend_cfg: Backend configuration dict
-        mla_dims: MLA dimension configuration
-        device: Target device
-        indexer: Optional MockIndexer for sparse backends
-
-    Returns:
-        BenchmarkResult with timing statistics
-    """
-    # Parse batch spec
-    requests = parse_batch_spec(config.batch_spec)
-    q_lens = [r.q_len for r in requests]
-    kv_lens = [r.kv_len for r in requests]
-    total_q = sum(q_lens)
-    max_kv_len = max(kv_lens)
-
-    # Determine block size
-    block_size = backend_cfg["block_size"] or config.block_size
-
-    # Build metadata
-    metadata, num_blocks = _build_attention_metadata(
-        requests, block_size, device, builder_instance
-    )
-
-    # Create KV cache
-    kv_cache = torch.zeros(
-        num_blocks,
-        block_size,
-        mla_dims["kv_lora_rank"] + mla_dims["qk_rope_head_dim"],
-        device=device,
-        dtype=torch.bfloat16,
-    )
-
-    # Create input tensors for both decode and prefill modes
-    decode_inputs, prefill_inputs = _create_input_tensors(
-        total_q,
-        mla_dims,
-        backend_cfg["query_format"],
-        device,
-        torch.bfloat16,
-    )
-
-    # Fill indexer with random indices for sparse backends
-    is_sparse = backend_cfg.get("is_sparse", False)
-    if is_sparse and indexer is not None:
-        indexer.fill_random_indices(total_q, max_kv_len)
-
-    # Determine which forward method to use
-    if is_sparse:
-        # Sparse backends use forward_mqa
-        forward_fn = lambda: impl.forward_mqa(decode_inputs, kv_cache, metadata, layer)
-    elif metadata.decode is not None:
-        forward_fn = lambda: impl._forward_decode(
-            decode_inputs, kv_cache, metadata, layer
-        )
-    elif metadata.prefill is not None:
-        forward_fn = lambda: impl._forward_prefill(
-            prefill_inputs["q"],
-            prefill_inputs["k_c_normed"],
-            prefill_inputs["k_pe"],
-            kv_cache,
-            metadata,
-            prefill_inputs["k_scale"],
-            prefill_inputs["output"],
-        )
-    else:
-        raise RuntimeError("Metadata has neither decode nor prefill metadata")
-
-    # Warmup
-    for _ in range(config.warmup_iters):
-        forward_fn()
-    torch.cuda.synchronize()
-
-    # Benchmark
-    times = []
-    for _ in range(config.repeats):
-        start = torch.cuda.Event(enable_timing=True)
-        end = torch.cuda.Event(enable_timing=True)
-
-        start.record()
-        for _ in range(config.num_layers):
-            forward_fn()
-        end.record()
-
-        torch.cuda.synchronize()
-        elapsed_ms = start.elapsed_time(end)
-        times.append(elapsed_ms / 1000.0 / config.num_layers)
-
-    mean_time = float(np.mean(times))
-    return BenchmarkResult(
-        config=config,
-        mean_time=mean_time,
-        std_time=float(np.std(times)),
-        min_time=float(np.min(times)),
-        max_time=float(np.max(times)),
-        throughput_tokens_per_sec=total_q / mean_time if mean_time > 0 else 0,
-    )
-
-
-def _run_mla_benchmark_batched(
-    backend: str,
-    configs_with_params: list[tuple],  # [(config, threshold, num_splits), ...]
-    index_topk: int = 2048,
-) -> list[BenchmarkResult]:
-    """
-    Unified batched MLA benchmark runner for all backends.
-
-    Works for: flashattn_mla, flashmla, flashinfer_mla, cutlass_mla,
-               flashinfer_mla_sparse, flashmla_sparse
-
-    This function reuses backend initialization across multiple benchmarks
-    to avoid setup/teardown overhead.
-
-    Args:
-        backend: Backend name
-        configs_with_params: List of (config, threshold, num_splits) tuples
-            - threshold: reorder_batch_threshold (FlashAttn/FlashMLA only)
-            - num_splits: num_kv_splits (CUTLASS only)
-        index_topk: Topk value for sparse MLA backends (default 2048)
-
-    Returns:
-        List of BenchmarkResult objects
-    """
-    if not configs_with_params:
-        return []
-
-    backend_cfg = _get_backend_config(backend)
-    device = torch.device(configs_with_params[0][0].device)
-    torch.cuda.set_device(device)
-
-    # Determine block size
-    config_block_size = configs_with_params[0][0].block_size
-    block_size = backend_cfg["block_size"] or config_block_size
-
-    # Extract MLA dimensions from the first config
-    first_config = configs_with_params[0][0]
-    mla_dims = _extract_mla_dims_from_config(first_config)
-
-    # If config didn't provide MLA dims, fall back to default model
-    if mla_dims is None:
-        mla_dims = setup_mla_dims("deepseek-v3")
-
-    # Determine if this is a sparse backend
-    is_sparse = backend_cfg.get("is_sparse", False)
-
-    # Create and set vLLM config for MLA (reused across all benchmarks)
-    vllm_config = create_minimal_vllm_config(
-        model_name="deepseek-v3",  # Used only for model path
-        block_size=block_size,
-        mla_dims=mla_dims,  # Use custom dims from config or default
-        index_topk=index_topk if is_sparse else None,
-    )
-
-    results = []
-
-    with set_current_vllm_config(vllm_config):
-        # Create backend impl, layer, builder, and indexer (reused across benchmarks)
-        impl, layer, builder_instance, indexer = _create_backend_impl(
-            backend_cfg,
-            mla_dims,
-            vllm_config,
-            device,
-            index_topk=index_topk if is_sparse else None,
-        )
-
-        # Run each benchmark with the shared impl
-        for config, threshold, num_splits in configs_with_params:
-            # Set threshold for this benchmark (FlashAttn/FlashMLA only)
-            original_threshold = None
-            if threshold is not None and builder_instance:
-                original_threshold = builder_instance.reorder_batch_threshold
-                builder_instance.reorder_batch_threshold = threshold
-
-            # Set num_splits for CUTLASS
-            original_num_splits = None
-            if num_splits is not None and hasattr(impl, "_num_kv_splits"):
-                original_num_splits = impl._num_kv_splits
-                impl._num_kv_splits = num_splits
-
-            try:
-                result = _run_single_benchmark(
-                    config,
-                    impl,
-                    layer,
-                    builder_instance,
-                    backend_cfg,
-                    mla_dims,
-                    device,
-                    indexer=indexer,
-                )
-                results.append(result)
-
-            finally:
-                # Restore original threshold
-                if original_threshold is not None:
-                    builder_instance.reorder_batch_threshold = original_threshold
-
-                # Restore original num_splits
-                if original_num_splits is not None:
-                    impl._num_kv_splits = original_num_splits
-
-    return results
-
-
-# ============================================================================
-# Public API
-# ============================================================================
-
-
-def run_mla_benchmark(
-    backend: str,
-    config,
-    reorder_batch_threshold: int | None = None,
-    num_kv_splits: int | None = None,
-    index_topk: int = 2048,
-) -> BenchmarkResult | list[BenchmarkResult]:
-    """
-    Unified MLA benchmark runner for all backends.
-
-    Works for: flashattn_mla, flashmla, flashinfer_mla, cutlass_mla,
-               flashinfer_mla_sparse, flashmla_sparse
-
-    Always uses batched execution internally for optimal performance.
-
-    Args:
-        backend: Backend name (flashattn_mla, flashmla, flashinfer_mla, cutlass_mla,
-                 flashinfer_mla_sparse, flashmla_sparse)
-        config: BenchmarkConfig or list of (BenchmarkConfig, param) tuples
-        reorder_batch_threshold: Threshold override for FlashAttn/FlashMLA
-                                 (single config mode only)
-        num_kv_splits: Number of KV splits for CUTLASS (single config mode only)
-        index_topk: Topk value for sparse MLA backends (default 2048)
-
-    Returns:
-        BenchmarkResult (single mode) or list of BenchmarkResult (batched mode)
-    """
-    # Normalize to batched mode: (config, threshold, num_splits)
-    if isinstance(config, list):
-        # Already in batched format
-        if len(config) > 0 and isinstance(config[0], tuple):
-            # Format: [(cfg, param), ...] where param is threshold or num_splits
-            if backend in ("flashattn_mla", "flashmla", "flashmla_sparse"):
-                configs_with_params = [(cfg, param, None) for cfg, param in config]
-            else:  # cutlass_mla, flashinfer_mla, or sparse backends
-                configs_with_params = [(cfg, None, param) for cfg, param in config]
-        else:
-            # Format: [cfg, ...] - just configs
-            configs_with_params = [(cfg, None, None) for cfg in config]
-        return_single = False
-    else:
-        # Single config: convert to batched format
-        configs_with_params = [(config, reorder_batch_threshold, num_kv_splits)]
-        return_single = True
-
-    # Use unified batched execution
-    results = _run_mla_benchmark_batched(backend, configs_with_params, index_topk)
-
-    # Return single result or list based on input
-    return results[0] if return_single else results

benchmarks/attention_benchmarks/runner.py

@@ -1,539 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""
-Standard attention benchmark runner - shared utilities for non-MLA benchmarks.
-
-This module provides helpers for running standard attention backends
-(FlashAttention, Triton, FlashInfer) with real vLLM integration.
-"""
-
-import logging
-import types
-from contextlib import contextmanager
-
-import numpy as np
-import torch
-from batch_spec import parse_batch_spec, reorder_for_flashinfer
-from common import BenchmarkConfig, BenchmarkResult, MockLayer, get_attention_scale
-
-from vllm.config import (
-    CacheConfig,
-    CompilationConfig,
-    DeviceConfig,
-    LoadConfig,
-    ModelConfig,
-    ParallelConfig,
-    SchedulerConfig,
-    VllmConfig,
-    set_current_vllm_config,
-)
-from vllm.v1.attention.backends.utils import (
-    CommonAttentionMetadata,
-    get_kv_cache_layout,
-    set_kv_cache_layout,
-)
-from vllm.v1.kv_cache_interface import FullAttentionSpec
-
-# ============================================================================
-# Backend Configuration
-# ============================================================================
-
-
-def _get_backend_config(backend: str) -> dict:
-    """
-    Get backend configuration from AttentionBackendEnum.
-
-    Args:
-        backend: Backend name matching AttentionBackendEnum exactly
-                 (e.g., "FLASH_ATTN", "TRITON_ATTN", "FLASHINFER")
-
-    Returns:
-        Dict with backend_class
-    """
-    from vllm.v1.attention.backends.registry import AttentionBackendEnum
-
-    try:
-        backend_enum = AttentionBackendEnum[backend]
-        backend_class = backend_enum.get_class()
-    except (KeyError, ValueError) as e:
-        valid_backends = [b.name for b in AttentionBackendEnum if b.name != "CUSTOM"]
-        raise ValueError(
-            f"Unknown backend: {backend}. Valid backends: {valid_backends}"
-        ) from e
-
-    return {"backend_class": backend_class}
-
-
-@contextmanager
-def log_warnings_and_errors_only():
-    """Temporarily set vLLM logger to WARNING level."""
-    logger = logging.getLogger("vllm")
-    old_level = logger.level
-    logger.setLevel(logging.WARNING)
-    try:
-        yield
-    finally:
-        logger.setLevel(old_level)
-
-
-# ============================================================================
-# Metadata Building Helpers
-# ============================================================================
-
-
-def _build_common_attn_metadata(
-    q_lens: list[int],
-    kv_lens: list[int],
-    block_size: int,
-    device: torch.device,
-) -> CommonAttentionMetadata:
-    """Build CommonAttentionMetadata from query/kv lengths."""
-    batch_size = len(q_lens)
-    total_tokens = sum(q_lens)
-
-    query_start_loc = torch.zeros(batch_size + 1, dtype=torch.int32, device=device)
-    query_start_loc[1:] = torch.tensor(q_lens, dtype=torch.int32, device=device).cumsum(
-        0
-    )
-    query_start_loc_cpu = query_start_loc.cpu()
-
-    seq_lens = torch.tensor(kv_lens, dtype=torch.int32, device=device)
-    max_seq_len = int(seq_lens.max().item())
-
-    max_blocks = (max(kv_lens) + block_size - 1) // block_size
-    num_blocks = batch_size * max_blocks
-    block_table_tensor = torch.arange(
-        num_blocks, dtype=torch.int32, device=device
-    ).view(batch_size, max_blocks)
-    slot_mapping = torch.arange(total_tokens, dtype=torch.int64, device=device)
-
-    max_query_len = max(q_lens)
-
-    return CommonAttentionMetadata(
-        query_start_loc=query_start_loc,
-        query_start_loc_cpu=query_start_loc_cpu,
-        seq_lens=seq_lens,
-        num_reqs=batch_size,
-        num_actual_tokens=total_tokens,
-        max_query_len=max_query_len,
-        max_seq_len=max_seq_len,
-        block_table_tensor=block_table_tensor,
-        slot_mapping=slot_mapping,
-        causal=True,
-    )
-
-
-def _create_vllm_config(
-    config: BenchmarkConfig,
-    max_num_blocks: int,
-) -> VllmConfig:
-    """Create a VllmConfig for benchmarking with mock model methods."""
-    model_config = ModelConfig(
-        model="meta-llama/Meta-Llama-3-8B",
-        tokenizer="meta-llama/Meta-Llama-3-8B",
-        trust_remote_code=False,
-        dtype="auto",  # Use model's native dtype
-        seed=0,
-        max_model_len=1024,
-    )
-
-    cache_config = CacheConfig(
-        block_size=config.block_size,
-        cache_dtype="auto",
-        swap_space=0,
-    )
-    cache_config.num_gpu_blocks = max_num_blocks
-    cache_config.num_cpu_blocks = 0
-
-    parallel_config = ParallelConfig(tensor_parallel_size=1)
-    scheduler_config = SchedulerConfig(
-        max_num_seqs=256,
-        max_num_batched_tokens=8192,
-        max_model_len=8192,
-        is_encoder_decoder=False,
-        enable_chunked_prefill=True,
-    )
-    device_config = DeviceConfig()
-    load_config = LoadConfig()
-    compilation_config = CompilationConfig()
-
-    # Add mock methods for benchmark config values
-    model_config.get_num_layers = types.MethodType(
-        lambda self: config.num_layers, model_config
-    )
-    model_config.get_sliding_window_for_layer = types.MethodType(
-        lambda self, i: None, model_config
-    )
-    model_config.get_logits_soft_cap_for_layer = types.MethodType(
-        lambda self, i: 0.0, model_config
-    )
-    model_config.get_sm_scale_for_layer = types.MethodType(
-        lambda self, i: 1.0 / config.head_dim**0.5, model_config
-    )
-    model_config.get_num_attention_heads = types.MethodType(
-        lambda self, parallel_config=None: config.num_q_heads, model_config
-    )
-    model_config.get_num_kv_heads = types.MethodType(
-        lambda self, parallel_config=None: config.num_kv_heads, model_config
-    )
-    model_config.get_head_size = types.MethodType(
-        lambda self: config.head_dim, model_config
-    )
-    model_config.get_sliding_window = types.MethodType(lambda self: None, model_config)
-
-    return VllmConfig(
-        model_config=model_config,
-        cache_config=cache_config,
-        parallel_config=parallel_config,
-        scheduler_config=scheduler_config,
-        device_config=device_config,
-        load_config=load_config,
-        compilation_config=compilation_config,
-    )
-
-
-# ============================================================================
-# Backend Initialization
-# ============================================================================
-
-
-def _create_backend_impl(
-    backend_cfg: dict,
-    config: BenchmarkConfig,
-    device: torch.device,
-    dtype: torch.dtype,
-):
-    """Create backend implementation instance."""
-    backend_class = backend_cfg["backend_class"]
-
-    scale = get_attention_scale(config.head_dim)
-
-    impl = backend_class.get_impl_cls()(
-        num_heads=config.num_q_heads,
-        head_size=config.head_dim,
-        scale=scale,
-        num_kv_heads=config.num_kv_heads,
-        alibi_slopes=None,
-        sliding_window=None,
-        kv_cache_dtype="auto",
-    )
-
-    kv_cache_spec = FullAttentionSpec(
-        block_size=config.block_size,
-        num_kv_heads=config.num_kv_heads,
-        head_size=config.head_dim,
-        dtype=dtype,
-    )
-
-    layer = MockLayer(device, kv_cache_spec=kv_cache_spec)
-
-    return backend_class, impl, layer
-
-
-def _create_metadata_builder(
-    backend_class,
-    kv_cache_spec: FullAttentionSpec,
-    vllm_config: VllmConfig,
-    device: torch.device,
-    backend_name: str = "",
-):
-    """Create metadata builder instance."""
-    layer_names = ["layer_0"]
-    builder_cls = backend_class.get_builder_cls()
-
-    # Flashinfer needs get_per_layer_parameters mocked since we don't have
-    # real model layers registered
-    if backend_name == "FLASHINFER":
-        import unittest.mock
-
-        from vllm.v1.attention.backends.utils import PerLayerParameters
-
-        def mock_get_per_layer_parameters(vllm_config, layer_names, impl_cls):
-            head_size = vllm_config.model_config.get_head_size()
-            return {
-                layer_name: PerLayerParameters(
-                    window_left=-1,  # No sliding window
-                    logits_soft_cap=0.0,  # No soft cap
-                    sm_scale=1.0 / (head_size**0.5),  # Standard scale
-                )
-                for layer_name in layer_names
-            }
-
-        with unittest.mock.patch(
-            "vllm.v1.attention.backends.flashinfer.get_per_layer_parameters",
-            mock_get_per_layer_parameters,
-        ):
-            return builder_cls(
-                kv_cache_spec=kv_cache_spec,
-                layer_names=layer_names,
-                vllm_config=vllm_config,
-                device=device,
-            )
-
-    return builder_cls(
-        kv_cache_spec=kv_cache_spec,
-        layer_names=layer_names,
-        vllm_config=vllm_config,
-        device=device,
-    )
-
-
-# ============================================================================
-# Tensor Creation Helpers
-# ============================================================================
-
-
-def _create_input_tensors(
-    config: BenchmarkConfig,
-    total_q: int,
-    device: torch.device,
-    dtype: torch.dtype,
-) -> tuple:
-    """Create Q, K, V input tensors for all layers."""
-    q_list = [
-        torch.randn(
-            total_q, config.num_q_heads, config.head_dim, device=device, dtype=dtype
-        )
-        for _ in range(config.num_layers)
-    ]
-    k_list = [
-        torch.randn(
-            total_q, config.num_kv_heads, config.head_dim, device=device, dtype=dtype
-        )
-        for _ in range(config.num_layers)
-    ]
-    v_list = [
-        torch.randn(
-            total_q, config.num_kv_heads, config.head_dim, device=device, dtype=dtype
-        )
-        for _ in range(config.num_layers)
-    ]
-    return q_list, k_list, v_list
-
-
-def _create_kv_cache(
-    config: BenchmarkConfig,
-    max_num_blocks: int,
-    backend_class,
-    device: torch.device,
-    dtype: torch.dtype,
-) -> list:
-    """Create KV cache tensors for all layers using the backend's methods.
-
-    Uses the backend's get_kv_cache_shape() and get_kv_cache_stride_order()
-    to create the cache with the correct shape and memory layout.
-    """
-    # Get the logical shape from the backend
-    cache_shape = backend_class.get_kv_cache_shape(
-        num_blocks=max_num_blocks,
-        block_size=config.block_size,
-        num_kv_heads=config.num_kv_heads,
-        head_size=config.head_dim,
-    )
-
-    # Get the stride order for custom memory layout
-    try:
-        stride_order = backend_class.get_kv_cache_stride_order()
-        assert len(stride_order) == len(cache_shape)
-    except (AttributeError, NotImplementedError):
-        stride_order = tuple(range(len(cache_shape)))
-
-    # Permute shape to physical layout order
-    physical_shape = tuple(cache_shape[i] for i in stride_order)
-
-    # Compute inverse permutation to get back to logical view
-    inv_order = [stride_order.index(i) for i in range(len(stride_order))]
-
-    cache_list = []
-    for _ in range(config.num_layers):
-        # Allocate in physical layout order (contiguous in memory)
-        cache = torch.zeros(*physical_shape, device=device, dtype=dtype)
-        # Permute to logical view
-        cache = cache.permute(*inv_order)
-        cache_list.append(cache)
-
-    return cache_list
-
-
-# ============================================================================
-# Benchmark Execution
-# ============================================================================
-
-
-def _run_single_benchmark(
-    config: BenchmarkConfig,
-    impl,
-    layer,
-    q_list: list,
-    k_list: list,
-    v_list: list,
-    cache_list: list,
-    attn_metadata,
-    device: torch.device,
-    dtype: torch.dtype,
-) -> tuple:
-    """Run single benchmark iteration with warmup and timing loop."""
-    total_q = q_list[0].shape[0]
-    out = torch.empty(
-        total_q, config.num_q_heads, config.head_dim, device=device, dtype=dtype
-    )
-
-    # Warmup
-    for _ in range(config.warmup_iters):
-        for i in range(config.num_layers):
-            impl.forward(
-                layer,
-                q_list[i],
-                k_list[i],
-                v_list[i],
-                cache_list[i],
-                attn_metadata,
-                output=out,
-            )
-    torch.cuda.synchronize()
-
-    # Benchmark
-    times = []
-    for _ in range(config.repeats):
-        start = torch.cuda.Event(enable_timing=True)
-        end = torch.cuda.Event(enable_timing=True)
-
-        start.record()
-        for i in range(config.num_layers):
-            impl.forward(
-                layer,
-                q_list[i],
-                k_list[i],
-                v_list[i],
-                cache_list[i],
-                attn_metadata,
-                output=out,
-            )
-        end.record()
-
-        torch.cuda.synchronize()
-        elapsed_ms = start.elapsed_time(end)
-        times.append(elapsed_ms / 1000.0 / config.num_layers)  # seconds per layer
-
-    mem_stats = {}
-    if config.profile_memory:
-        mem_stats = {
-            "allocated_mb": torch.cuda.memory_allocated(device) / 1024**2,
-            "reserved_mb": torch.cuda.memory_reserved(device) / 1024**2,
-        }
-
-    return times, mem_stats
-
-
-# ============================================================================
-# Public API
-# ============================================================================
-
-
-def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
-    """
-    Run standard attention benchmark with real kernels.
-
-    Supports: FLASH_ATTN, TRITON_ATTN, FLASHINFER
-
-    Args:
-        config: Benchmark configuration
-
-    Returns:
-        BenchmarkResult with timing and memory statistics
-    """
-    device = torch.device(config.device)
-    torch.cuda.set_device(device)
-
-    backend_cfg = _get_backend_config(config.backend)
-
-    requests = parse_batch_spec(config.batch_spec)
-
-    if config.backend == "FLASHINFER":
-        requests = reorder_for_flashinfer(requests)
-
-    q_lens = [r.q_len for r in requests]
-    kv_lens = [r.kv_len for r in requests]
-    total_q = sum(q_lens)
-    max_kv = max(kv_lens)
-    batch_size = len(q_lens)
-
-    # Calculate total blocks needed: batch_size * max_blocks_per_request
-    max_blocks_per_request = (max_kv + config.block_size - 1) // config.block_size
-    max_num_blocks = batch_size * max_blocks_per_request
-
-    # Suppress vLLM logs during setup to reduce spam
-    with log_warnings_and_errors_only():
-        # Create vllm_config first - uses model's native dtype via "auto"
-        vllm_config = _create_vllm_config(config, max_num_blocks)
-        dtype = vllm_config.model_config.dtype
-
-        # Wrap everything in set_current_vllm_config context
-        # This is required for backends like flashinfer that need global config
-        with set_current_vllm_config(vllm_config):
-            backend_class, impl, layer = _create_backend_impl(
-                backend_cfg, config, device, dtype
-            )
-
-            # Set KV cache layout if the backend requires a specific one
-            # (e.g., FlashInfer requires HND on SM100/Blackwell for TRTLLM attention)
-            required_layout = backend_class.get_required_kv_cache_layout()
-            if required_layout is not None:
-                set_kv_cache_layout(required_layout)
-                get_kv_cache_layout.cache_clear()
-
-            common_metadata = _build_common_attn_metadata(
-                q_lens, kv_lens, config.block_size, device
-            )
-
-            kv_cache_spec = FullAttentionSpec(
-                block_size=config.block_size,
-                num_kv_heads=config.num_kv_heads,
-                head_size=config.head_dim,
-                dtype=dtype,
-            )
-
-            builder = _create_metadata_builder(
-                backend_class, kv_cache_spec, vllm_config, device, config.backend
-            )
-
-            attn_metadata = builder.build(
-                common_prefix_len=0,
-                common_attn_metadata=common_metadata,
-            )
-
-            q_list, k_list, v_list = _create_input_tensors(
-                config, total_q, device, dtype
-            )
-
-            cache_list = _create_kv_cache(
-                config, max_num_blocks, backend_class, device, dtype
-            )
-
-            times, mem_stats = _run_single_benchmark(
-                config,
-                impl,
-                layer,
-                q_list,
-                k_list,
-                v_list,
-                cache_list,
-                attn_metadata,
-                device,
-                dtype,
-            )
-
-    mean_time = np.mean(times)
-    throughput = total_q / mean_time if mean_time > 0 else 0
-
-    return BenchmarkResult(
-        config=config,
-        mean_time=mean_time,
-        std_time=np.std(times),
-        min_time=np.min(times),
-        max_time=np.max(times),
-        throughput_tokens_per_sec=throughput,
-        memory_allocated_mb=mem_stats.get("allocated_mb"),
-        memory_reserved_mb=mem_stats.get("reserved_mb"),
-    )

benchmarks/kernels/benchmark_bitblas.py

@@ -0,0 +1,244 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from packaging import version
+
+from vllm.model_executor.layers.quantization.utils.bitblas_utils import (
+    MINIMUM_BITBLAS_VERSION,
+)
+
+try:
+    import bitblas
+
+    if version.parse(bitblas.__version__) < version.parse(MINIMUM_BITBLAS_VERSION):
+        raise ImportError(
+            "bitblas version is wrong. Please "
+            f"install bitblas>={MINIMUM_BITBLAS_VERSION}"
+        )
+except ImportError as e:
+    bitblas_import_exception = e
+    raise ValueError(
+        "Trying to use the bitblas backend, but could not import"
+        f"with the following error: {bitblas_import_exception}. "
+        "Please install bitblas through the following command: "
+        f"`pip install bitblas>={MINIMUM_BITBLAS_VERSION}`"
+    ) from bitblas_import_exception
+
+from bitblas import Matmul, MatmulConfig, auto_detect_nvidia_target
+
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+parser = FlexibleArgumentParser(
+    description="Benchmark BitBLAS int4 on a specific target."
+)
+
+# Add arguments to the parser
+parser.add_argument(
+    "--target",
+    type=str,
+    default=auto_detect_nvidia_target(),
+    help="Specify the target device for benchmarking.",
+)
+parser.add_argument(
+    "--group_size", type=int, default=None, help="Group size for grouped quantization."
+)
+parser.add_argument(
+    "--A_dtype",
+    type=str,
+    default="float16",
+    choices=["float16", "float32", "float64", "int32", "int8"],
+    help="Data type of activation A.",
+)
+parser.add_argument(
+    "--W_dtype",
+    type=str,
+    default="int4",
+    choices=[
+        "float16",
+        "float32",
+        "float64",
+        "int32",
+        "int8",
+        "int4",
+        "int2",
+        "int1",
+        "nf4",
+        "fp4_e2m1",
+    ],
+    help="Data type of weight W.",
+)
+parser.add_argument(
+    "--accum_dtype",
+    type=str,
+    default="float16",
+    choices=["float16", "int32"],
+    help="Data type for accumulation.",
+)
+parser.add_argument(
+    "--out_dtype",
+    type=str,
+    default="float16",
+    choices=["float16", "float32", "int32", "int8"],
+    help="Data type for output.",
+)
+parser.add_argument(
+    "--layout",
+    type=str,
+    default="nt",
+    choices=["nt", "nn"],
+    help="Matrix layout, 'nt' for non-transpose A and transpose W.",
+)
+parser.add_argument(
+    "--with_bias", action="store_true", help="Include bias in the benchmark."
+)
+parser.add_argument(
+    "--with_scaling",
+    action="store_true",
+    help="Include scaling factor in the quantization.",
+)
+parser.add_argument(
+    "--with_zeros", action="store_true", help="Include zeros in the quantization."
+)
+parser.add_argument(
+    "--zeros_mode",
+    type=str,
+    default=None,
+    choices=["original", "rescale", "quantized"],
+    help="Specify the mode for calculating zeros.",
+)
+
+# Parse the arguments
+args = parser.parse_args()
+
+# Assign arguments to variables
+target = args.target
+A_dtype = args.A_dtype
+W_dtype = args.W_dtype
+accum_dtype = args.accum_dtype
+out_dtype = args.out_dtype
+layout = args.layout
+with_bias = args.with_bias
+group_size = args.group_size
+with_scaling = args.with_scaling
+with_zeros = args.with_zeros
+zeros_mode = args.zeros_mode
+
+# Define a list of shared arguments that repeat in every config
+shared_args = [
+    A_dtype,
+    W_dtype,
+    out_dtype,
+    accum_dtype,
+    layout,
+    with_bias,
+    group_size,
+    with_scaling,
+    with_zeros,
+    zeros_mode,
+]
+
+# Define just the (M, K, N) shapes in a more compact list
+shapes = [
+    # square test
+    (1, 16384, 16384),
+    # BLOOM-176B
+    (1, 43008, 14336),
+    (1, 14336, 14336),
+    (1, 57344, 14336),
+    (1, 14336, 57344),
+    # OPT-65B
+    (1, 9216, 9216),
+    (1, 36864, 9216),
+    (1, 9216, 36864),
+    (1, 22016, 8192),
+    # LLAMA-70B/65B
+    (1, 8192, 22016),
+    (1, 8192, 8192),
+    (1, 28672, 8192),
+    (1, 8192, 28672),
+    # square test
+    (16384, 16384, 16384),
+    # BLOOM-176B
+    (8192, 43008, 14336),
+    (8192, 14336, 14336),
+    (8192, 57344, 14336),
+    (8192, 14336, 57344),
+    # OPT-65B
+    (8192, 9216, 9216),
+    (8192, 36864, 9216),
+    (8192, 9216, 36864),
+    (8192, 22016, 8192),
+    # LLAMA-70B/65B
+    (8192, 8192, 22016),
+    (8192, 8192, 8192),
+    (8192, 28672, 8192),
+    (8192, 8192, 28672),
+]
+
+# Build test shapes with all the shared arguments
+test_shapes = [(MatmulConfig, Matmul, (*shape, *shared_args)) for shape in shapes]
+
+benchmark_sets = []
+benchmark_sets.extend(test_shapes)
+
+benchmark_results = {}
+for config_class, operator, input_args in benchmark_sets:
+    config = config_class(*input_args)
+    matmul = operator(config, target=target, enable_tuning=True)
+    kernel_latency = matmul.profile_latency()
+
+    print("Time cost is: {:.3f} ms".format(kernel_latency))
+
+    profile_config = {
+        f"{operator.__name__}-{'-'.join([str(i) for i in input_args])}": {
+            "BitBLAS_top20_latency": kernel_latency,
+        }
+    }
+
+    benchmark_results.update(profile_config)
+
+# Define headers for the table
+headers = [
+    "PrimFunc",
+    "Input Arguments",
+    "BitBLAS Top20 Latency",
+]
+
+# Calculate column widths for pretty printing
+col_widths = [0, 0, 0]
+for config_key, values in benchmark_results.items():
+    args_split = config_key.split("-")
+    func_name = args_split[0]
+    input_args_str = "-".join(args_split[1:])
+    col_widths[0] = max(col_widths[0], len(func_name) + 2, len(headers[0]) + 2)
+    col_widths[1] = max(col_widths[1], len(input_args_str) + 2, len(headers[1]) + 2)
+    col_widths[2] = max(
+        col_widths[2],
+        len(f"{values['BitBLAS_top20_latency']:.3f} ms") + 2,
+        len(headers[2]) + 2,
+    )
+    # break only if you want to measure widths from a single example;
+    # otherwise, let it loop over all items.
+
+# Print header
+for i, header in enumerate(headers):
+    headers[i] = header.ljust(col_widths[i])
+print("".join(headers))
+print("-" * sum(col_widths))
+
+# Print rows
+for config_key, values in benchmark_results.items():
+    args_split = config_key.split("-")
+    func_name = args_split[0]
+    input_args_str = "-".join(args_split[1:])
+    row = [
+        func_name,
+        input_args_str,
+        f"{values['BitBLAS_top20_latency']:.3f} ms",
+    ]
+    row_str = "".join(
+        [str(cell).ljust(col_widths[idx]) for idx, cell in enumerate(row)]
+    )
+    print(row_str)

benchmarks/kernels/benchmark_cutlass_moe_fp8.py

@@ -11,7 +11,6 @@ import torch
 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.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
@@ -162,7 +161,7 @@ def bench_run(
                 w2_fp8q_cutlass,
                 topk_weights,
                 topk_ids,
-                activation=MoEActivation.SILU,
+                activation="silu",
                 global_num_experts=num_experts,
             )
     torch.cuda.synchronize()

benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py

@@ -197,7 +197,7 @@ def bench_run(
         )
 
         kernel = mk.FusedMoEModularKernel(
-            MoEPrepareAndFinalizeNoEP(),
+            MoEPrepareAndFinalizeNoEP(defer_input_quant=True),
             CutlassExpertsFp4(
                 make_dummy_moe_config(),
                 quant_config=quant_config,
@@ -242,7 +242,7 @@ def bench_run(
         )
 
         kernel = mk.FusedMoEModularKernel(
-            MoEPrepareAndFinalizeNoEP(),
+            MoEPrepareAndFinalizeNoEP(defer_input_quant=True),
             CutlassExpertsFp4(
                 make_dummy_moe_config(),
                 quant_config=quant_config,

benchmarks/kernels/benchmark_fused_collective.py

@@ -5,7 +5,7 @@
 Benchmark for FlashInfer fused collective operations vs standard operations.
 
 This benchmark compares:
-1. FlashInfer's allreduce_fusion (fused allreduce + rmsnorm + optional quant)
+1. FlashInfer's trtllm_allreduce_fusion (fused allreduce + rmsnorm + optional quant)
 2. Standard tensor_model_parallel_all_reduce + separate rmsnorm/quant operations
 
 Usage with torchrun:
@@ -24,6 +24,7 @@ import torch.distributed as dist  # type: ignore
 
 from vllm.config.vllm import CompilationConfig, VllmConfig, set_current_vllm_config
 from vllm.distributed import (
+    get_tp_group,
     tensor_model_parallel_all_reduce,
 )
 from vllm.distributed.parallel_state import (
@@ -51,12 +52,11 @@ logger = init_logger(__name__)
 try:
     import flashinfer.comm as flashinfer_comm  # type: ignore
 
-    if not (
-        hasattr(flashinfer_comm, "allreduce_fusion")
-        and hasattr(flashinfer_comm, "create_allreduce_fusion_workspace")
-    ):
+    if not hasattr(flashinfer_comm, "trtllm_allreduce_fusion"):
         flashinfer_comm = None
-        logger.warning("FlashInfer comm module found but missing allreduce_fusion API")
+        logger.warning(
+            "FlashInfer comm module found but missing trtllm_allreduce_fusion"
+        )
 except ImportError:
     flashinfer_comm = None
     logger.warning("FlashInfer not found, only benchmarking standard operations")
@@ -75,7 +75,7 @@ _FI_MAX_SIZES = {
 }
 
 # Global workspace tensor for FlashInfer
-_FI_WORKSPACE = None
+_FI_WORKSPACE_TENSOR = None
 
 
 def setup_flashinfer_workspace(
@@ -83,10 +83,10 @@ def setup_flashinfer_workspace(
     rank: int,
     hidden_dim: int,
     max_token_num: int,
-    dtype: torch.dtype,
+    use_fp32_lamport: bool = False,
 ):
     """Setup FlashInfer workspace for fused allreduce operations."""
-    global _FI_WORKSPACE
+    global _FI_WORKSPACE_TENSOR
 
     if flashinfer_comm is None:
         return None, None
@@ -96,29 +96,33 @@ def setup_flashinfer_workspace(
         return None, None
 
     try:
-        workspace = flashinfer_comm.create_allreduce_fusion_workspace(
-            backend="trtllm",
-            world_size=world_size,
-            rank=rank,
-            max_token_num=max_token_num,
-            hidden_dim=hidden_dim,
-            dtype=dtype,
+        # Create IPC workspace
+        ipc_handles, workspace_tensor = (
+            flashinfer_comm.trtllm_create_ipc_workspace_for_all_reduce_fusion(
+                tp_rank=rank,
+                tp_size=world_size,
+                max_token_num=max_token_num,
+                hidden_dim=hidden_dim,
+                group=get_tp_group().device_group,
+                use_fp32_lamport=use_fp32_lamport,
+            )
         )
 
-        _FI_WORKSPACE = workspace
-        return workspace
+        _FI_WORKSPACE_TENSOR = workspace_tensor
+        return ipc_handles, workspace_tensor
     except Exception as e:
         logger.error("Failed to setup FlashInfer workspace: %s", e)
-        return None
+        return None, None
 
 
-def cleanup_flashinfer_workspace(workspace):
+def cleanup_flashinfer_workspace(ipc_handles):
     """Cleanup FlashInfer workspace."""
-    if flashinfer_comm is None or workspace is None:
+    if flashinfer_comm is None or ipc_handles is None:
         return
 
     try:
-        workspace.destroy()
+        group = get_tp_group().device_group
+        flashinfer_comm.trtllm_destroy_ipc_workspace_for_all_reduce(ipc_handles, group)
     except Exception as e:
         logger.error("Failed to cleanup FlashInfer workspace: %s", e)
 
@@ -128,15 +132,25 @@ class FlashInferFusedAllReduceParams:
 
     def __init__(
         self,
+        rank: int,
+        world_size: int,
+        use_fp32_lamport: bool = False,
         max_token_num: int = 1024,
     ):
+        self.rank = rank
+        self.world_size = world_size
+        self.use_fp32_lamport = use_fp32_lamport
+        self.trigger_completion_at_end = True
         self.launch_with_pdl = True
         self.fp32_acc = True
         self.max_token_num = max_token_num
 
     def get_trtllm_fused_allreduce_kwargs(self):
         return {
+            "world_rank": self.rank,
+            "world_size": self.world_size,
             "launch_with_pdl": self.launch_with_pdl,
+            "trigger_completion_at_end": self.trigger_completion_at_end,
             "fp32_acc": self.fp32_acc,
         }
 
@@ -151,7 +165,7 @@ def flashinfer_fused_allreduce_rmsnorm(
     norm_out: torch.Tensor | None = None,
 ):
     """FlashInfer fused allreduce + rmsnorm operation."""
-    if flashinfer_comm is None or _FI_WORKSPACE is None:
+    if flashinfer_comm is None or _FI_WORKSPACE_TENSOR is None:
         raise RuntimeError("FlashInfer not available or workspace not initialized")
 
     if norm_out is None:
@@ -160,15 +174,18 @@ def flashinfer_fused_allreduce_rmsnorm(
     else:
         residual_out = input_tensor
 
-    flashinfer_comm.allreduce_fusion(
-        input=input_tensor,
-        workspace=_FI_WORKSPACE,
-        pattern=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNorm,
+    flashinfer_comm.trtllm_allreduce_fusion(
+        allreduce_in=input_tensor,
+        token_num=input_tensor.shape[0],
         residual_in=residual,
         residual_out=residual_out,
         norm_out=norm_out,
         rms_gamma=rms_gamma,
         rms_eps=rms_eps,
+        hidden_dim=input_tensor.shape[-1],
+        workspace_ptrs=_FI_WORKSPACE_TENSOR,
+        pattern_code=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNorm,
+        allreduce_out=None,
         quant_out=None,
         scale_out=None,
         layout_code=flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4,
@@ -190,7 +207,7 @@ def flashinfer_fused_allreduce_rmsnorm_fp8_quant(
     quant_out: torch.Tensor | None = None,
 ):
     """FlashInfer fused allreduce + rmsnorm + FP8 quantization."""
-    if flashinfer_comm is None or _FI_WORKSPACE is None:
+    if flashinfer_comm is None or _FI_WORKSPACE_TENSOR is None:
         raise RuntimeError("FlashInfer not available or workspace not initialized")
 
     if norm_out is None:
@@ -199,15 +216,18 @@ def flashinfer_fused_allreduce_rmsnorm_fp8_quant(
     else:
         residual_out = input_tensor
 
-    flashinfer_comm.allreduce_fusion(
-        input=input_tensor,
-        workspace=_FI_WORKSPACE,
-        pattern=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP8Quant,
+    flashinfer_comm.trtllm_allreduce_fusion(
+        allreduce_in=input_tensor,
+        token_num=input_tensor.shape[0],
         residual_in=residual,
         residual_out=residual_out,
         norm_out=norm_out,
         rms_gamma=rms_gamma,
         rms_eps=rms_eps,
+        hidden_dim=input_tensor.shape[-1],
+        workspace_ptrs=_FI_WORKSPACE_TENSOR,
+        pattern_code=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP8Quant,
+        allreduce_out=None,
         quant_out=quant_out,
         scale_out=None,
         layout_code=flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4,
@@ -230,7 +250,7 @@ def flashinfer_fused_allreduce_rmsnorm_fp4_quant(
     norm_out: torch.Tensor | None = None,
 ):
     """FlashInfer fused allreduce + rmsnorm + FP4 quantization."""
-    if flashinfer_comm is None or _FI_WORKSPACE is None:
+    if flashinfer_comm is None or _FI_WORKSPACE_TENSOR is None:
         raise RuntimeError("FlashInfer not available or workspace not initialized")
 
     if norm_out is None:
@@ -239,15 +259,18 @@ def flashinfer_fused_allreduce_rmsnorm_fp4_quant(
     else:
         residual_out = input_tensor
 
-    flashinfer_comm.allreduce_fusion(
-        input=input_tensor,
-        workspace=_FI_WORKSPACE,
-        pattern=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP4Quant,
+    flashinfer_comm.trtllm_allreduce_fusion(
+        allreduce_in=input_tensor,
+        token_num=input_tensor.shape[0],
         residual_in=residual,
         residual_out=residual_out,
         norm_out=norm_out,
         rms_gamma=rms_gamma,
         rms_eps=rms_eps,
+        hidden_dim=input_tensor.shape[-1],
+        workspace_ptrs=_FI_WORKSPACE_TENSOR,
+        pattern_code=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP4Quant,
+        allreduce_out=None,
         quant_out=quant_out,
         scale_out=output_scale,
         layout_code=flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4,
@@ -1017,31 +1040,23 @@ def main():
     configs = list(itertools.product(args.num_tokens, dtypes, residual_options))
 
     # Setup FlashInfer workspace if available
-    workspace = None
+    ipc_handles = None
     allreduce_params = None
 
     if flashinfer_comm is not None:
         # Use the largest hidden dimension for workspace setup
-        max_element_size = max(torch.finfo(dt).bits // 8 for dt in dtypes)
-        workspace_dtype = (
-            torch.float32
-            if max_element_size == 4
-            else (torch.bfloat16 if torch.bfloat16 in dtypes else torch.float16)
-        )
         max_num_token = _FI_MAX_SIZES.get(world_size) // (
-            args.hidden_dim * max_element_size
+            args.hidden_dim * world_size * 2
         )
 
-        workspace = setup_flashinfer_workspace(
-            world_size,
-            rank,
-            args.hidden_dim,
-            max_num_token,
-            dtype=workspace_dtype,
+        ipc_handles, workspace_tensor = setup_flashinfer_workspace(
+            world_size, rank, args.hidden_dim, max_num_token
         )
 
-        if workspace is not None:
+        if workspace_tensor is not None:
             allreduce_params = FlashInferFusedAllReduceParams(
+                rank=rank,
+                world_size=world_size,
                 max_token_num=max_num_token,
             )
 
@@ -1104,8 +1119,8 @@ def main():
 
     finally:
         # Cleanup
-        if workspace is not None:
-            cleanup_flashinfer_workspace(workspace)
+        if ipc_handles is not None:
+            cleanup_flashinfer_workspace(ipc_handles)
 
         dist.barrier()
 

benchmarks/kernels/benchmark_lora.py

@@ -842,7 +842,6 @@ class BenchmarkTensors:
             "sorted_token_ids": sorted_token_ids,
             "expert_ids": expert_ids,
             "num_tokens_post_padded": num_tokens_post_padded,
-            "token_lora_mapping": self.lora_kernel_meta.token_lora_mapping,
             "top_k_num": ctx.top_k_num,
             "device": self.input.device,
             "N": lora_rank,
@@ -916,7 +915,6 @@ class BenchmarkTensors:
             "sorted_token_ids": sorted_token_ids,
             "expert_ids": expert_ids,
             "num_tokens_post_padded": num_tokens_post_padded,
-            "token_lora_mapping": self.lora_kernel_meta.token_lora_mapping,
             "top_k_num": ctx.top_k_num,
             "device": self.input.device,
             "N": lora_rank,

benchmarks/kernels/benchmark_marlin.py

@@ -6,6 +6,12 @@ import torch.utils.benchmark as benchmark
 from benchmark_shapes import WEIGHT_SHAPES
 
 from vllm import _custom_ops as ops
+from vllm.model_executor.layers.quantization.gptq_marlin_24 import (
+    GPTQ_MARLIN_24_MAX_PARALLEL,
+    GPTQ_MARLIN_24_MIN_THREAD_N,
+    GPTQ_MARLIN_24_SUPPORTED_GROUP_SIZES,
+    GPTQ_MARLIN_24_SUPPORTED_QUANT_TYPES,
+)
 from vllm.model_executor.layers.quantization.utils.allspark_utils import (
     ALLSPARK_AMPERE_M_CUBLAS_THRESHOLD,
     ALLSPARK_SUPPORTED_QUANT_TYPES,
@@ -28,6 +34,9 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils_test import (
     awq_marlin_quantize,
     marlin_quantize,
 )
+from vllm.model_executor.layers.quantization.utils.marlin_utils_test_24 import (
+    marlin_24_quantize,
+)
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     gptq_pack,
     gptq_quantize_weights,
@@ -69,7 +78,14 @@ def bench_run(
     if size_k % group_size != 0:
         return
 
-    repack_supported = group_size in MARLIN_SUPPORTED_GROUP_SIZES
+    marlin_24_supported = (
+        quant_type in GPTQ_MARLIN_24_SUPPORTED_QUANT_TYPES
+        and group_size in GPTQ_MARLIN_24_SUPPORTED_GROUP_SIZES
+    )
+    repack_supported = (
+        quant_type in GPTQ_MARLIN_24_SUPPORTED_QUANT_TYPES
+        and group_size in MARLIN_SUPPORTED_GROUP_SIZES
+    )
     allspark_supported = (
         quant_type in ALLSPARK_SUPPORTED_QUANT_TYPES
         and group_size == -1
@@ -110,6 +126,14 @@ def bench_run(
             marlin_sort_indices,
         )
 
+    def gen_marlin_24_params():
+        marlin_24_w_ref = marlin_24_q_w_comp = marlin_24_meta = marlin_24_s = None
+        if marlin_24_supported:
+            (marlin_24_w_ref, marlin_24_q_w_comp, marlin_24_meta, marlin_24_s) = (
+                marlin_24_quantize(b, quant_type, group_size)
+            )
+        return (marlin_24_w_ref, marlin_24_q_w_comp, marlin_24_meta, marlin_24_s)
+
     def gen_repack_params():
         q_w_gptq = None
         repack_sort_indices = None
@@ -164,6 +188,9 @@ def bench_run(
         marlin_g_idx,
         marlin_sort_indices,
     ) = gen_marlin_params()
+    marlin_24_w_ref, marlin_24_q_w_comp, marlin_24_meta, marlin_24_s = (
+        gen_marlin_24_params()
+    )
     q_w_gptq, repack_sort_indices = gen_repack_params()
     qw_reorder, s_reorder, zp_reorder, sm_count, sm_version, CUBLAS_M_THRESHOLD = (
         gen_allspark_params()
@@ -173,6 +200,9 @@ def bench_run(
     marlin_workspace = MarlinWorkspace(
         size_n, GPTQ_MARLIN_MIN_THREAD_N, GPTQ_MARLIN_MAX_PARALLEL
     )
+    marlin_24_workspace = MarlinWorkspace(
+        size_n, GPTQ_MARLIN_24_MIN_THREAD_N, GPTQ_MARLIN_24_MAX_PARALLEL
+    )
 
     globals = {
         # Gen params
@@ -192,6 +222,12 @@ def bench_run(
         "marlin_sort_indices": marlin_sort_indices,
         "marlin_workspace": marlin_workspace,
         "is_k_full": is_k_full,
+        # Marlin_24 params
+        "marlin_24_w_ref": marlin_24_w_ref,
+        "marlin_24_q_w_comp": marlin_24_q_w_comp,
+        "marlin_24_meta": marlin_24_meta,
+        "marlin_24_s": marlin_24_s,
+        "marlin_24_workspace": marlin_24_workspace,
         # GPTQ params
         "q_w_gptq": q_w_gptq,
         "repack_sort_indices": repack_sort_indices,
@@ -204,6 +240,7 @@ def bench_run(
         "CUBLAS_M_THRESHOLD": CUBLAS_M_THRESHOLD,
         # Kernels
         "marlin_gemm": ops.marlin_gemm,
+        "gptq_marlin_24_gemm": ops.gptq_marlin_24_gemm,
         "gptq_marlin_repack": ops.gptq_marlin_repack,
         "allspark_w8a16_gemm": ops.allspark_w8a16_gemm,
     }
@@ -244,6 +281,17 @@ def bench_run(
         ).blocked_autorange(min_run_time=min_run_time)
     )
 
+    if marlin_24_supported:
+        results.append(
+            benchmark.Timer(
+                stmt="output = gptq_marlin_24_gemm(a, marlin_24_q_w_comp, marlin_24_meta, marlin_24_s, marlin_24_workspace.scratch, quant_type, size_m, size_n, size_k)",  # noqa: E501
+                globals=globals,
+                label=label,
+                sub_label=sub_label,
+                description="gptq_marlin_24_gemm",
+            ).blocked_autorange(min_run_time=min_run_time)
+        )
+
     if repack_supported:
         results.append(
             benchmark.Timer(

benchmarks/kernels/benchmark_moe.py

@@ -16,7 +16,6 @@ import torch
 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.config import (
     FusedMoEConfig,
     FusedMoEParallelConfig,
@@ -28,6 +27,7 @@ from vllm.model_executor.layers.fused_moe.fused_moe import *
 from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (
     TritonOrDeepGemmExperts,
 )
+from vllm.platforms import current_platform
 from vllm.transformers_utils.config import get_config
 from vllm.triton_utils import triton
 from vllm.utils.argparse_utils import FlexibleArgumentParser
@@ -100,38 +100,13 @@ def benchmark_config(
     dtype: torch.dtype,
     use_fp8_w8a8: bool,
     use_int8_w8a16: bool,
-    use_int4_w4a16: bool = False,
     num_iters: int = 100,
     block_quant_shape: list[int] = None,
     use_deep_gemm: bool = False,
 ) -> float:
     init_dtype = torch.float16 if use_fp8_w8a8 else dtype
     x = torch.randn(num_tokens, hidden_size, dtype=dtype)
-    if use_int4_w4a16:
-        # Int4 packed weights: 2 int4 values per uint8 byte
-        # K dimension is packed (halved)
-        intermediate_size = shard_intermediate_size // 2  # after silu_and_mul
-        w1 = torch.randint(
-            0,
-            255,
-            (
-                num_experts,
-                shard_intermediate_size,
-                hidden_size // 2,  # int4 packing
-            ),
-            dtype=torch.uint8,
-        )
-        w2 = torch.randint(
-            0,
-            255,
-            (
-                num_experts,
-                hidden_size,
-                intermediate_size // 2,  # int4 packing
-            ),
-            dtype=torch.uint8,
-        )
-    elif use_int8_w8a16:
+    if use_int8_w8a16:
         w1 = torch.randint(
             -127,
             127,
@@ -165,20 +140,7 @@ def benchmark_config(
     w2_scale = None
     a1_scale = None
     a2_scale = None
-    if use_int4_w4a16:
-        if block_quant_shape is None:
-            raise ValueError("block_quant_shape is required for int4_w4a16")
-        group_size = block_quant_shape[1]
-        # Scales shape: (E, N, K // group_size) in fp16
-        w1_scale = torch.rand(
-            (num_experts, shard_intermediate_size, hidden_size // group_size),
-            dtype=dtype,
-        )
-        w2_scale = torch.rand(
-            (num_experts, hidden_size, intermediate_size // group_size),
-            dtype=dtype,
-        )
-    elif use_int8_w8a16:
+    if use_int8_w8a16:
         w1_scale = torch.randn(
             (num_experts, 2 * shard_intermediate_size), dtype=torch.float32
         )
@@ -237,7 +199,6 @@ def benchmark_config(
             a1_scale=a1_scale,
             a2_scale=a2_scale,
             block_shape=block_quant_shape,
-            weight_dtype="int4" if use_int4_w4a16 else None,
         )
 
         deep_gemm_experts = None
@@ -251,8 +212,7 @@ def benchmark_config(
                         hidden_dim=hidden_size,
                         intermediate_size_per_partition=shard_intermediate_size,
                         num_local_experts=num_experts,
-                        num_logical_experts=num_experts,
-                        activation=MoEActivation.SILU,
+                        activation="silu",
                         moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
                         in_dtype=init_dtype,
                         routing_method=RoutingMethodType.TopK,
@@ -267,10 +227,9 @@ def benchmark_config(
                 x, input_gating, topk, renormalize=not use_deep_gemm
             )
 
-            inplace = not disable_inplace()
             if use_deep_gemm:
                 return deep_gemm_experts(
-                    x, w1, w2, topk_weights, topk_ids, inplace=inplace
+                    x, w1, w2, topk_weights, topk_ids, inplace=True
                 )
             return fused_experts(
                 x,
@@ -278,7 +237,7 @@ def benchmark_config(
                 w2,
                 topk_weights,
                 topk_ids,
-                inplace=inplace,
+                inplace=True,
                 quant_config=quant_config,
             )
 
@@ -520,18 +479,12 @@ class BenchmarkWorker:
         dtype: torch.dtype,
         use_fp8_w8a8: bool,
         use_int8_w8a16: bool,
-        use_int4_w4a16: bool = False,
         block_quant_shape: list[int] = None,
         use_deep_gemm: bool = False,
     ) -> tuple[dict[str, int], float]:
-        # local import to allow serialization by ray
-
         set_random_seed(self.seed)
         dtype_str = _get_config_dtype_str(
-            dtype,
-            use_int8_w8a16=use_int8_w8a16,
-            use_fp8_w8a8=use_fp8_w8a8,
-            use_int4_w4a16=use_int4_w4a16,
+            dtype, use_int8_w8a16=use_int8_w8a16, use_fp8_w8a8=use_fp8_w8a8
         )
         # NOTE(woosuk): The current naming convention uses w2.shape[2], which
         # is the intermediate size after silu_and_mul.
@@ -562,7 +515,6 @@ class BenchmarkWorker:
             dtype,
             use_fp8_w8a8,
             use_int8_w8a16,
-            use_int4_w4a16=use_int4_w4a16,
             num_iters=100,
             block_quant_shape=block_quant_shape,
             use_deep_gemm=use_deep_gemm,
@@ -579,18 +531,14 @@ class BenchmarkWorker:
         dtype: torch.dtype,
         use_fp8_w8a8: bool,
         use_int8_w8a16: bool,
-        use_int4_w4a16: bool,
         search_space: list[dict[str, int]],
         block_quant_shape: list[int],
         use_deep_gemm: bool,
     ) -> dict[str, int]:
-        # local import to allow serialization by ray
-        from vllm.platforms import current_platform
-
         best_config = None
         best_time = float("inf")
         if current_platform.is_rocm():
-            is_fp16 = not (use_fp8_w8a8 or use_int8_w8a16 or use_int4_w4a16)
+            is_fp16 = not (use_fp8_w8a8 or use_int8_w8a16)
             search_space = prune_rocm_search_space(
                 num_tokens,
                 shard_intermediate_size,
@@ -619,7 +567,6 @@ class BenchmarkWorker:
                         dtype,
                         use_fp8_w8a8,
                         use_int8_w8a16,
-                        use_int4_w4a16,
                         num_iters=20,
                         block_quant_shape=block_quant_shape,
                         use_deep_gemm=use_deep_gemm,
@@ -667,7 +614,6 @@ def sort_config(config: BenchmarkConfig) -> BenchmarkConfig:
             else {}
         ),
         **({"kpack": config["kpack"]} if "kpack" in config else {}),
-        **({"SPLIT_K": config["SPLIT_K"]} if "SPLIT_K" in config else {}),
     }
 
 
@@ -680,15 +626,11 @@ def save_configs(
     dtype: torch.dtype,
     use_fp8_w8a8: bool,
     use_int8_w8a16: bool,
-    use_int4_w4a16: bool,
     block_quant_shape: list[int],
     save_dir: str,
 ) -> None:
     dtype_str = _get_config_dtype_str(
-        dtype,
-        use_int8_w8a16=use_int8_w8a16,
-        use_fp8_w8a8=use_fp8_w8a8,
-        use_int4_w4a16=use_int4_w4a16,
+        dtype, use_int8_w8a16=use_int8_w8a16, use_fp8_w8a8=use_fp8_w8a8
     )
 
     # NOTE(woosuk): The current naming convention uses w2.shape[2], which
@@ -704,28 +646,20 @@ def save_configs(
         f.write("\n")
 
 
-def get_compressed_tensors_block_structure(config, default_value=None):
-    config_groups = config.get("config_groups", {})
-    if len(config_groups) != 1:
-        return default_value
-    group = next(iter(config_groups.values()))
-    weights = group.get("weights", {})
-    block_structure = weights.get("block_structure", default_value)
-    return block_structure
-
-
 def get_weight_block_size_safety(config, default_value=None):
     quantization_config = getattr(config, "quantization_config", {})
     if isinstance(quantization_config, dict):
-        if "weight_block_size" in quantization_config:
-            return quantization_config["weight_block_size"]
-        return get_compressed_tensors_block_structure(
-            quantization_config, default_value
-        )
+        return quantization_config.get("weight_block_size", default_value)
     return default_value
 
 
-def get_model_params(config):
+def main(args: argparse.Namespace):
+    print(args)
+
+    config = get_config(model=args.model, trust_remote_code=args.trust_remote_code)
+    if args.model_prefix:
+        config = getattr(config, args.model_prefix)
+
     if config.architectures[0] == "DbrxForCausalLM":
         E = config.ffn_config.moe_num_experts
         topk = config.ffn_config.moe_top_k
@@ -740,11 +674,9 @@ def get_model_params(config):
         "DeepseekV2ForCausalLM",
         "DeepseekV3ForCausalLM",
         "DeepseekV32ForCausalLM",
-        "GlmMoeDsaForCausalLM",
         "Glm4MoeForCausalLM",
         "Glm4MoeLiteForCausalLM",
         "NemotronHForCausalLM",
-        "MistralLarge3ForCausalLM",
     ):
         E = config.n_routed_experts
         topk = config.num_experts_per_tok
@@ -765,20 +697,16 @@ def get_model_params(config):
         topk = text_config.num_experts_per_tok
         intermediate_size = text_config.moe_intermediate_size
         hidden_size = text_config.hidden_size
-    elif config.architectures[0] == "HunYuanMoEV1ForCausalLM":
+    elif config.architectures[0] in ("HunYuanMoEV1ForCausalLM"):
         E = config.num_experts
         topk = config.moe_topk[0]
         intermediate_size = config.moe_intermediate_size[0]
         hidden_size = config.hidden_size
-    elif config.architectures[0] == "Qwen3OmniMoeForConditionalGeneration":
+    elif config.architectures[0] in ["Qwen3OmniMoeForConditionalGeneration"]:
         E = config.thinker_config.text_config.num_experts
         topk = config.thinker_config.text_config.num_experts_per_tok
         intermediate_size = config.thinker_config.text_config.moe_intermediate_size
         hidden_size = config.thinker_config.text_config.hidden_size
-    elif config.architectures[0] == "PixtralForConditionalGeneration":
-        # Pixtral can contain different LLM architectures,
-        # recurse to get their parameters
-        return get_model_params(config.get_text_config())
     else:
         # Support for llama4
         config = config.get_text_config()
@@ -787,48 +715,6 @@ def get_model_params(config):
         topk = config.num_experts_per_tok
         intermediate_size = config.intermediate_size
         hidden_size = config.hidden_size
-    return E, topk, intermediate_size, hidden_size
-
-
-def get_quantization_group_size(config) -> int | None:
-    """Extract the quantization group size from the HF model config.
-
-    This reads directly from the HuggingFace config object (as returned by
-    ``get_config()``), not from vLLM's quantization config classes.
-
-    Supports AWQ/GPTQ-style configs (direct 'group_size' key) and
-    compressed-tensors configs (nested inside 'config_groups').
-    """
-    quantization_config = getattr(config, "quantization_config", {})
-    if not isinstance(quantization_config, dict):
-        return None
-    # AWQ / GPTQ style: group_size is a top-level key
-    gs = quantization_config.get("group_size")
-    if gs is not None:
-        return gs
-    # compressed-tensors style: group_size is nested in config_groups
-    config_groups = quantization_config.get("config_groups", {})
-    if not isinstance(config_groups, dict):
-        return None
-    for group_cfg in config_groups.values():
-        if not isinstance(group_cfg, dict):
-            continue
-        weights = group_cfg.get("weights", {})
-        if not isinstance(weights, dict):
-            continue
-        gs = weights.get("group_size")
-        if gs is not None:
-            return gs
-    return None
-
-
-def main(args: argparse.Namespace):
-    print(args)
-
-    config = get_config(model=args.model, trust_remote_code=args.trust_remote_code)
-    if args.model_prefix:
-        config = getattr(config, args.model_prefix)
-    E, topk, intermediate_size, hidden_size = get_model_params(config)
     enable_ep = bool(args.enable_expert_parallel)
     if enable_ep:
         ensure_divisibility(E, args.tp_size, "Number of experts")
@@ -840,20 +726,7 @@ def main(args: argparse.Namespace):
     dtype = torch.float16 if current_platform.is_rocm() else config.dtype
     use_fp8_w8a8 = args.dtype == "fp8_w8a8"
     use_int8_w8a16 = args.dtype == "int8_w8a16"
-    use_int4_w4a16 = args.dtype == "int4_w4a16"
     block_quant_shape = get_weight_block_size_safety(config)
-    if use_int4_w4a16:
-        group_size = get_quantization_group_size(config)
-        if group_size is None:
-            raise ValueError(
-                "Could not determine group_size from model config. "
-                "The model's quantization_config must contain a 'group_size' "
-                "field (AWQ/GPTQ) or 'config_groups.*.weights.group_size' "
-                "(compressed-tensors)."
-            )
-        # For int4_w4a16, block_shape = [0, group_size]
-        # block_shape[0]=0 means no block quantization on N dimension
-        block_quant_shape = [0, group_size]
 
     if args.batch_size is None:
         batch_sizes = [
@@ -907,20 +780,8 @@ def main(args: argparse.Namespace):
         return ray.get(outputs)
 
     if args.tune:
-        # int4_w4a16 weights are uint8-packed, not fp16; treat like fp8 for
-        # search space generation (no matrix_instr_nonkdim/kpack exploration).
-        is_fp16 = not (use_fp8_w8a8 or use_int8_w8a16 or use_int4_w4a16)
-        # For int4_w4a16, the group_size constraint on BLOCK_SIZE_K does not
-        # apply: the gptq_awq kernel handles arbitrary BLOCK_SIZE_K regardless
-        # of group_size. Skip block_quant_shape filtering to keep the full
-        # search space (e.g. BLOCK_SIZE_K=64 with group_size=128).
-        tune_block_quant_shape = None if use_int4_w4a16 else block_quant_shape
-        search_space = get_configs_compute_bound(is_fp16, tune_block_quant_shape)
-        if use_int4_w4a16:
-            # SPLIT_K is a required kernel constexpr for gptq_awq kernel;
-            # only SPLIT_K=1 is used at runtime, so fix it during tuning.
-            for cfg in search_space:
-                cfg["SPLIT_K"] = 1
+        is_fp16 = not (use_fp8_w8a8 or use_int8_w8a16)
+        search_space = get_configs_compute_bound(is_fp16, block_quant_shape)
         print(f"Start tuning over {len(search_space)} configurations...")
         if use_deep_gemm:
             raise ValueError(
@@ -940,7 +801,6 @@ def main(args: argparse.Namespace):
                     dtype,
                     use_fp8_w8a8,
                     use_int8_w8a16,
-                    use_int4_w4a16,
                     search_space,
                     block_quant_shape,
                     use_deep_gemm,
@@ -960,7 +820,6 @@ def main(args: argparse.Namespace):
             dtype,
             use_fp8_w8a8,
             use_int8_w8a16,
-            use_int4_w4a16,
             block_quant_shape,
             args.save_dir,
         )
@@ -979,7 +838,6 @@ def main(args: argparse.Namespace):
                     dtype,
                     use_fp8_w8a8,
                     use_int8_w8a16,
-                    use_int4_w4a16,
                     block_quant_shape,
                     use_deep_gemm,
                 )
@@ -1002,10 +860,7 @@ if __name__ == "__main__":
     )
     parser.add_argument("--enable-expert-parallel", "-enable-ep", action="store_true")
     parser.add_argument(
-        "--dtype",
-        type=str,
-        choices=["auto", "fp8_w8a8", "int8_w8a16", "int4_w4a16"],
-        default="auto",
+        "--dtype", type=str, choices=["auto", "fp8_w8a8", "int8_w8a16"], default="auto"
     )
     parser.add_argument("--use-deep-gemm", action="store_true")
     parser.add_argument(

benchmarks/kernels/benchmark_moe_permute_unpermute.py

@@ -10,6 +10,8 @@ from transformers import AutoConfig
 
 from vllm.model_executor.layers.fused_moe import fused_topk
 from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import (
+    _moe_permute,
+    _moe_unpermute_and_reduce,
     moe_permute,
     moe_unpermute,
 )
@@ -39,13 +41,16 @@ def benchmark_permute(
     use_fp8_w8a8: bool,
     use_int8_w8a16: bool,
     num_iters: int = 100,
+    use_customized_permute: bool = False,
 ) -> float:
     # init_dtype = torch.float16 if use_fp8_w8a8 else dtype
     hidden_states = torch.randn(num_tokens, hidden_size, dtype=dtype)
     # output_hidden_states = torch.empty_like(hidden_states)
     if use_fp8_w8a8:
+        align_block_size = 128  # deepgemm needs 128 m aligned block
         qhidden_states, scale = _fp8_quantize(hidden_states, None, None)
     else:
+        align_block_size = None
         qhidden_states = hidden_states
 
     gating_output = torch.randn(num_iters, num_tokens, num_experts, dtype=torch.float32)
@@ -59,13 +64,29 @@ def benchmark_permute(
         input_gating.copy_(gating_output[i])
 
     def run():
-        moe_permute(
-            qhidden_states,
-            a1q_scale=None,
-            topk_ids=topk_ids,
-            n_expert=num_experts,
-            expert_map=None,
-        )
+        if use_customized_permute:
+            (
+                permuted_hidden_states,
+                a1q_scale,
+                first_token_off,
+                inv_perm_idx,
+                m_indices,
+            ) = moe_permute(
+                qhidden_states,
+                a1q_scale=None,
+                topk_ids=topk_ids,
+                n_expert=num_experts,
+                expert_map=None,
+                align_block_size=align_block_size,
+            )
+        else:
+            (
+                permuted_hidden_states,
+                a1q_scale,
+                sorted_token_ids,
+                expert_ids,
+                inv_perm,
+            ) = _moe_permute(qhidden_states, None, topk_ids, num_experts, None, 16)
 
     # JIT compilation & warmup
     run()
@@ -110,12 +131,16 @@ def benchmark_unpermute(
     use_fp8_w8a8: bool,
     use_int8_w8a16: bool,
     num_iters: int = 100,
+    use_customized_permute: bool = False,
 ) -> float:
     # init_dtype = torch.float16 if use_fp8_w8a8 else dtype
     hidden_states = torch.randn(num_tokens, hidden_size, dtype=dtype)
+    output_hidden_states = torch.empty_like(hidden_states)
     if use_fp8_w8a8:
+        align_block_size = 128  # deepgemm needs 128 m aligned block
         qhidden_states, scale = _fp8_quantize(hidden_states, None, None)
     else:
+        align_block_size = None
         qhidden_states = hidden_states
 
     input_gating = torch.randn(num_tokens, num_experts, dtype=torch.float32)
@@ -125,36 +150,78 @@ def benchmark_unpermute(
     )
 
     def prepare():
-        (
-            permuted_hidden_states,
-            _,
-            first_token_off,
-            inv_perm_idx,
-            _,
-        ) = moe_permute(
-            qhidden_states,
-            a1q_scale=None,
-            topk_ids=topk_ids,
-            n_expert=num_experts,
-            expert_map=None,
-        )
-        # convert to fp16/bf16 as gemm output
-        return (
-            permuted_hidden_states.to(dtype),
-            first_token_off,
-            inv_perm_idx,
-        )
+        if use_customized_permute:
+            (
+                permuted_hidden_states,
+                a1q_scale,
+                first_token_off,
+                inv_perm_idx,
+                m_indices,
+            ) = moe_permute(
+                qhidden_states,
+                a1q_scale=None,
+                topk_ids=topk_ids,
+                n_expert=num_experts,
+                expert_map=None,
+                align_block_size=align_block_size,
+            )
+            # convert to fp16/bf16 as gemm output
+            return (
+                permuted_hidden_states.to(dtype),
+                first_token_off,
+                inv_perm_idx,
+                m_indices,
+            )
+        else:
+            (
+                permuted_qhidden_states,
+                a1q_scale,
+                sorted_token_ids,
+                expert_ids,
+                inv_perm,
+            ) = _moe_permute(
+                qhidden_states, None, topk_ids, num_experts, None, block_m=16
+            )
+            # convert to fp16/bf16 as gemm output
+            return (
+                permuted_qhidden_states.to(dtype),
+                a1q_scale,
+                sorted_token_ids,
+                expert_ids,
+                inv_perm,
+            )
 
     def run(input: tuple):
-        (permuted_hidden_states, first_token_off, inv_perm_idx) = input
-        output = torch.empty_like(hidden_states)
-        moe_unpermute(
-            output,
-            permuted_hidden_states,
-            topk_weights,
-            inv_perm_idx,
-            first_token_off,
-        )
+        if use_customized_permute:
+            (
+                permuted_hidden_states,
+                first_token_off,
+                inv_perm_idx,
+                m_indices,
+            ) = input
+            output = torch.empty_like(hidden_states)
+            moe_unpermute(
+                output,
+                permuted_hidden_states,
+                topk_weights,
+                inv_perm_idx,
+                first_token_off,
+            )
+        else:
+            (
+                permuted_hidden_states,
+                a1q_scale,
+                sorted_token_ids,
+                expert_ids,
+                inv_perm,
+            ) = input
+            _moe_unpermute_and_reduce(
+                output_hidden_states,
+                permuted_hidden_states,
+                inv_perm,
+                topk_weights,
+                True,
+            )
 
     # JIT compilation & warmup
     input = prepare()
@@ -209,7 +276,8 @@ class BenchmarkWorker:
         dtype: torch.dtype,
         use_fp8_w8a8: bool,
         use_int8_w8a16: bool,
-    ) -> tuple[float, float]:
+        use_customized_permute: bool = False,
+    ) -> tuple[dict[str, int], float]:
         set_random_seed(self.seed)
 
         permute_time = benchmark_permute(
@@ -221,6 +289,7 @@ class BenchmarkWorker:
             use_fp8_w8a8,
             use_int8_w8a16,
             num_iters=100,
+            use_customized_permute=use_customized_permute,
         )
         unpermute_time = benchmark_unpermute(
             num_tokens,
@@ -231,6 +300,7 @@ class BenchmarkWorker:
             use_fp8_w8a8,
             use_int8_w8a16,
             num_iters=100,
+            use_customized_permute=use_customized_permute,
         )
         return permute_time, unpermute_time
 
@@ -277,6 +347,7 @@ def main(args: argparse.Namespace):
     dtype = torch.float16 if current_platform.is_rocm() else config.dtype
     use_fp8_w8a8 = args.dtype == "fp8_w8a8"
     use_int8_w8a16 = args.dtype == "int8_w8a16"
+    use_customized_permute = args.use_customized_permute
 
     if args.batch_size is None:
         batch_sizes = [
@@ -328,6 +399,7 @@ def main(args: argparse.Namespace):
                 dtype,
                 use_fp8_w8a8,
                 use_int8_w8a16,
+                use_customized_permute,
             )
             for batch_size in batch_sizes
         ],
@@ -347,6 +419,7 @@ if __name__ == "__main__":
     parser.add_argument(
         "--dtype", type=str, choices=["auto", "fp8_w8a8", "int8_w8a16"], default="auto"
     )
+    parser.add_argument("--use-customized-permute", action="store_true")
     parser.add_argument("--seed", type=int, default=0)
     parser.add_argument("--batch-size", type=int, required=False)
     parser.add_argument("--trust-remote-code", action="store_true")

benchmarks/kernels/benchmark_w8a8_block_fp8.py

@@ -22,8 +22,8 @@ from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 mp.set_start_method("spawn", force=True)
 
-assert current_platform.is_cuda() or current_platform.is_rocm(), (
-    "Only support tune w8a8 block fp8 kernel on CUDA/ROCm device."
+assert current_platform.is_cuda(), (
+    "Only support tune w8a8 block fp8 kernel on CUDA device."
 )
 
 DTYPE_MAP = {

benchmarks/kernels/cpu/benchmark_cpu_attn.py

@@ -14,7 +14,7 @@ from vllm._custom_ops import (
 )
 from vllm.platforms import CpuArchEnum, current_platform
 from vllm.utils.argparse_utils import FlexibleArgumentParser
-from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, set_random_seed
+from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
 from vllm.v1.attention.backends.cpu_attn import CPUAttentionBackend, _get_attn_isa
 
 
@@ -58,7 +58,7 @@ def main(
     seed: int = 0,
     iters: int = 20,
 ) -> None:
-    set_random_seed(seed)
+    current_platform.seed_everything(seed)
     num_seqs = len(seq_lens)
     query_lens = [x[0] for x in seq_lens]
     kv_lens = [x[1] for x in seq_lens]

benchmarks/kernels/cpu/benchmark_cpu_fused_moe.py

@@ -7,8 +7,8 @@ import time
 import numpy as np
 import torch
 
+from vllm.platforms import current_platform
 from vllm.utils.argparse_utils import FlexibleArgumentParser
-from vllm.utils.torch_utils import set_random_seed
 
 # Check if CPU MoE operations are available
 try:
@@ -41,7 +41,7 @@ def main(
     seed: int = 0,
     iters: int = 20,
 ) -> None:
-    set_random_seed(seed)
+    current_platform.seed_everything(seed)
     # up_dim = 2 * intermediate_size for gate + up projection
     up_dim = 2 * intermediate_size
 

cmake/cpu_extension.cmake

@@ -359,19 +359,6 @@ else()
     add_compile_definitions(-DVLLM_NUMA_DISABLED)
 endif()
 
-#
-# Generate CPU attention dispatch header
-#
-message(STATUS "Generating CPU attention dispatch header")
-execute_process(
-    COMMAND ${Python_EXECUTABLE} ${CMAKE_SOURCE_DIR}/csrc/cpu/generate_cpu_attn_dispatch.py
-    WORKING_DIRECTORY ${CMAKE_SOURCE_DIR}/csrc/cpu
-    RESULT_VARIABLE GEN_RESULT
-)
-if(NOT GEN_RESULT EQUAL 0)
-    message(FATAL_ERROR "Failed to generate CPU attention dispatch header")
-endif()
-
 #
 # _C extension
 #

cmake/external_projects/triton_kernels.cmake

@@ -1,9 +1,9 @@
 # Install OpenAI triton_kernels from https://github.com/triton-lang/triton/tree/main/python/triton_kernels
 
-set(DEFAULT_TRITON_KERNELS_TAG "v3.6.0")
+set(DEFAULT_TRITON_KERNELS_TAG "v3.5.0")
 
 # Set TRITON_KERNELS_SRC_DIR for use with local development with vLLM. We expect TRITON_KERNELS_SRC_DIR to
-# be directly set to the triton_kernels python directory.
+# be directly set to the triton_kernels python directory. 
 if (DEFINED ENV{TRITON_KERNELS_SRC_DIR})
   message(STATUS "[triton_kernels] Fetch from $ENV{TRITON_KERNELS_SRC_DIR}")
   FetchContent_Declare(
@@ -24,7 +24,7 @@ else()
   )
 endif()
 
-# Fetch content
+# Fetch content 
 FetchContent_MakeAvailable(triton_kernels)
 
 if (NOT triton_kernels_SOURCE_DIR)
@@ -47,7 +47,7 @@ install(CODE "file(MAKE_DIRECTORY \"\${CMAKE_INSTALL_PREFIX}/vllm/third_party/tr
 ## Copy .py files to install directory.
 install(DIRECTORY
         ${TRITON_KERNELS_PYTHON_DIR}
-        DESTINATION
+        DESTINATION 
         vllm/third_party/triton_kernels/
         COMPONENT triton_kernels
         FILES_MATCHING PATTERN "*.py")

cmake/external_projects/vllm_flash_attn.cmake

@@ -38,7 +38,7 @@ else()
   FetchContent_Declare(
           vllm-flash-attn
           GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
-          GIT_TAG 5824e6e2008271063c3229ab3e7032bd74abbbc6
+          GIT_TAG 188be16520ceefdc625fdf71365585d2ee348fe2
           GIT_PROGRESS TRUE
           # Don't share the vllm-flash-attn build between build types
           BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn

csrc/activation_kernels.cu

@@ -9,111 +9,6 @@
 
 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,
@@ -121,69 +16,52 @@ __device__ __forceinline__ scalar_t compute(const scalar_t& x,
   return act_first ? ACT_FN(x) * y : x * ACT_FN(y);
 }
 
-template <typename packed_t, packed_t (*PACKED_ACT_FN)(const packed_t&),
-          bool act_first>
-__device__ __forceinline__ packed_t packed_compute(const packed_t& x,
-                                                   const packed_t& y) {
-  return act_first ? packed_mul(PACKED_ACT_FN(x), y)
-                   : 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) {
+__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&),
-          packed_t (*PACKED_ACT_FN)(const packed_t&), bool act_first,
-          bool use_vec, bool use_256b = false>
+template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
+          bool act_first>
 __global__ void act_and_mul_kernel(
     scalar_t* __restrict__ out,          // [..., d]
     const scalar_t* __restrict__ input,  // [..., 2, d]
     const int d) {
-  const scalar_t* x_ptr = input + blockIdx.x * 2 * d;
+  constexpr int VEC_SIZE = 16 / sizeof(scalar_t);
+  const int64_t token_idx = blockIdx.x;
+  const scalar_t* x_ptr = input + token_idx * 2 * d;
   const scalar_t* y_ptr = x_ptr + d;
-  scalar_t* out_ptr = out + blockIdx.x * d;
+  scalar_t* out_ptr = out + token_idx * d;
 
-  if constexpr (use_vec) {
-    // 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);
+  // Check alignment for 128-bit vectorized access.
+  // All three pointers must be 16-byte aligned for safe int4 operations.
+  const bool aligned = is_16byte_aligned(x_ptr) && is_16byte_aligned(y_ptr) &&
+                       is_16byte_aligned(out_ptr);
 
-    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;
+  if (aligned && d >= VEC_SIZE) {
+    // Fast path: 128-bit vectorized loop
+    const int4* x_vec = reinterpret_cast<const int4*>(x_ptr);
+    const int4* y_vec = reinterpret_cast<const int4*>(y_ptr);
+    int4* out_vec = reinterpret_cast<int4*>(out_ptr);
+    const int num_vecs = d / VEC_SIZE;
+    const int vec_end = num_vecs * VEC_SIZE;
 
     for (int i = threadIdx.x; i < num_vecs; i += blockDim.x) {
-      vec_t x, y;
-      if constexpr (use_256b) {
-        ld256(x, &x_vec[i]);
-        ld256(y, &y_vec[i]);
-      } else {
-        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);
+      int4 x = VLLM_LDG(&x_vec[i]), y = VLLM_LDG(&y_vec[i]), r;
+      auto* xp = reinterpret_cast<scalar_t*>(&x);
+      auto* yp = reinterpret_cast<scalar_t*>(&y);
+      auto* rp = reinterpret_cast<scalar_t*>(&r);
 #pragma unroll
       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 {
-        out_vec[i] = x;
+        rp[j] = compute<scalar_t, ACT_FN, act_first>(xp[j], yp[j]);
       }
+      out_vec[i] = r;
+    }
+    // Scalar cleanup for remaining elements
+    for (int i = vec_end + threadIdx.x; i < d; i += blockDim.x) {
+      out_ptr[i] = compute<scalar_t, ACT_FN, act_first>(VLLM_LDG(&x_ptr[i]),
+                                                        VLLM_LDG(&y_ptr[i]));
     }
   } else {
     // Scalar fallback for unaligned data or small d
@@ -201,15 +79,6 @@ __device__ __forceinline__ T silu_kernel(const T& x) {
   return (T)(((float)x) / (1.0f + expf((float)-x)));
 }
 
-template <typename packed_t>
-__device__ __forceinline__ packed_t packed_silu_kernel(const packed_t& val) {
-  // x * sigmoid(x)
-  float2 fval = cast_to_float2(val);
-  fval.x = fval.x / (1.0f + expf(-fval.x));
-  fval.y = fval.y / (1.0f + expf(-fval.y));
-  return cast_to_packed<packed_t>(fval);
-}
-
 template <typename T>
 __device__ __forceinline__ T gelu_kernel(const T& x) {
   // Equivalent to PyTorch GELU with 'none' approximation.
@@ -220,18 +89,6 @@ __device__ __forceinline__ T gelu_kernel(const T& x) {
   return (T)(f * 0.5f * (1.0f + ::erf(f * ALPHA)));
 }
 
-template <typename packed_t>
-__device__ __forceinline__ packed_t packed_gelu_kernel(const packed_t& val) {
-  // Equivalent to PyTorch GELU with 'none' approximation.
-  // Refer to:
-  // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L36-L38
-  constexpr float ALPHA = M_SQRT1_2;
-  float2 fval = cast_to_float2(val);
-  fval.x = fval.x * 0.5f * (1.0f + ::erf(fval.x * ALPHA));
-  fval.y = fval.y * 0.5f * (1.0f + ::erf(fval.y * ALPHA));
-  return cast_to_packed<packed_t>(fval);
-}
-
 template <typename T>
 __device__ __forceinline__ T gelu_tanh_kernel(const T& x) {
   // Equivalent to PyTorch GELU with 'tanh' approximation.
@@ -245,83 +102,32 @@ __device__ __forceinline__ T gelu_tanh_kernel(const T& x) {
   return (T)(0.5f * f * (1.0f + ::tanhf(inner)));
 }
 
-template <typename packed_t>
-__device__ __forceinline__ packed_t
-packed_gelu_tanh_kernel(const packed_t& val) {
-  // Equivalent to PyTorch GELU with 'tanh' approximation.
-  // Refer to:
-  // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L25-L30
-  float2 fval = cast_to_float2(val);
-  constexpr float BETA = M_SQRT2 * M_2_SQRTPI * 0.5f;
-  constexpr float KAPPA = 0.044715;
-
-  float x_cube = fval.x * fval.x * fval.x;
-  float inner = BETA * (fval.x + KAPPA * x_cube);
-  fval.x = 0.5f * fval.x * (1.0f + ::tanhf(inner));
-
-  x_cube = fval.y * fval.y * fval.y;
-  inner = BETA * (fval.y + KAPPA * x_cube);
-  fval.y = 0.5f * fval.y * (1.0f + ::tanhf(inner));
-  return cast_to_packed<packed_t>(fval);
-}
-
 }  // namespace vllm
 
 // 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 = (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);         \
-    });                                                                     \
-  }
+#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, ACT_FIRST)                 \
+  int d = input.size(-1) / 2;                                            \
+  int64_t num_tokens = input.numel() / input.size(-1);                   \
+  dim3 grid(num_tokens);                                                 \
+  dim3 block(std::min(d, 1024));                                         \
+  if (num_tokens == 0) {                                                 \
+    return;                                                              \
+  }                                                                      \
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));      \
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();          \
+  VLLM_DISPATCH_FLOATING_TYPES(                                          \
+      input.scalar_type(), "act_and_mul_kernel", [&] {                   \
+        vllm::act_and_mul_kernel<scalar_t, KERNEL<scalar_t>, ACT_FIRST>  \
+            <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(),       \
+                                         input.data_ptr<scalar_t>(), d); \
+      });
 
 void silu_and_mul(torch::Tensor& out,    // [..., d]
                   torch::Tensor& input)  // [..., 2 * d]
 {
-  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
-                                true);
+  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, true);
 }
 
 void mul_and_silu(torch::Tensor& out,    // [..., d]
@@ -329,22 +135,19 @@ void mul_and_silu(torch::Tensor& out,    // [..., d]
 {
   // The difference between mul_and_silu and silu_and_mul is that mul_and_silu
   // applies the silu to the latter half of the input.
-  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
-                                false);
+  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, false);
 }
 
 void gelu_and_mul(torch::Tensor& out,    // [..., d]
                   torch::Tensor& input)  // [..., 2 * d]
 {
-  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_kernel, vllm::packed_gelu_kernel,
-                                true);
+  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_kernel, true);
 }
 
 void gelu_tanh_and_mul(torch::Tensor& out,    // [..., d]
                        torch::Tensor& input)  // [..., 2 * d]
 {
-  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_tanh_kernel,
-                                vllm::packed_gelu_tanh_kernel, true);
+  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_tanh_kernel, true);
 }
 
 namespace vllm {
@@ -355,57 +158,42 @@ __device__ __forceinline__ T fatrelu_kernel(const T& x, const float threshold) {
   return (T)(f > threshold ? f : 0.0f);
 }
 
-template <typename packed_t>
-__device__ __forceinline__ packed_t
-packed_fatrelu_kernel(const packed_t& val, const float threshold) {
-  float2 fval = cast_to_float2(val);
-  fval.x = fval.x > threshold ? fval.x : 0.0f;
-  fval.y = fval.y > threshold ? fval.y : 0.0f;
-  return cast_to_packed<packed_t>(fval);
-}
-
-template <typename scalar_t, typename packed_t,
-          scalar_t (*ACT_FN)(const scalar_t&, const float),
-          packed_t (*PACKED_ACT_FN)(const packed_t&, const float), bool use_vec,
-          bool use_256b = false>
+template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&, const float)>
 __global__ void act_and_mul_kernel_with_param(
     scalar_t* __restrict__ out, const scalar_t* __restrict__ input, const int d,
     const float param) {
-  const scalar_t* x_ptr = input + blockIdx.x * 2 * d;
+  constexpr int VEC_SIZE = 16 / sizeof(scalar_t);
+  const int64_t token_idx = blockIdx.x;
+  const scalar_t* x_ptr = input + token_idx * 2 * d;
   const scalar_t* y_ptr = x_ptr + d;
-  scalar_t* out_ptr = out + blockIdx.x * d;
+  scalar_t* out_ptr = out + token_idx * d;
 
-  if constexpr (use_vec) {
-    // 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);
+  // Check alignment for 128-bit vectorized access
+  const bool aligned = is_16byte_aligned(x_ptr) && is_16byte_aligned(y_ptr) &&
+                       is_16byte_aligned(out_ptr);
 
-    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;
+  if (aligned && d >= VEC_SIZE) {
+    // Fast path: 128-bit vectorized loop
+    const int4* x_vec = reinterpret_cast<const int4*>(x_ptr);
+    const int4* y_vec = reinterpret_cast<const int4*>(y_ptr);
+    int4* out_vec = reinterpret_cast<int4*>(out_ptr);
+    const int num_vecs = d / VEC_SIZE;
+    const int vec_end = num_vecs * VEC_SIZE;
 
     for (int i = threadIdx.x; i < num_vecs; i += blockDim.x) {
-      vec_t x, y;
-      if constexpr (use_256b) {
-        ld256(x, &x_vec[i]);
-        ld256(y, &y_vec[i]);
-      } else {
-        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);
+      int4 x = VLLM_LDG(&x_vec[i]), y = VLLM_LDG(&y_vec[i]), r;
+      auto* xp = reinterpret_cast<scalar_t*>(&x);
+      auto* yp = reinterpret_cast<scalar_t*>(&y);
+      auto* rp = reinterpret_cast<scalar_t*>(&r);
 #pragma unroll
       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 {
-        out_vec[i] = x;
+        rp[j] = ACT_FN(xp[j], param) * yp[j];
       }
+      out_vec[i] = r;
+    }
+    // Scalar cleanup for remaining elements
+    for (int i = vec_end + threadIdx.x; i < d; i += blockDim.x) {
+      out_ptr[i] = ACT_FN(VLLM_LDG(&x_ptr[i]), param) * VLLM_LDG(&y_ptr[i]);
     }
   } else {
     // Scalar fallback for unaligned data or small d
@@ -488,58 +276,20 @@ __global__ void swigluoai_and_mul_kernel(
 
 }  // namespace vllm
 
-#define LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM(KERNEL, PACKED_KERNEL, PARAM) \
-  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_with_param", [&] {                        \
-            vllm::act_and_mul_kernel_with_param<                               \
-                scalar_t, typename vllm::PackedTraits<scalar_t>::packed_t,     \
-                KERNEL<scalar_t>,                                              \
-                PACKED_KERNEL<                                                 \
-                    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);                                                        \
-          });                                                                  \
-    } else {                                                                   \
-      VLLM_DISPATCH_FLOATING_TYPES(                                            \
-          dtype, "act_and_mul_kernel_with_param", [&] {                        \
-            vllm::act_and_mul_kernel_with_param<                               \
-                scalar_t, typename vllm::PackedTraits<scalar_t>::packed_t,     \
-                KERNEL<scalar_t>,                                              \
-                PACKED_KERNEL<                                                 \
-                    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);                                                        \
-          });                                                                  \
-    }                                                                          \
-  } else {                                                                     \
-    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::PackedTraits<scalar_t>::packed_t,           \
-          KERNEL<scalar_t>,                                                    \
-          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);     \
-    });                                                                        \
-  }
+#define LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM(KERNEL, PARAM)         \
+  int d = input.size(-1) / 2;                                           \
+  int64_t num_tokens = input.numel() / input.size(-1);                  \
+  dim3 grid(num_tokens);                                                \
+  dim3 block(std::min(d, 1024));                                        \
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));     \
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();         \
+  VLLM_DISPATCH_FLOATING_TYPES(                                         \
+      input.scalar_type(), "act_and_mul_kernel_with_param", [&] {       \
+        vllm::act_and_mul_kernel_with_param<scalar_t, KERNEL<scalar_t>> \
+            <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(),      \
+                                         input.data_ptr<scalar_t>(), d, \
+                                         PARAM);                        \
+      });
 
 #define LAUNCH_SIGLUOAI_AND_MUL(KERNEL, ALPHA, LIMIT)                          \
   int d = input.size(-1) / 2;                                                  \
@@ -559,8 +309,7 @@ __global__ void swigluoai_and_mul_kernel(
 void fatrelu_and_mul(torch::Tensor& out,    // [..., d],
                      torch::Tensor& input,  // [..., 2 * d]
                      double threshold) {
-  LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM(
-      vllm::fatrelu_kernel, vllm::packed_fatrelu_kernel, threshold);
+  LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM(vllm::fatrelu_kernel, threshold);
 }
 void swigluoai_and_mul(torch::Tensor& out,    // [..., d]
                        torch::Tensor& input,  // [..., 2 * d]
@@ -570,41 +319,39 @@ void swigluoai_and_mul(torch::Tensor& out,    // [..., d]
 namespace vllm {
 
 // Element-wise activation kernel template.
-template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&), bool use_vec,
-          bool use_256b = false>
+template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&)>
 __global__ void activation_kernel(
     scalar_t* __restrict__ out,          // [..., d]
     const scalar_t* __restrict__ input,  // [..., d]
     const int d) {
-  const scalar_t* in_ptr = input + blockIdx.x * d;
-  scalar_t* out_ptr = out + blockIdx.x * d;
+  constexpr int VEC_SIZE = 16 / sizeof(scalar_t);
+  const int64_t token_idx = blockIdx.x;
+  const scalar_t* in_ptr = input + token_idx * d;
+  scalar_t* out_ptr = out + token_idx * d;
 
-  if constexpr (use_vec) {
-    // 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(scalar_t);
-    const vec_t* in_vec = reinterpret_cast<const vec_t*>(in_ptr);
-    vec_t* out_vec = reinterpret_cast<vec_t*>(out_ptr);
+  // Check alignment for 128-bit vectorized access
+  const bool aligned = is_16byte_aligned(in_ptr) && is_16byte_aligned(out_ptr);
+
+  if (aligned && d >= VEC_SIZE) {
+    // Fast path: 128-bit vectorized loop
+    const int4* in_vec = reinterpret_cast<const int4*>(in_ptr);
+    int4* out_vec = reinterpret_cast<int4*>(out_ptr);
     const int num_vecs = d / VEC_SIZE;
+    const int vec_end = num_vecs * VEC_SIZE;
 
     for (int i = threadIdx.x; i < num_vecs; i += blockDim.x) {
-      vec_t v;
-      if constexpr (use_256b) {
-        ld256(v, &in_vec[i]);
-      } else {
-        v = VLLM_LDG(&in_vec[i]);
-      }
+      int4 v = VLLM_LDG(&in_vec[i]), r;
       auto* vp = reinterpret_cast<scalar_t*>(&v);
+      auto* rp = reinterpret_cast<scalar_t*>(&r);
 #pragma unroll
       for (int j = 0; j < VEC_SIZE; j++) {
-        vp[j] = ACT_FN(vp[j]);
-      }
-      if constexpr (use_256b) {
-        st256(v, &out_vec[i]);
-      } else {
-        out_vec[i] = v;
+        rp[j] = ACT_FN(vp[j]);
       }
+      out_vec[i] = r;
+    }
+    // Scalar cleanup for remaining elements
+    for (int i = vec_end + threadIdx.x; i < d; i += blockDim.x) {
+      out_ptr[i] = ACT_FN(VLLM_LDG(&in_ptr[i]));
     }
   } else {
     // Scalar fallback for unaligned data or small d
@@ -618,43 +365,18 @@ __global__ void activation_kernel(
 }  // namespace vllm
 
 // Launch element-wise activation kernel.
-#define LAUNCH_ACTIVATION_KERNEL(KERNEL)                                 \
-  auto dtype = input.scalar_type();                                      \
-  int d = input.size(-1);                                                \
-  int64_t num_tokens = input.numel() / input.size(-1);                   \
-  if (num_tokens == 0) {                                                 \
-    return;                                                              \
-  }                                                                      \
-  dim3 grid(num_tokens);                                                 \
-  int cc_major = 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, "activation_kernel", [&] {     \
-        vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, true, true>  \
-            <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(),       \
-                                         input.data_ptr<scalar_t>(), d); \
-      });                                                                \
-    } else {                                                             \
-      VLLM_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] {     \
-        vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, true, false> \
-            <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(),       \
-                                         input.data_ptr<scalar_t>(), d); \
-      });                                                                \
-    }                                                                    \
-  } else {                                                               \
-    dim3 block(std::min(d, 1024));                                       \
-    VLLM_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] {       \
-      vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, false>         \
-          <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(),         \
-                                       input.data_ptr<scalar_t>(), d);   \
-    });                                                                  \
-  }
+#define LAUNCH_ACTIVATION_KERNEL(KERNEL)                                       \
+  int d = input.size(-1);                                                      \
+  int64_t num_tokens = input.numel() / d;                                      \
+  dim3 grid(num_tokens);                                                       \
+  dim3 block(std::min(d, 1024));                                               \
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));            \
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();                \
+  VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "activation_kernel", [&] { \
+    vllm::activation_kernel<scalar_t, KERNEL<scalar_t>>                        \
+        <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(),                 \
+                                     input.data_ptr<scalar_t>(), d);           \
+  });
 
 namespace vllm {
 

csrc/cache_kernels.cu

@@ -24,12 +24,6 @@
 typedef __hip_bfloat16 __nv_bfloat16;
 #endif
 
-#if defined(__gfx942__)
-constexpr float kFp8ScaleDivisor = 224.f;
-#else
-constexpr float kFp8ScaleDivisor = 448.f;
-#endif
-
 void swap_blocks(torch::Tensor& src, torch::Tensor& dst,
                  int64_t block_size_in_bytes,
                  const torch::Tensor& block_mapping) {
@@ -407,7 +401,8 @@ __global__ void concat_and_cache_ds_mla_kernel(
   }
 
   // Compute the scale for the tile
-  float tile_scale = fmaxf(max_abs / kFp8ScaleDivisor, FLT_MIN);
+  float tile_scale = max_abs / 448.f;
+  tile_scale = fmaxf(tile_scale, FLT_MIN);
 
   // The first lane of each half-warp writes the scale to kv_cache
   if ((lane_idx == 0) || (lane_idx == 16)) {
@@ -476,8 +471,11 @@ __global__ void indexer_k_quant_and_cache_kernel(
 #endif
   }
 
-  float scale = fmaxf(amax, 1e-4) / kFp8ScaleDivisor;
-
+#if defined(__gfx942__)
+  float scale = fmaxf(amax, 1e-4) / 224.0f;
+#else
+  float scale = fmaxf(amax, 1e-4) / 448.0f;
+#endif
   if (use_ue8m0) {
     scale = exp2f(ceilf(log2f(scale)));
   }
@@ -1234,13 +1232,8 @@ void cp_gather_and_upconvert_fp8_kv_cache(
               "src_cache and seq_lens must be on the same device");
   TORCH_CHECK(src_cache.device() == workspace_starts.device(),
               "src_cache and workspace_starts must be on the same device");
-  auto dtype = src_cache.scalar_type();
-  TORCH_CHECK(
-      dtype == at::ScalarType::Byte ||               // uint8
-          dtype == at::ScalarType::Float8_e4m3fn ||  // fp8 e4m3
-          dtype == at::ScalarType::Float8_e5m2,      // fp8 e5m2
-      "src_cache must be uint8, float8_e4m3fn, or float8_e5m2, but got ",
-      src_cache.dtype());
+
+  TORCH_CHECK(src_cache.dtype() == torch::kUInt8, "src_cache must be uint8");
   TORCH_CHECK(dst.dtype() == torch::kBFloat16, "dst must be bfloat16");
   TORCH_CHECK(head_dim == 576, "head_dim must be 576 for MLA");
 
@@ -1249,21 +1242,14 @@ void cp_gather_and_upconvert_fp8_kv_cache(
   int64_t cache_entry_stride = src_cache.stride(1);
   int64_t dst_entry_stride = dst.stride(0);
 
-  const uint8_t* src_ptr = nullptr;
-  if (dtype == at::ScalarType::Byte) {
-    src_ptr = src_cache.data_ptr<uint8_t>();
-  } else {
-    // float8_e4m3fn or float8_e5m2
-    src_ptr = reinterpret_cast<const uint8_t*>(src_cache.data_ptr());
-  }
-
   // Decide on the number of splits based on the batch size
   int num_splits = batch_size > 128 ? 2 : batch_size > 64 ? 4 : 16;
   dim3 grid(batch_size, num_splits);
   dim3 block(576);
 
   vllm::cp_gather_and_upconvert_fp8_kv_cache<<<grid, block, 0, stream>>>(
-      src_ptr, reinterpret_cast<__nv_bfloat16*>(dst.data_ptr()),
+      src_cache.data_ptr<uint8_t>(),
+      reinterpret_cast<__nv_bfloat16*>(dst.data_ptr()),
       block_table.data_ptr<int32_t>(), seq_lens.data_ptr<int32_t>(),
       workspace_starts.data_ptr<int32_t>(), block_size, head_dim,
       block_table_stride, cache_block_stride, cache_entry_stride,

csrc/cpu/cpu_attn.cpp

@@ -1,4 +1,79 @@
-#include "cpu_attn_dispatch_generated.h"
+#include "cpu_attn_vec.hpp"
+#include "cpu_attn_vec16.hpp"
+
+#ifdef CPU_CAPABILITY_AMXBF16
+  #include "cpu_attn_amx.hpp"
+  #define AMX_DISPATCH(...)                                                   \
+    case cpu_attention::ISA::AMX: {                                           \
+      using attn_impl = cpu_attention::AttentionImpl<cpu_attention::ISA::AMX, \
+                                                     scalar_t, head_dim>;     \
+      return __VA_ARGS__();                                                   \
+    }
+#else
+  #define AMX_DISPATCH(...) case cpu_attention::ISA::AMX:
+#endif
+
+#ifdef __aarch64__
+  #include "cpu_attn_neon.hpp"
+  // NEON requires head_dim to be a multiple of 32
+  #define NEON_DISPATCH(...)                                                   \
+    case cpu_attention::ISA::NEON: {                                           \
+      using attn_impl = cpu_attention::AttentionImpl<cpu_attention::ISA::NEON, \
+                                                     scalar_t, head_dim>;      \
+      return __VA_ARGS__();                                                    \
+    }
+#else
+  #define NEON_DISPATCH(...) case cpu_attention::ISA::NEON:
+#endif  // #ifdef __aarch64__
+
+#define CPU_ATTN_DISPATCH_CASE(HEAD_DIM, ...) \
+  case HEAD_DIM: {                            \
+    constexpr size_t head_dim = HEAD_DIM;     \
+    return __VA_ARGS__();                     \
+  }
+
+#define CPU_ATTN_DISPATCH_CASE_HEADDIM(HEAD_DIM, ...)           \
+  [&] {                                                         \
+    switch (HEAD_DIM) {                                         \
+      CPU_ATTN_DISPATCH_CASE(32, __VA_ARGS__)                   \
+      CPU_ATTN_DISPATCH_CASE(64, __VA_ARGS__)                   \
+      CPU_ATTN_DISPATCH_CASE(80, __VA_ARGS__)                   \
+      CPU_ATTN_DISPATCH_CASE(96, __VA_ARGS__)                   \
+      CPU_ATTN_DISPATCH_CASE(112, __VA_ARGS__)                  \
+      CPU_ATTN_DISPATCH_CASE(128, __VA_ARGS__)                  \
+      CPU_ATTN_DISPATCH_CASE(160, __VA_ARGS__)                  \
+      CPU_ATTN_DISPATCH_CASE(192, __VA_ARGS__)                  \
+      CPU_ATTN_DISPATCH_CASE(224, __VA_ARGS__)                  \
+      CPU_ATTN_DISPATCH_CASE(256, __VA_ARGS__)                  \
+      default: {                                                \
+        TORCH_CHECK(false, "Invalid CPU attention head_dim: " + \
+                               std::to_string(HEAD_DIM));       \
+      }                                                         \
+    }                                                           \
+  }()
+
+#define CPU_ATTN_DISPATCH_IMPL(ISA_TYPE, ...)                                 \
+  [&] {                                                                       \
+    switch (ISA_TYPE) {                                                       \
+      AMX_DISPATCH(__VA_ARGS__)                                               \
+      NEON_DISPATCH(__VA_ARGS__)                                              \
+      case cpu_attention::ISA::VEC: {                                         \
+        using attn_impl =                                                     \
+            cpu_attention::AttentionImpl<cpu_attention::ISA::VEC, scalar_t,   \
+                                         head_dim>;                           \
+        return __VA_ARGS__();                                                 \
+      }                                                                       \
+      case cpu_attention::ISA::VEC16: {                                       \
+        using attn_impl =                                                     \
+            cpu_attention::AttentionImpl<cpu_attention::ISA::VEC16, scalar_t, \
+                                         head_dim>;                           \
+        return __VA_ARGS__();                                                 \
+      }                                                                       \
+      default: {                                                              \
+        TORCH_CHECK(false, "Invalid CPU attention ISA type.");                \
+      }                                                                       \
+    }                                                                         \
+  }()
 
 torch::Tensor get_scheduler_metadata(
     const int64_t num_req, const int64_t num_heads_q,
@@ -47,14 +122,16 @@ torch::Tensor get_scheduler_metadata(
   input.enable_kv_split = enable_kv_split;
 
   VLLM_DISPATCH_FLOATING_TYPES(dtype, "get_scheduler_metadata", [&]() {
-    CPU_ATTN_DISPATCH(head_dim, isa, [&]() {
-      input.elem_size = sizeof(scalar_t);
-      input.q_buffer_elem_size = sizeof(attn_impl::q_buffer_t);
-      input.logits_buffer_elem_size = sizeof(attn_impl::logits_buffer_t);
-      input.output_buffer_elem_size =
-          sizeof(attn_impl::partial_output_buffer_t);
-      input.max_num_q_per_iter = attn_impl::MaxQHeadNumPerIteration;
-      input.kv_block_alignment = attn_impl::BlockSizeAlignment;
+    CPU_ATTN_DISPATCH_CASE_HEADDIM(head_dim, [&] {
+      CPU_ATTN_DISPATCH_IMPL(isa, [&]() {
+        input.elem_size = sizeof(scalar_t);
+        input.q_buffer_elem_size = sizeof(attn_impl::q_buffer_t);
+        input.logits_buffer_elem_size = sizeof(attn_impl::logits_buffer_t);
+        input.output_buffer_elem_size =
+            sizeof(attn_impl::partial_output_buffer_t);
+        input.max_num_q_per_iter = attn_impl::MaxQHeadNumPerIteration;
+        input.kv_block_alignment = attn_impl::BlockSizeAlignment;
+      });
     });
   });
 
@@ -107,14 +184,18 @@ void cpu_attn_reshape_and_cache(
 
   VLLM_DISPATCH_FLOATING_TYPES(
       key.scalar_type(), "cpu_attn_reshape_and_cache", [&]() {
-        CPU_ATTN_DISPATCH(head_dim, isa_tag, [&]() {
-          attn_impl::reshape_and_cache(
-              key.data_ptr<scalar_t>(), value.data_ptr<scalar_t>(),
-              key_cache.data_ptr<scalar_t>(), value_cache.data_ptr<scalar_t>(),
-              slot_mapping.data_ptr<int64_t>(), token_num, key_token_num_stride,
-              value_token_num_stride, head_num, key_head_num_stride,
-              value_head_num_stride, num_blocks, num_blocks_stride,
-              cache_head_num_stride, block_size, block_size_stride);
+        CPU_ATTN_DISPATCH_CASE_HEADDIM(head_dim, [&] {
+          CPU_ATTN_DISPATCH_IMPL(isa_tag, [&]() {
+            attn_impl::reshape_and_cache(
+                key.data_ptr<scalar_t>(), value.data_ptr<scalar_t>(),
+                key_cache.data_ptr<scalar_t>(),
+                value_cache.data_ptr<scalar_t>(),
+                slot_mapping.data_ptr<int64_t>(), token_num,
+                key_token_num_stride, value_token_num_stride, head_num,
+                key_head_num_stride, value_head_num_stride, num_blocks,
+                num_blocks_stride, cache_head_num_stride, block_size,
+                block_size_stride);
+          });
         });
       });
 }
@@ -176,10 +257,12 @@ void cpu_attention_with_kv_cache(
 
   VLLM_DISPATCH_FLOATING_TYPES(
       query.scalar_type(), "cpu_attention_with_kv_cache", [&]() {
-        CPU_ATTN_DISPATCH(query.size(2), input.metadata->isa, [&]() {
-          TORCH_CHECK_EQ(input.block_size % attn_impl::BlockSizeAlignment, 0);
-          cpu_attention::AttentionMainLoop<attn_impl> mainloop;
-          mainloop(&input);
+        CPU_ATTN_DISPATCH_CASE_HEADDIM(query.size(2), [&] {
+          CPU_ATTN_DISPATCH_IMPL(input.metadata->isa, [&]() {
+            TORCH_CHECK_EQ(input.block_size % attn_impl::BlockSizeAlignment, 0);
+            cpu_attention::AttentionMainLoop<attn_impl> mainloop;
+            mainloop(&input);
+          });
         });
       });
 }

csrc/cpu/cpu_attn_amx.hpp

@@ -377,7 +377,7 @@ class AttentionImpl<ISA::AMX, scalar_t, head_dim> {
       const int32_t q_heads_per_kv, const int64_t q_num_stride,
       const int64_t q_head_stride, const float scale) {
     constexpr int64_t bytes_per_head = head_dim * sizeof(scalar_t);
-    static_assert(bytes_per_head % AMX_TILE_ROW_BYTES == 0);
+    // static_assert(bytes_per_head % AMX_TILE_ROW_BYTES == 0);
     constexpr int64_t head_size_block_num = bytes_per_head / AMX_TILE_ROW_BYTES;
     constexpr int64_t head_elem_num_pre_block =
         AMX_TILE_ROW_BYTES / sizeof(scalar_t);

csrc/cpu/cpu_attn_impl.hpp

@@ -816,12 +816,16 @@ struct VecTypeTrait<float> {
   using vec_t = vec_op::FP32Vec16;
 };
 
+// ARM only supports BF16 with ARMv8.6-A extension
+#if (defined(__aarch64__) && !defined(ARM_BF16_SUPPORT))
+#else
 template <>
 struct VecTypeTrait<c10::BFloat16> {
   using vec_t = vec_op::BF16Vec16;
 };
+#endif
 
-#if !defined(__powerpc__)
+#if !defined(__powerpc__) && !defined(__s390x__)
 template <>
 struct VecTypeTrait<c10::Half> {
   using vec_t = vec_op::FP16Vec16;
@@ -1107,8 +1111,7 @@ class AttentionMainLoop {
           if (sliding_window_left != -1) {
             pos = std::max(pos, curr_token_pos - sliding_window_left);
           }
-          // Clamp to tile end to avoid OOB when window starts past the tile
-          return std::min(pos, kv_tile_end_pos);
+          return pos;
         }();
 
         int32_t right_kv_pos = [&]() {
@@ -1582,10 +1585,17 @@ class AttentionMainLoop {
 
               if (use_sink) {
                 alignas(64) float s_aux_fp32[16];
+#if defined(__aarch64__) && !defined(ARM_BF16_SUPPORT)
+                // ARM without native BF16 support: manual conversion
+                for (int i = 0; i < 16; ++i) {
+                  s_aux_fp32[i] = static_cast<float>(curr_s_aux[i]);
+                }
+#else
                 // All other platforms have BF16Vec16 available
                 vec_op::BF16Vec16 vec_bf16(curr_s_aux);
                 vec_op::FP32Vec16 vec_fp32(vec_bf16);
                 vec_fp32.save(s_aux_fp32);
+#endif
 
                 float* __restrict__ curr_sum_buffer = sum_buffer;
                 float* __restrict__ curr_max_buffer = max_buffer;

csrc/cpu/cpu_attn_neon.hpp

@@ -4,9 +4,6 @@
 #include "cpu_attn_impl.hpp"
 #include <arm_neon.h>
 #include <type_traits>
-#ifdef ARM_BF16_SUPPORT
-  #include "cpu_attn_neon_bfmmla.hpp"
-#endif
 namespace cpu_attention {
 
 namespace {
@@ -60,7 +57,7 @@ FORCE_INLINE void load_row8_B_as_f32<c10::BFloat16>(const c10::BFloat16* p,
 #endif
 }
 
-// Mx8, with 1 <= M <= 8 , K streamed, unroll-by-4 with ASIMD FMLAs
+// Mx8, with 1 <= M <= 8 , K streamed, unroll-by-4 with NEON FMLAs
 // #Loads = (K // 4) * (M + 4 * sizeof(kv_cache_t) / 2)
 // #FMLAs = (K // 4) * (4 * 2 * M)
 // We have (4 * 2 * M) FMLAs for (M + 4 * sizeof(kv_cache_t) / 2) loads
@@ -267,7 +264,7 @@ class AttentionImpl<ISA::NEON, scalar_t, head_dim> {
   constexpr static ISA ISAType = ISA::NEON;
   constexpr static bool scale_on_logits = false;  // apply scale on q_buffer
 
-  static_assert(HeadDim % HeadDimAlignment == 0);
+  //  static_assert(HeadDim % HeadDimAlignment == 0);
   // the gemm micro kernel is Mx8
   static_assert(HeadDimAlignment % 8 == 0);
   static_assert(BlockSizeAlignment % 8 == 0);
@@ -384,18 +381,6 @@ class AttentionImpl<ISA::NEON, scalar_t, head_dim> {
     }
   }
 };
-
-#ifdef ARM_BF16_SUPPORT
-// For BF16 on Arm, reuse the BFMMLA kernels with 32-token alignment.
-template <int64_t head_dim>
-class AttentionImpl<ISA::NEON, c10::BFloat16, head_dim>
-    : public AttentionImplNEONBFMMLA<BLOCK_SIZE_ALIGNMENT, ISA::NEON,
-                                     head_dim> {};
-#endif
 }  // namespace cpu_attention
 
-#undef BLOCK_SIZE_ALIGNMENT
-#undef HEAD_SIZE_ALIGNMENT
-#undef MAX_Q_HEAD_NUM_PER_ITER
-
-#endif  // #ifndef CPU_ATTN_ASIMD_HPP
+#endif  // #ifndef CPU_ATTN_NEON_HPP

csrc/cpu/cpu_attn_neon_bfmmla.hpp

@@ -1,682 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-#ifndef CPU_ATTN_NEON_BFMMLA_HPP
-#define CPU_ATTN_NEON_BFMMLA_HPP
-
-#include "cpu_attn_impl.hpp"
-
-#include <arm_neon.h>
-
-#include <cstdint>
-#include <vector>
-
-namespace cpu_attention {
-
-namespace {
-
-// BFMMLA tile dimensions
-constexpr int32_t TILE_ROWS = 2;  // M dimension
-constexpr int32_t TILE_K = 4;     // K reduction
-constexpr int32_t TILE_COLS = 2;  // N dimension (column-pair)
-
-// Derived constants
-constexpr int32_t OUTPUT_COLS_PER_BLOCK = 8;   // 4 column-pairs
-constexpr int32_t K_TOKENS_PER_GROUP = 8;      // Tokens grouped in K cache
-constexpr int32_t V_TOKENS_PER_ROW_BLOCK = 4;  // Tokens per V cache row block
-constexpr int32_t K_INNER_STRIDE = K_TOKENS_PER_GROUP * TILE_K;
-constexpr int32_t V_INNER_STRIDE = V_TOKENS_PER_ROW_BLOCK * TILE_COLS;
-constexpr int32_t PACK_ELEMENTS_PER_K_CHUNK = TILE_ROWS * TILE_K;  // A packing
-
-// Matrix Packing and Accumulator
-// Reshape two rows of Q into BFMMLA-friendly interleaved
-// Input:  row0 = [a0,a1,a2,a3], row1 = [b0,b1,b2,b3]
-// Output: [a0,a1,a2,a3,b0,b1,b2,b3, a4,a5,a6,a7,b4,b5,b6,b7]
-// For K tail (K % TILE_K != 0): pads with zeros to complete the final chunk
-FORCE_INLINE void reshape_Q_2xK_for_bfmmla(const c10::BFloat16* __restrict r0,
-                                           const c10::BFloat16* __restrict r1,
-                                           c10::BFloat16* __restrict dst,
-                                           int32_t K) {
-  const uint16_t* s0 = reinterpret_cast<const uint16_t*>(r0);
-  const uint16_t* s1 = reinterpret_cast<const uint16_t*>(r1);
-  uint16_t* d = reinterpret_cast<uint16_t*>(dst);
-
-  // Process TILE_K elements at a time (PACK_ELEMENTS_PER_K_CHUNK output)
-  int32_t k = 0;
-  for (; k + TILE_K <= K; k += TILE_K, d += PACK_ELEMENTS_PER_K_CHUNK) {
-    vst1q_u16(d, vcombine_u16(vld1_u16(s0 + k), vld1_u16(s1 + k)));
-  }
-
-  // Handle K tail: pack remaining elements with zero-padding
-  const int32_t tail = K - k;
-  if (tail > 0) {
-    // Pack remaining tail elements: [r0[k..k+tail-1], pad, r1[k..k+tail-1],
-    // pad]
-    for (int32_t t = 0; t < tail; ++t) {
-      d[t] = s0[k + t];
-      d[t + TILE_K] = s1[k + t];
-    }
-    // Zero-pad the rest
-    for (int32_t t = tail; t < TILE_K; ++t) {
-      d[t] = 0;
-      d[t + TILE_K] = 0;
-    }
-  }
-}
-
-// 2x2 accumulator load/store with compile-time row count
-template <int32_t m_rows>
-FORCE_INLINE float32x4_t load_acc_2x2(float* base, int64_t ldc, int col_off) {
-  static_assert(m_rows == 1 || m_rows == 2);
-  float32x2_t row0 = vld1_f32(base + col_off);
-  float32x2_t row1 =
-      (m_rows == 2) ? vld1_f32(base + ldc + col_off) : vdup_n_f32(0.f);
-  return vcombine_f32(row0, row1);
-}
-
-template <int32_t m_rows>
-FORCE_INLINE void store_acc_2x2(float32x4_t acc, float* base, int64_t ldc,
-                                int col_off) {
-  static_assert(m_rows == 1 || m_rows == 2);
-  vst1_f32(base + col_off, vget_low_f32(acc));
-  if constexpr (m_rows == 2) {
-    vst1_f32(base + ldc + col_off, vget_high_f32(acc));
-  }
-}
-
-// Initialize 4 column-pair accumulators for 2 rows (8 columns total)
-#define INIT_ACC_ROWPAIR_4(a0, a1, a2, a3, Crow, ldc, m_rows, accum) \
-  do {                                                               \
-    if (accum) {                                                     \
-      if (m_rows == 2) {                                             \
-        a0 = load_acc_2x2<2>(Crow, ldc, 0);                          \
-        a1 = load_acc_2x2<2>(Crow, ldc, 2);                          \
-        a2 = load_acc_2x2<2>(Crow, ldc, 4);                          \
-        a3 = load_acc_2x2<2>(Crow, ldc, 6);                          \
-      } else {                                                       \
-        a0 = load_acc_2x2<1>(Crow, ldc, 0);                          \
-        a1 = load_acc_2x2<1>(Crow, ldc, 2);                          \
-        a2 = load_acc_2x2<1>(Crow, ldc, 4);                          \
-        a3 = load_acc_2x2<1>(Crow, ldc, 6);                          \
-      }                                                              \
-    } else {                                                         \
-      a0 = a1 = a2 = a3 = vdupq_n_f32(0.f);                          \
-    }                                                                \
-  } while (0)
-
-// Store 4 column-pair accumulators back to C matrix
-#define STORE_ACC_ROWPAIR_4(a0, a1, a2, a3, Crow, ldc, m_rows) \
-  do {                                                         \
-    if (m_rows == 2) {                                         \
-      store_acc_2x2<2>(a0, Crow, ldc, 0);                      \
-      store_acc_2x2<2>(a1, Crow, ldc, 2);                      \
-      store_acc_2x2<2>(a2, Crow, ldc, 4);                      \
-      store_acc_2x2<2>(a3, Crow, ldc, 6);                      \
-    } else {                                                   \
-      store_acc_2x2<1>(a0, Crow, ldc, 0);                      \
-      store_acc_2x2<1>(a1, Crow, ldc, 2);                      \
-      store_acc_2x2<1>(a2, Crow, ldc, 4);                      \
-      store_acc_2x2<1>(a3, Crow, ldc, 6);                      \
-    }                                                          \
-  } while (0)
-
-// Perform 4 BFMMLA operations: acc += A @ B for 4 column-pairs
-#define BFMMLA_COMPUTE_4(r0, r1, r2, r3, a, b0, b1, b2, b3) \
-  do {                                                      \
-    r0 = vbfmmlaq_f32(r0, a, b0);                           \
-    r1 = vbfmmlaq_f32(r1, a, b1);                           \
-    r2 = vbfmmlaq_f32(r2, a, b2);                           \
-    r3 = vbfmmlaq_f32(r3, a, b3);                           \
-  } while (0)
-
-// Micro-kernel: updates a small fixed tile using BFMMLA.
-// RP = number of row-pairs (1,2,4)
-// Computes C[TILE_ROWS*RP, OUTPUT_COLS_PER_BLOCK] += A_packed @ B.
-// A_packed interleaves RP row-pairs; B layout is driven by the attention phase:
-// - AttentionGemmPhase::QK -> token-column layout (Q @ K^T)
-// - AttentionGemmPhase::PV -> token-row layout (P @ V)
-// K_static < 0 enables runtime K (PV only)
-template <int32_t RP, int32_t K_static, AttentionGemmPhase phase>
-FORCE_INLINE void gemm_rowpairs_x8_bfmmla_neon(
-    const bfloat16_t* const* __restrict A_packed_rp,
-    const int32_t* __restrict m_rows_rp, const bfloat16_t* __restrict B_blk,
-    float* __restrict C, int64_t ldc, bool accumulate, int64_t b_stride,
-    int32_t K_runtime = 0) {
-  static_assert(RP == 1 || RP == 2 || RP == 4, "RP must be 1,2,4");
-  static_assert(K_static < 0 || K_static % TILE_K == 0,
-                "K must be divisible by TILE_K");
-  static_assert(K_static >= 0 || phase == AttentionGemmPhase::PV,
-                "Runtime K only supported for PV");
-
-  constexpr bool runtime_k = (K_static < 0);
-  const int32_t K_iters =
-      runtime_k ? (K_runtime / TILE_K) : (K_static / TILE_K);
-  const int32_t K_tail = runtime_k ? (K_runtime % TILE_K) : 0;
-
-  if (!runtime_k) {
-    // Help the compiler fold away unused K_runtime when K is compile-time
-    (void)K_runtime;
-  }
-
-  auto* C_al = C;
-  const auto* B_al = B_blk;
-
-  // Setup A pointers
-  const bfloat16_t* a_ptr[4] = {
-      A_packed_rp[0],
-      (RP >= 2) ? A_packed_rp[1] : nullptr,
-      (RP >= 4) ? A_packed_rp[2] : nullptr,
-      (RP >= 4) ? A_packed_rp[3] : nullptr,
-  };
-
-  // Setup B pointers based on layout
-  const bfloat16_t* b_ptr[4];
-  if constexpr (phase == AttentionGemmPhase::PV) {
-    b_ptr[0] = B_blk + 0 * b_stride;
-    b_ptr[1] = B_blk + 1 * b_stride;
-    b_ptr[2] = B_blk + 2 * b_stride;
-    b_ptr[3] = B_blk + 3 * b_stride;
-  }
-
-  float32x4_t acc[4][4];
-
-// Initialize accumulators
-#define INIT_RP(rp)                                                            \
-  if constexpr (RP > rp) {                                                     \
-    INIT_ACC_ROWPAIR_4(acc[rp][0], acc[rp][1], acc[rp][2], acc[rp][3],         \
-                       C_al + (rp * 2) * ldc, ldc, m_rows_rp[rp], accumulate); \
-  }
-  INIT_RP(0);
-  INIT_RP(1);
-  INIT_RP(2);
-  INIT_RP(3);
-#undef INIT_RP
-
-  // Main compute loop
-  for (int32_t ki = 0; ki < K_iters; ++ki) {
-    bfloat16x8_t b0, b1, b2, b3;
-    if constexpr (phase == AttentionGemmPhase::PV) {
-      b0 = vld1q_bf16(b_ptr[0] + ki * V_INNER_STRIDE);
-      b1 = vld1q_bf16(b_ptr[1] + ki * V_INNER_STRIDE);
-      b2 = vld1q_bf16(b_ptr[2] + ki * V_INNER_STRIDE);
-      b3 = vld1q_bf16(b_ptr[3] + ki * V_INNER_STRIDE);
-    } else {
-      const bfloat16_t* b_base = B_al + ki * b_stride;
-      b0 = vld1q_bf16(b_base + 0 * V_INNER_STRIDE);
-      b1 = vld1q_bf16(b_base + 1 * V_INNER_STRIDE);
-      b2 = vld1q_bf16(b_base + 2 * V_INNER_STRIDE);
-      b3 = vld1q_bf16(b_base + 3 * V_INNER_STRIDE);
-    }
-
-#define COMPUTE_RP(rp)                                                       \
-  if constexpr (RP > rp) {                                                   \
-    bfloat16x8_t a = vld1q_bf16(a_ptr[rp] + ki * PACK_ELEMENTS_PER_K_CHUNK); \
-    BFMMLA_COMPUTE_4(acc[rp][0], acc[rp][1], acc[rp][2], acc[rp][3], a, b0,  \
-                     b1, b2, b3);                                            \
-  }
-    COMPUTE_RP(0);
-    COMPUTE_RP(1);
-    COMPUTE_RP(2);
-    COMPUTE_RP(3);
-#undef COMPUTE_RP
-  }
-
-  // K tail for runtime PV: fallback path
-  if constexpr (runtime_k) {
-    if (K_tail > 0) {
-      const int32_t tail_offset = K_iters * V_INNER_STRIDE;
-      const int32_t a_tail_offset = K_iters * PACK_ELEMENTS_PER_K_CHUNK;
-      for (int32_t kt = 0; kt < K_tail; ++kt) {
-        float32x4_t b_vecs[4];
-        for (int32_t p = 0; p < 4; ++p) {
-          const bfloat16_t* bp = b_ptr[p] + tail_offset + kt * TILE_COLS;
-          const float b0 = vcvtah_f32_bf16(bp[0]);
-          const float b1 = vcvtah_f32_bf16(bp[1]);
-          const float32x2_t b_pair = vset_lane_f32(b1, vdup_n_f32(b0), 1);
-          b_vecs[p] = vcombine_f32(b_pair, b_pair);
-        }
-
-#define TAIL_RP(rp)                                                     \
-  if constexpr (RP > rp) {                                              \
-    const bfloat16_t* ap = A_packed_rp[rp] + a_tail_offset;             \
-    float a_row0 = vcvtah_f32_bf16(ap[kt]);                             \
-    float a_row1 =                                                      \
-        (m_rows_rp[rp] == 2) ? vcvtah_f32_bf16(ap[kt + TILE_K]) : 0.0f; \
-    const float32x4_t a_vec =                                           \
-        vcombine_f32(vdup_n_f32(a_row0), vdup_n_f32(a_row1));           \
-    for (int32_t p = 0; p < 4; ++p) {                                   \
-      acc[rp][p] = vmlaq_f32(acc[rp][p], a_vec, b_vecs[p]);             \
-    }                                                                   \
-  }
-        TAIL_RP(0);
-        TAIL_RP(1);
-        TAIL_RP(2);
-        TAIL_RP(3);
-#undef TAIL_RP
-      }
-    }
-  }
-
-  // Store results
-#define STORE_RP(rp)                                                    \
-  if constexpr (RP > rp) {                                              \
-    STORE_ACC_ROWPAIR_4(acc[rp][0], acc[rp][1], acc[rp][2], acc[rp][3], \
-                        C_al + (rp * 2) * ldc, ldc, m_rows_rp[rp]);     \
-  }
-  STORE_RP(0);
-  STORE_RP(1);
-  STORE_RP(2);
-  STORE_RP(3);
-#undef STORE_RP
-}
-
-// Meso-kernel: packs a small MBxK slice of A, then tiles over N and calls the
-// micro-kernel for each OUTPUT_COLS_PER_BLOCK chunk. K_static < 0 enables
-// runtime K (PV only).
-template <int32_t MB, int32_t N, int32_t K_static, AttentionGemmPhase phase>
-FORCE_INLINE void gemm_packA_compute_MB_xN(
-    const c10::BFloat16* __restrict A, const c10::BFloat16* __restrict B,
-    float* __restrict C, int32_t K_runtime, int64_t lda, int64_t ldc,
-    int64_t b_layout_stride, int64_t b_reduction_stride, bool accumulate) {
-  static_assert(MB >= 1 && MB <= 8, "MB must be in [1,8]");
-  static_assert(N % OUTPUT_COLS_PER_BLOCK == 0,
-                "N must be a multiple of OUTPUT_COLS_PER_BLOCK");
-  static_assert(K_static < 0 || K_static % TILE_K == 0,
-                "K must be divisible by TILE_K");
-  static_assert(K_static >= 0 || phase == AttentionGemmPhase::PV,
-                "Runtime K only supported for PV");
-
-  constexpr bool runtime_k = (K_static < 0);
-  const int32_t K_val = runtime_k ? K_runtime : K_static;
-
-  // Keep small packs on-stack to avoid heap churn
-  constexpr int32_t STACK_PACK_STRIDE =
-      (1024 / TILE_K) * PACK_ELEMENTS_PER_K_CHUNK;
-
-  constexpr int32_t ROW_PAIRS = (MB + 1) / TILE_ROWS;
-  const int32_t pack_stride =
-      runtime_k ? ((K_val + TILE_K - 1) / TILE_K) * PACK_ELEMENTS_PER_K_CHUNK
-                : (K_static / TILE_K) * PACK_ELEMENTS_PER_K_CHUNK;
-
-  alignas(64) c10::BFloat16 A_packed_stack[ROW_PAIRS * STACK_PACK_STRIDE];
-  std::vector<c10::BFloat16> A_packed_heap;
-  c10::BFloat16* A_packed =
-      (pack_stride <= STACK_PACK_STRIDE)
-          ? A_packed_stack
-          : (A_packed_heap.resize(ROW_PAIRS * pack_stride),
-             A_packed_heap.data());
-
-  for (int32_t rp = 0; rp < ROW_PAIRS; ++rp) {
-    const int32_t m = rp * TILE_ROWS;
-    const int32_t m_rows = (m + 1 < MB) ? TILE_ROWS : 1;
-    const c10::BFloat16* A0 = A + m * lda;
-    const c10::BFloat16* A1 = (m_rows == TILE_ROWS) ? (A + (m + 1) * lda) : A0;
-    reshape_Q_2xK_for_bfmmla(A0, A1, A_packed + rp * pack_stride, K_val);
-  }
-
-  for (int32_t n = 0; n < N; n += OUTPUT_COLS_PER_BLOCK) {
-    const c10::BFloat16* B_blk_c10 =
-        (phase == AttentionGemmPhase::PV)
-            ? (B + (n / TILE_COLS) * b_layout_stride)
-            : (B + (n / OUTPUT_COLS_PER_BLOCK) * b_layout_stride);
-    const bfloat16_t* B_blk = reinterpret_cast<const bfloat16_t*>(B_blk_c10);
-
-    // Process row-pairs in groups of 4, 2, then 1
-    int32_t row_pair_idx = 0;
-
-#define PROCESS_RP_GROUP(group_size)                                       \
-  for (; row_pair_idx + (group_size - 1) < ROW_PAIRS;                      \
-       row_pair_idx += group_size) {                                       \
-    const bfloat16_t* Ap[group_size];                                      \
-    int32_t mr[group_size];                                                \
-    for (int32_t i = 0; i < group_size; ++i) {                             \
-      Ap[i] = reinterpret_cast<const bfloat16_t*>(                         \
-          A_packed + (row_pair_idx + i) * pack_stride);                    \
-      mr[i] = (((row_pair_idx + i) * TILE_ROWS + 1) < MB) ? TILE_ROWS : 1; \
-    }                                                                      \
-    float* C_blk = C + (row_pair_idx * TILE_ROWS) * ldc + n;               \
-    if constexpr (runtime_k) {                                             \
-      gemm_rowpairs_x8_bfmmla_neon<group_size, -1, phase>(                 \
-          Ap, mr, B_blk, C_blk, ldc, accumulate, b_layout_stride, K_val);  \
-    } else {                                                               \
-      gemm_rowpairs_x8_bfmmla_neon<group_size, K_static, phase>(           \
-          Ap, mr, B_blk, C_blk, ldc, accumulate,                           \
-          (phase == AttentionGemmPhase::PV) ? b_layout_stride              \
-                                            : b_reduction_stride);         \
-    }                                                                      \
-  }
-
-    PROCESS_RP_GROUP(4);
-    PROCESS_RP_GROUP(2);
-    PROCESS_RP_GROUP(1);
-#undef PROCESS_RP_GROUP
-  }
-}
-
-// Macro-kernel: iterates over M in MB={8,4,2,1} chunks.
-// Supports compile-time K specialization when K >= 0; otherwise uses runtime K
-// (runtime K path is only supported for PV).
-template <AttentionGemmPhase phase, int32_t N, int32_t K = -1>
-FORCE_INLINE void gemm_macro_neon_bfmmla(
-    const c10::BFloat16* __restrict A, const c10::BFloat16* __restrict B,
-    float* __restrict C, int32_t M, int32_t K_runtime, int64_t lda, int64_t ldc,
-    int64_t b_layout_stride, int64_t b_reduction_stride, bool accumulate) {
-  static_assert(N % OUTPUT_COLS_PER_BLOCK == 0,
-                "N must be a multiple of OUTPUT_COLS_PER_BLOCK");
-
-  if constexpr (K >= 0) {
-    static_assert(K % TILE_K == 0, "K must be divisible by TILE_K");
-    for (int32_t m = 0; m < M;) {
-      const int32_t rem = M - m;
-      const c10::BFloat16* A_blk = A + m * lda;
-      float* C_blk = C + m * ldc;
-
-#define DISPATCH_MB(mb)                                                   \
-  gemm_packA_compute_MB_xN<mb, N, K, phase>(A_blk, B, C_blk, 0, lda, ldc, \
-                                            b_layout_stride,              \
-                                            b_reduction_stride, accumulate)
-
-      if (rem >= 8) {
-        DISPATCH_MB(8);
-        m += 8;
-      } else if (rem >= 4) {
-        DISPATCH_MB(4);
-        m += 4;
-      } else if (rem >= 2) {
-        DISPATCH_MB(2);
-        m += 2;
-      } else {
-        DISPATCH_MB(1);
-        m += 1;
-      }
-#undef DISPATCH_MB
-    }
-  } else {
-    static_assert(phase == AttentionGemmPhase::PV,
-                  "Runtime K specialization only supported for PV.");
-    const int32_t K_val = K_runtime;
-
-    for (int32_t m = 0; m < M;) {
-      const int32_t rem = M - m;
-      const c10::BFloat16* A_blk = A + m * lda;
-      float* C_blk = C + m * ldc;
-
-#define DISPATCH_MB_RUNTIME(mb)                                                \
-  gemm_packA_compute_MB_xN<mb, N, -1, phase>(A_blk, B, C_blk, K_val, lda, ldc, \
-                                             b_layout_stride,                  \
-                                             b_reduction_stride, accumulate)
-
-      if (rem >= 8) {
-        DISPATCH_MB_RUNTIME(8);
-        m += 8;
-      } else if (rem >= 4) {
-        DISPATCH_MB_RUNTIME(4);
-        m += 4;
-      } else if (rem >= 2) {
-        DISPATCH_MB_RUNTIME(2);
-        m += 2;
-      } else {
-        DISPATCH_MB_RUNTIME(1);
-        m += 1;
-      }
-#undef DISPATCH_MB_RUNTIME
-    }
-  }
-}
-
-#undef INIT_ACC_ROWPAIR_4
-#undef STORE_ACC_ROWPAIR_4
-#undef BFMMLA_COMPUTE_4
-
-}  // namespace
-
-// TileGemm Adapter for Attention
-
-template <typename kv_cache_t, int32_t BlockTokens, int32_t HeadDim>
-class TileGemmNEONBFMMLA {
- public:
-  template <AttentionGemmPhase phase, int32_t head_dim_ct>
-  FORCE_INLINE static void gemm(const int32_t m_size, void* __restrict__ a_tile,
-                                kv_cache_t* __restrict__ b_tile,
-                                float* __restrict__ c_tile, const int64_t lda,
-                                [[maybe_unused]] const int64_t ldb,
-                                const int64_t ldc,
-                                [[maybe_unused]] const int32_t block_size,
-                                [[maybe_unused]] const int32_t dynamic_k_size,
-                                const bool accum_c) {
-    static_assert(BlockTokens % OUTPUT_COLS_PER_BLOCK == 0);
-    // BFMMLA kernels require compile-time head_dim; keep head_dim_ct only for
-    // API parity with other tile_gemm implementations.
-    if constexpr (head_dim_ct >= 0) {
-      static_assert(head_dim_ct == HeadDim,
-                    "BFMMLA expects head_dim_ct to match HeadDim; PV passes "
-                    "-1 for API parity.");
-    }
-
-    if constexpr (phase == AttentionGemmPhase::QK) {
-      const int64_t b_reduction_stride = K_INNER_STRIDE;
-      const int64_t b_token_block_stride = (HeadDim / TILE_K) * K_INNER_STRIDE;
-
-      gemm_macro_neon_bfmmla<AttentionGemmPhase::QK, BlockTokens, HeadDim>(
-          reinterpret_cast<const c10::BFloat16*>(a_tile), b_tile, c_tile,
-          m_size, 0, lda, ldc, b_token_block_stride, b_reduction_stride,
-          accum_c);
-    } else {
-      const int64_t b_pair_stride =
-          (block_size / V_TOKENS_PER_ROW_BLOCK) * V_INNER_STRIDE;
-
-      // PV gemm with runtime K specialization
-      switch (dynamic_k_size) {
-        case 32:
-          gemm_macro_neon_bfmmla<AttentionGemmPhase::PV, HeadDim, 32>(
-              reinterpret_cast<const c10::BFloat16*>(a_tile), b_tile, c_tile,
-              m_size, 32, lda, ldc, b_pair_stride, 0, accum_c);
-          break;
-        case 128:
-          gemm_macro_neon_bfmmla<AttentionGemmPhase::PV, HeadDim, 128>(
-              reinterpret_cast<const c10::BFloat16*>(a_tile), b_tile, c_tile,
-              m_size, 128, lda, ldc, b_pair_stride, 0, accum_c);
-          break;
-        case 256:
-          gemm_macro_neon_bfmmla<AttentionGemmPhase::PV, HeadDim, 256>(
-              reinterpret_cast<const c10::BFloat16*>(a_tile), b_tile, c_tile,
-              m_size, 256, lda, ldc, b_pair_stride, 0, accum_c);
-          break;
-        default:
-          gemm_macro_neon_bfmmla<AttentionGemmPhase::PV, HeadDim>(
-              reinterpret_cast<const c10::BFloat16*>(a_tile), b_tile, c_tile,
-              m_size, dynamic_k_size, lda, ldc, b_pair_stride, 0, accum_c);
-          break;
-      }
-    }
-  }
-};
-
-// Shared ASIMD BFMMLA implementation (BF16 only). The block size alignment and
-// ISA tag are template parameters so we can reuse the same kernels for
-// different NEON configurations.
-template <int64_t block_size_alignment, ISA isa_type, int64_t head_dim>
-class AttentionImplNEONBFMMLA {
- public:
-  using query_t = c10::BFloat16;
-  using q_buffer_t = c10::BFloat16;
-  using kv_cache_t = c10::BFloat16;
-  using logits_buffer_t = float;
-  using partial_output_buffer_t = float;
-  using prob_buffer_t = c10::BFloat16;
-
-  static constexpr int64_t BlockSizeAlignment = block_size_alignment;
-  // HeadDimAlignment equals head_dim so that the PV phase processes
-  // the full head dimension in a single gemm call.
-  static constexpr int64_t HeadDimAlignment = head_dim;
-  static constexpr int64_t MaxQHeadNumPerIteration = 16;
-  static constexpr int64_t HeadDim = head_dim;
-  static constexpr ISA ISAType = isa_type;
-  static constexpr bool scale_on_logits = false;
-
-  static_assert(HeadDim % OUTPUT_COLS_PER_BLOCK == 0);
-  static_assert(BlockSizeAlignment % OUTPUT_COLS_PER_BLOCK == 0);
-  static_assert(HeadDim % TILE_K == 0, "HeadDim must be a multiple of TILE_K");
-
- public:
-  template <template <typename tile_gemm_t> typename attention>
-  FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) {
-    attention<
-        TileGemmNEONBFMMLA<kv_cache_t, static_cast<int32_t>(BlockSizeAlignment),
-                           static_cast<int32_t>(HeadDim)>>
-        attention_iteration;
-    attention_iteration(CPU_ATTENTION_PARAMS);
-  }
-
-  // Key cache stride per token group (TokenColumn layout; QK)
-  static constexpr int64_t k_cache_token_group_stride(
-      [[maybe_unused]] const int32_t block_size) {
-    static_assert(BlockSizeAlignment % K_TOKENS_PER_GROUP == 0);
-    return (BlockSizeAlignment / K_TOKENS_PER_GROUP) *
-           ((head_dim / TILE_K) * K_INNER_STRIDE);
-  }
-
-  // Value cache stride per token group (TokenRow layout; PV)
-  static constexpr int64_t v_cache_token_group_stride(
-      [[maybe_unused]] const int32_t block_size) {
-    static_assert(BlockSizeAlignment % V_TOKENS_PER_ROW_BLOCK == 0);
-    return (BlockSizeAlignment / V_TOKENS_PER_ROW_BLOCK) * V_INNER_STRIDE;
-  }
-
-  // The stride to move to the "next" head_dim group
-  // is the full V cache size per head, since HeadDimAlignment == head_dim.
-  // Hence, the stride is not used in this case
-  static constexpr int64_t v_cache_head_group_stride(
-      [[maybe_unused]] const int32_t block_size) {
-    return head_dim * block_size;
-  }
-
-  // Convert Q heads to BF16 and apply scale factor using native BF16 intrinsics
-  static void copy_q_heads_tile(c10::BFloat16* __restrict__ src,
-                                c10::BFloat16* __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) {
-    constexpr int32_t dim = static_cast<int32_t>(head_dim);
-    const float32x4_t scale_vec = vdupq_n_f32(scale);
-
-    for (int32_t qi = 0; qi < q_num; ++qi) {
-      for (int32_t hi = 0; hi < q_heads_per_kv; ++hi) {
-        c10::BFloat16* __restrict__ curr_q =
-            src + qi * q_num_stride + hi * q_head_stride;
-        c10::BFloat16* __restrict__ dst =
-            q_buffer + qi * q_heads_per_kv * head_dim + hi * head_dim;
-
-        for (int32_t i = 0; i < dim; i += OUTPUT_COLS_PER_BLOCK) {
-          bfloat16x8_t in8 =
-              vld1q_bf16(reinterpret_cast<const bfloat16_t*>(curr_q + i));
-          float32x4_t lo = vmulq_f32(vcvtq_low_f32_bf16(in8), scale_vec);
-          float32x4_t hi = vmulq_f32(vcvtq_high_f32_bf16(in8), scale_vec);
-
-          bfloat16x4_t lo_b = vcvt_bf16_f32(lo);
-          bfloat16x4_t hi_b = vcvt_bf16_f32(hi);
-          bfloat16x8_t out = vcombine_bf16(lo_b, hi_b);
-          vst1q_bf16(reinterpret_cast<bfloat16_t*>(dst + i), out);
-        }
-      }
-    }
-  }
-
- public:
-  // Reshape and cache K/V into BFMMLA-optimized layouts
-  // K cache:
-  // [block_size/K_TOKENS_PER_GROUP][head_dim/TILE_K][K_INNER_STRIDE]
-  // - TokenColumn
-  // V cache:
-  // [head_dim/TILE_COLS][block_size/V_TOKENS_PER_ROW_BLOCK][V_INNER_STRIDE]
-  // - TokenRows
-  static void reshape_and_cache(
-      const c10::BFloat16* __restrict__ key,
-      const c10::BFloat16* __restrict__ value,
-      c10::BFloat16* __restrict__ key_cache,
-      c10::BFloat16* __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,
-      [[maybe_unused]] const int64_t num_blocks,
-      const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
-      const int64_t block_size,
-      [[maybe_unused]] const int64_t block_size_stride) {
-    const int64_t k_block_stride = (head_dim / TILE_K) * K_INNER_STRIDE;
-    const int64_t v_pair_stride =
-        (block_size / V_TOKENS_PER_ROW_BLOCK) * V_INNER_STRIDE;
-
-#pragma omp parallel for
-    for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
-      for (int64_t token_idx = 0; token_idx < token_num; ++token_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;
-
-        // Key cache: TokenColumn QK
-        {
-          const c10::BFloat16* __restrict key_src =
-              key + token_idx * key_token_num_stride +
-              head_idx * key_head_num_stride;
-
-          c10::BFloat16* __restrict key_base = key_cache +
-                                               block_idx * num_blocks_stride +
-                                               head_idx * cache_head_num_stride;
-
-          const int64_t block_in_block = block_offset / K_TOKENS_PER_GROUP;
-          const int64_t pair_in_block =
-              (block_offset % K_TOKENS_PER_GROUP) / TILE_COLS;
-          const int64_t lane_base = (block_offset & 1) ? TILE_K : 0;
-
-          c10::BFloat16* __restrict block_base =
-              key_base + block_in_block * k_block_stride;
-
-          for (int64_t hd4 = 0; hd4 < head_dim / TILE_K; ++hd4) {
-            uint16_t* dst_u16 = reinterpret_cast<uint16_t*>(
-                block_base + hd4 * K_INNER_STRIDE +
-                pair_in_block * V_INNER_STRIDE + lane_base);
-            const uint16_t* src_u16 =
-                reinterpret_cast<const uint16_t*>(key_src + hd4 * TILE_K);
-            vst1_u16(dst_u16, vld1_u16(src_u16));
-          }
-        }
-
-        // Value cache: TokenRow PV
-        {
-          const c10::BFloat16* __restrict value_src =
-              value + token_idx * value_token_num_stride +
-              head_idx * value_head_num_stride;
-
-          c10::BFloat16* __restrict value_base =
-              value_cache + block_idx * num_blocks_stride +
-              head_idx * cache_head_num_stride;
-
-          const int64_t row_block = block_offset / V_TOKENS_PER_ROW_BLOCK;
-          const int64_t lane = block_offset & (V_TOKENS_PER_ROW_BLOCK - 1);
-
-          c10::BFloat16* __restrict row_block_base =
-              value_base + row_block * V_INNER_STRIDE;
-
-          for (int64_t hd2 = 0; hd2 < head_dim / TILE_COLS; ++hd2) {
-            c10::BFloat16* __restrict dst_val =
-                row_block_base + hd2 * v_pair_stride;
-
-            const uint16_t* src_u16 =
-                reinterpret_cast<const uint16_t*>(value_src);
-            uint16_t* dst_u16 = reinterpret_cast<uint16_t*>(dst_val);
-            dst_u16[lane] = src_u16[hd2 * TILE_COLS + 0];
-            dst_u16[lane + V_TOKENS_PER_ROW_BLOCK] =
-                src_u16[hd2 * TILE_COLS + 1];
-          }
-        }
-      }
-    }
-  }
-};
-
-}  // namespace cpu_attention
-
-#endif  // CPU_ATTN_ASIMD_BFMMLA_HPP

csrc/cpu/cpu_fused_moe.cpp

@@ -147,7 +147,7 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input,
                     const int32_t token_num, const int32_t expert_num,
                     const int32_t topk_num, const int32_t input_size_13,
                     const int32_t output_size_13, const int32_t input_size_2,
-                    const int32_t output_size_2, const bool skip_weighted) {
+                    const int32_t output_size_2) {
   using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
   constexpr int32_t gemm_n_tile_size = gemm_t::NSize;
   constexpr int32_t gemm_m_tile_size = gemm_t::MaxMSize;
@@ -582,11 +582,6 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input,
         scalar_t* __restrict__ curr_output_buffer =
             output + token_id * output_size_2;
 
-        if (skip_weighted) {
-          // Only for topk_num == 1
-          *curr_weight = 1.0f;
-        }
-
         if (topk_num > 1) {
           {
             int32_t w2_output_idx = curr_expand_token_id_index_buffer[0];
@@ -704,7 +699,7 @@ void cpu_fused_moe(
     const std::optional<torch::Tensor>& w2_bias,  // [expert_num, output_size_2]
     const torch::Tensor& topk_weights,            // [token_num, k], float32
     const torch::Tensor& topk_id,                 // [token_num, k], int32
-    const bool skip_weighted, const std::string& act, const std::string& isa) {
+    const std::string& act, const std::string& isa) {
   const int32_t token_num = input.size(0);
   const int32_t input_size_13 = input.size(1);
   const int64_t input_stride = input.stride(0);
@@ -716,8 +711,6 @@ void cpu_fused_moe(
   const int32_t topk_num = topk_id.size(1);
   const FusedMOEAct act_type = get_act_type(act);
   cpu_utils::ISA isa_type = cpu_utils::get_isa(isa);
-  TORCH_CHECK(!skip_weighted || topk_num == 1,
-              "skip_weighted is only supported for topk=1 on CPU");
 
   VLLM_DISPATCH_FLOATING_TYPES(w13.scalar_type(), "cpu_fused_moe", [&]() {
     CPU_ISA_DISPATCH_IMPL(isa_type, [&]() {
@@ -728,7 +721,7 @@ void cpu_fused_moe(
           w2_bias.has_value() ? w2_bias->data_ptr<scalar_t>() : nullptr,
           topk_weights.data_ptr<float>(), topk_id.data_ptr<int32_t>(), act_type,
           token_num, expert_num, topk_num, input_size_13, output_size_13,
-          input_size_2, output_size_2, skip_weighted);
+          input_size_2, output_size_2);
     });
   });
 }

csrc/cpu/cpu_types_vxe.hpp

@@ -16,12 +16,10 @@ namespace vec_op {
 #define vec_sr(a, b) ((a) >> (b))  // Vector Shift Right Algebraic
 #define vec_sl(a, b) ((a) << (b))  // Vector Shift Left
 
-// NOTE: FP16 (Half) is supported on s390x via custom bit-manipulation
-// conversion. PyTorch itself lacks native s390x FP16 support.
-#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...)            \
-  AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__)    \
-  AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) \
-  AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)
+// FIXME: FP16 is not fully supported in Torch-CPU
+#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...)         \
+  AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
+  AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
 
 #define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
   AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
@@ -88,39 +86,6 @@ struct BF16Vec8 : public Vec<BF16Vec8> {
   }
 };
 
-struct FP16Vec8 : public Vec<FP16Vec8> {
-  constexpr static int VEC_ELEM_NUM = 8;
-
-  __vector signed short reg;
-
-  explicit FP16Vec8(const void* ptr) : reg(*(__vector signed short*)ptr) {}
-  explicit FP16Vec8(const FP32Vec8&);
-
-  void save(void* ptr) const {
-    *reinterpret_cast<__vector signed short*>(ptr) = reg;
-  }
-};
-
-struct FP16Vec16 : public Vec<FP16Vec16> {
-  constexpr static int VEC_ELEM_NUM = 16;
-
-  ss16x8x2_t reg;
-
-  explicit FP16Vec16(const void* ptr) {
-    // Load 256 bits (16 FP16 values) in two parts
-    reg.val[0] = (__vector signed short)vec_xl(0, (signed short*)ptr);
-    reg.val[1] = (__vector signed short)vec_xl(16, (signed short*)ptr);
-  }
-
-  explicit FP16Vec16(const FP32Vec16&);
-
-  void save(void* ptr) const {
-    // Save 256 bits in two parts
-    vec_xst(reg.val[0], 0, (signed short*)ptr);
-    vec_xst(reg.val[1], 16, (signed short*)ptr);
-  }
-};
-
 struct BF16Vec16 : public Vec<BF16Vec16> {
   constexpr static int VEC_ELEM_NUM = 16;
 
@@ -143,92 +108,6 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
 
 const static __vector signed short zero = vec_splats((signed short)0);
 
-FORCE_INLINE __vector float fp16_to_fp32_bits(__vector unsigned int x) {
-  const __vector unsigned int mask_sign = {0x8000, 0x8000, 0x8000, 0x8000};
-  const __vector unsigned int mask_exp = {0x7C00, 0x7C00, 0x7C00, 0x7C00};
-  const __vector unsigned int mask_mant = {0x03FF, 0x03FF, 0x03FF, 0x03FF};
-  const __vector unsigned int bias_adj = {112, 112, 112, 112};
-  const __vector unsigned int exp_max_fp16 = {0x1F, 0x1F, 0x1F,
-                                              0x1F};  // FP16 NaN/Inf exponent
-  const __vector unsigned int exp_max_fp32 = {0xFF, 0xFF, 0xFF,
-                                              0xFF};  // FP32 NaN/Inf exponent
-
-  __vector unsigned int s = (x & mask_sign) << 16;
-  __vector unsigned int e = (x & mask_exp) >> 10;
-  __vector unsigned int m = (x & mask_mant) << 13;
-
-  // Check for NaN/Inf: exponent = 0x1F in FP16
-  __vector __bool int is_nan_inf = vec_cmpeq(e, exp_max_fp16);
-
-  // Normal: adjust bias; NaN/Inf: set to 0xFF
-  __vector unsigned int e_normal = e + bias_adj;
-  e = vec_sel(e_normal, exp_max_fp32, is_nan_inf);
-
-  return (__vector float)(s | (e << 23) | m);
-}
-
-FORCE_INLINE __vector unsigned int fp32_to_fp16_bits(__vector float f_in) {
-  __vector unsigned int in = (__vector unsigned int)f_in;
-
-  const __vector unsigned int mask_sign_32 = {0x80000000, 0x80000000,
-                                              0x80000000, 0x80000000};
-  const __vector unsigned int mask_exp_32 = {0x7F800000, 0x7F800000, 0x7F800000,
-                                             0x7F800000};
-  const __vector unsigned int mask_mant_32 = {0x007FFFFF, 0x007FFFFF,
-                                              0x007FFFFF, 0x007FFFFF};
-
-  // Use SIGNED integers for exponent math to handle underflow check
-  const __vector signed int bias_adj = {112, 112, 112, 112};
-  const __vector signed int zero = {0, 0, 0, 0};
-  const __vector signed int max_exp = {31, 31, 31, 31};  // Max FP16 exp
-  const __vector unsigned int exp_max_fp32 = {0xFF, 0xFF, 0xFF, 0xFF};
-  const __vector unsigned int exp_max_fp16 = {0x1F, 0x1F, 0x1F, 0x1F};
-
-  __vector unsigned int s = (in & mask_sign_32) >> 16;
-  __vector unsigned int e_u = (in & mask_exp_32) >> 23;
-
-  // Check for NaN/Inf: exponent = 0xFF in FP32
-  __vector __bool int is_nan_inf = vec_cmpeq(e_u, exp_max_fp32);
-
-  __vector signed int e_s = (__vector signed int)e_u;
-  e_s = vec_sub(e_s, bias_adj);
-  e_s = vec_max(e_s, zero);
-  e_s = vec_min(e_s, max_exp);
-  __vector unsigned int e_normal = (__vector unsigned int)e_s;
-
-  __vector unsigned int e_final = vec_sel(e_normal, exp_max_fp16, is_nan_inf);
-
-  const __vector unsigned int one_v = {1, 1, 1, 1};
-  const __vector unsigned int mask_sticky = {0xFFF, 0xFFF, 0xFFF, 0xFFF};
-
-  __vector unsigned int round_bit = (in >> 12) & one_v;
-  __vector unsigned int sticky = in & mask_sticky;
-  __vector unsigned int m = (in & mask_mant_32) >> 13;
-  __vector unsigned int lsb = m & one_v;  // LSB of mantissa for tie-breaking
-
-  // Round up if: round_bit && (sticky || lsb)
-  __vector __bool int sticky_nonzero =
-      vec_cmpgt(sticky, (__vector unsigned int){0, 0, 0, 0});
-  __vector __bool int lsb_set = vec_cmpeq(lsb, one_v);
-  __vector __bool int round_up =
-      vec_and(vec_cmpeq(round_bit, one_v), vec_or(sticky_nonzero, lsb_set));
-
-  m = vec_sel(m, m + one_v, round_up);
-
-  const __vector unsigned int mant_mask = {0x3FF, 0x3FF, 0x3FF, 0x3FF};
-  const __vector unsigned int max_normal_exp = {0x1E, 0x1E, 0x1E, 0x1E};
-  __vector __bool int mant_overflows = vec_cmpgt(m, mant_mask);
-  __vector __bool int would_overflow_to_inf =
-      vec_and(mant_overflows, vec_cmpeq(e_final, max_normal_exp));
-  __vector unsigned int e_inc = vec_min(e_final + one_v, exp_max_fp16);
-  e_final = vec_sel(e_final, e_inc, mant_overflows);
-  m = vec_and(m, mant_mask);
-  e_final = vec_sel(e_final, max_normal_exp, would_overflow_to_inf);
-  m = vec_sel(m, mant_mask, would_overflow_to_inf);
-
-  return s | (e_final << 10) | m;
-}
-
 struct BF16Vec32 : public Vec<BF16Vec32> {
   constexpr static int VEC_ELEM_NUM = 32;
 
@@ -301,18 +180,6 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
     reg.val[1] = (__vector float)vec_mergel(v.reg, zero);
   }
 
-  explicit FP32Vec8(const FP16Vec8& v) {
-    // Cast to UNSIGNED short vector to prevent sign-extension during unpack
-    __vector unsigned short raw_u = (__vector unsigned short)v.reg;
-
-    // Unpack 8x16-bit to two 4x32-bit vectors (Zero extended)
-    __vector unsigned int raw_hi = (__vector unsigned int)vec_unpackh(raw_u);
-    __vector unsigned int raw_lo = (__vector unsigned int)vec_unpackl(raw_u);
-
-    reg.val[0] = fp16_to_fp32_bits(raw_hi);
-    reg.val[1] = fp16_to_fp32_bits(raw_lo);
-  }
-
   float reduce_sum() const {
     AliasReg ar;
     ar.reg = reg;
@@ -664,22 +531,6 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
     reg.val[3] = (__vector float)vec_mergel(v.reg.val[1], zero);
   }
 
-  explicit FP32Vec16(const FP16Vec16& v) {
-    __vector unsigned int raw_hi_0 =
-        (__vector unsigned int)vec_unpackh(v.reg.val[0]);
-    __vector unsigned int raw_lo_0 =
-        (__vector unsigned int)vec_unpackl(v.reg.val[0]);
-    reg.val[0] = fp16_to_fp32_bits(raw_hi_0);
-    reg.val[1] = fp16_to_fp32_bits(raw_lo_0);
-
-    __vector unsigned int raw_hi_1 =
-        (__vector unsigned int)vec_unpackh(v.reg.val[1]);
-    __vector unsigned int raw_lo_1 =
-        (__vector unsigned int)vec_unpackl(v.reg.val[1]);
-    reg.val[2] = fp16_to_fp32_bits(raw_hi_1);
-    reg.val[3] = fp16_to_fp32_bits(raw_lo_1);
-  }
-
   explicit FP32Vec16(const BF16Vec8& v) : FP32Vec16(FP32Vec8(v)) {}
 
   FP32Vec16 operator*(const FP32Vec16& b) const {
@@ -777,10 +628,8 @@ struct VecType<c10::BFloat16> {
   using vec_type = BF16Vec8;
 };
 
-template <>
-struct VecType<c10::Half> {
-  using vec_type = FP16Vec8;
-};
+// On s390x, FP16 (Half) is not natively supported, use FP32 vectors instead
+using FP16Vec16 = FP32Vec16;
 
 template <typename T>
 void storeFP32(float v, T* ptr) {
@@ -801,52 +650,6 @@ inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
   *ptr = *(v_ptr + 1);
 }
 
-template <>
-inline void storeFP32<::c10::Half>(float v, ::c10::Half* ptr) {
-  // Use bit-manipulation for IEEE FP32 to FP16 conversion since vector
-  // intrinsics for FP32 to FP16 conversion does not use IEEE rounding and can
-  // produce incorrect results for some inputs. Process each of the 4 vectors
-  // separately.
-  uint32_t in;
-  std::memcpy(&in, &v, sizeof(in));
-
-  uint32_t s = (in & 0x80000000) >> 16;  // Sign
-  uint32_t e = (in & 0x7F800000) >> 23;  // Exponent
-  uint32_t round_bit = (in >> 12) & 1;
-  uint32_t sticky = (in & 0xFFF) != 0;  // Any bits in [11..0]
-  uint32_t m = (in & 0x007FFFFF) >> 13;
-  uint32_t lsb = m & 1;  // LSB of mantissa for tie-breaking
-
-  // Check for NaN/Inf before rounding
-  bool is_nan_inf = (e == 0xFF);
-
-  if (round_bit && (sticky || lsb)) {
-    m++;
-    // Handle mantissa overflow: if m overflows 10 bits, increment exponent
-    if (m > 0x3FF) {
-      m = 0;
-      e++;
-    }
-  }
-
-  if (is_nan_inf) {
-    // NaN/Inf: preserve it
-    e = 0x1F;
-  } else {
-    // Normal: adjust bias (127 - 15), flush subnormals to zero
-    e = (e >= 112) ? (e - 112) : 0;
-    // If exponent overflows to Inf range, saturate to max normal FP16 value
-    if (e > 0x1E) {
-      e = 0x1E;   // Max normal exponent
-      m = 0x3FF;  // Max mantissa
-    }
-  }
-
-  uint16_t fp16 = (uint16_t)(s | (e << 10) | m);
-
-  *reinterpret_cast<uint16_t*>(ptr) = fp16;
-}
-
 #ifndef __VEC_CLASS_FP_NAN
   #define __VEC_CLASS_FP_NAN (1 << 6)
 #endif
@@ -1000,44 +803,6 @@ inline BF16Vec16::BF16Vec16(const FP32Vec16& v) {
   reg.val[1] = (__vector signed short)vec_perm(inp2, inp3, omask);
 }
 
-inline FP16Vec8::FP16Vec8(const FP32Vec8& v) {
-  // Use bit-manipulation for IEEE FP32 to FP16 conversion since vector
-  // intrinsics for FP32 to FP16 conversion does not use IEEE rounding and can
-  // produce incorrect results for some inputs. Process each of the 4 vectors
-  // separately.
-  __vector unsigned int res_hi = fp32_to_fp16_bits(v.reg.val[0]);
-  __vector unsigned int res_lo = fp32_to_fp16_bits(v.reg.val[1]);
-
-  const __vector unsigned char perm_pack = {
-      2,  3,  6,  7,  10, 11, 14, 15,  // Select lower 2 bytes from res_hi
-      18, 19, 22, 23, 26, 27, 30, 31   // Select lower 2 bytes from res_lo
-  };
-
-  reg = vec_perm((__vector signed short)res_hi, (__vector signed short)res_lo,
-                 perm_pack);
-}
-
-inline FP16Vec16::FP16Vec16(const FP32Vec16& v) {
-  // Use bit-manipulation for IEEE FP32 to FP16 conversion since vector
-  // intrinsics for FP32 to FP16 conversion does not use IEEE rounding and can
-  // produce incorrect results for some inputs. Process each of the 4 vectors
-  // separately.
-  __vector unsigned int res_0 = fp32_to_fp16_bits(v.reg.val[0]);
-  __vector unsigned int res_1 = fp32_to_fp16_bits(v.reg.val[1]);
-  __vector unsigned int res_2 = fp32_to_fp16_bits(v.reg.val[2]);
-  __vector unsigned int res_3 = fp32_to_fp16_bits(v.reg.val[3]);
-
-  const __vector unsigned char perm_pack = {
-      2,  3,  6,  7,  10, 11, 14, 15,  // Lower 2 bytes from first vector
-      18, 19, 22, 23, 26, 27, 30, 31   // Lower 2 bytes from second vector
-  };
-
-  reg.val[0] = vec_perm((__vector signed short)res_0,
-                        (__vector signed short)res_1, perm_pack);
-  reg.val[1] = vec_perm((__vector signed short)res_2,
-                        (__vector signed short)res_3, perm_pack);
-}
-
 // 1D softmax over `n` elements in `input`, writes result to `output`.
 // Uses FP32Vec8 for main body, scalar tail handling.
 // Requirement: n > 0

csrc/cpu/cpu_wna16.cpp

@@ -116,7 +116,7 @@ class Dequantizer4b {
       scalar_vec_t output_vec_0(wb_0);
       scalar_vec_t output_vec_1(wb_1);
 
-      // AMX needs to interleave K elements to pack as 32 bits
+      // AMX needs to interlave K elements to pack as 32 bits
       if constexpr (isa == ISA::AMX) {
         vec_op::interleave_save(output_vec_0, output_vec_1, curr_weight);
       } else {

csrc/cpu/dnnl_helper.cpp

@@ -237,20 +237,12 @@ 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(
-                     MSizeCacheKey{.a_m_size = kProbeM,
+                     MSizeCacheKey{.a_m_size = DNNL_RUNTIME_DIM_VAL,
                                    .use_bias = false,
                                    .bias_type = dnnl::memory::data_type::undef},
-                     /*first_time=*/true)
+                     true)
                      .weights_desc());
   init_runtime_memory_cache(args);
 }