add file

Signed-off-by: Alexander Matveev <amatveev@redhat.com>
aligning
2025-11-26 13:56:58 +00:00 · 2025-11-25 15:31:48 +00:00 · 2025-11-25 15:31:46 +00:00 · 2025-11-25 15:28:46 +00:00 · 2025-11-25 15:28:17 +00:00
1395 changed files with 33055 additions and 80871 deletions
@@ -1,24 +0,0 @@
-name: vllm_ci
-job_dirs:
-  - ".buildkite/test_areas"
-  - ".buildkite/image_build"
-run_all_patterns:
-  - "docker/Dockerfile"
-  - "CMakeLists.txt"
-  - "requirements/common.txt"
-  - "requirements/cuda.txt"
-  - "requirements/build.txt"
-  - "requirements/test.txt"
-  - "setup.py"
-  - "csrc/"
-  - "cmake/"
-run_all_exclude_patterns:
-  - "docker/Dockerfile."
-  - "csrc/cpu/"
-  - "csrc/rocm/"
-  - "cmake/hipify.py"
-  - "cmake/cpu_extension.cmake"
-registries: public.ecr.aws/q9t5s3a7
-repositories:
-  main: "vllm-ci-postmerge-repo"
-  premerge: "vllm-ci-test-repo"
@@ -0,0 +1,46 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import argparse
+import os
+
+template = """<!DOCTYPE html>
+<html>
+    <body>
+    <h1>Links for vLLM</h1/>
+        <a href="../{x86_wheel_html_escaped}">{x86_wheel}</a><br/>
+        <a href="../{arm_wheel_html_escaped}">{arm_wheel}</a><br/>
+    </body>
+</html>
+"""
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--wheel", help="The wheel path.", required=True)
+args = parser.parse_args()
+
+filename = os.path.basename(args.wheel)
+
+with open("index.html", "w") as f:
+    print(f"Generated index.html for {args.wheel}")
+    # sync the abi tag with .buildkite/scripts/upload-wheels.sh
+    if "x86_64" in filename:
+        x86_wheel = filename
+        arm_wheel = filename.replace("x86_64", "aarch64").replace(
+            "manylinux1", "manylinux2014"
+        )
+    elif "aarch64" in filename:
+        x86_wheel = filename.replace("aarch64", "x86_64").replace(
+            "manylinux2014", "manylinux1"
+        )
+        arm_wheel = filename
+    else:
+        raise ValueError(f"Unsupported wheel: {filename}")
+    # cloudfront requires escaping the '+' character
+    f.write(
+        template.format(
+            x86_wheel=x86_wheel,
+            x86_wheel_html_escaped=x86_wheel.replace("+", "%2B"),
+            arm_wheel=arm_wheel,
+            arm_wheel_html_escaped=arm_wheel.replace("+", "%2B"),
+        )
+    )
@@ -1,56 +0,0 @@
-#!/bin/bash
-set -e
-
-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
-
-REGISTRY=$1
-REPO=$2
-BUILDKITE_COMMIT=$3
-BRANCH=$4
-VLLM_USE_PRECOMPILED=$5
-VLLM_MERGE_BASE_COMMIT=$6
-CACHE_FROM=$7
-CACHE_TO=$8
-
-# 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
-
-# docker buildx 
-docker buildx create --name vllm-builder --driver docker-container --use
-docker buildx inspect --bootstrap
-docker buildx ls
-
-# 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
-
-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
-
-# 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 .
@@ -1,57 +0,0 @@
-group: Abuild
-steps:
-  - label: ":docker: Build image"
-    key: image-build
-    depends_on: []
-    commands:
-    - .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
-          limit: 2
-        - exit_status: -10  # Agent was lost
-          limit: 2
-
-  - label: ":docker: Build CPU image"
-    key: image-build-cpu
-    depends_on: []
-    commands:
-    - .buildkite/image_build/image_build_cpu.sh $REGISTRY $REPO $BUILDKITE_COMMIT
-    env:
-      DOCKER_BUILDKIT: "1"
-    retry:
-      automatic:
-        - exit_status: -1  # Agent was lost
-          limit: 2
-        - exit_status: -10  # Agent was lost
-          limit: 2
-
-  - label: ":docker: Build HPU image"
-    soft_fail: true
-    depends_on: []
-    key: image-build-hpu
-    commands:
-    - .buildkite/image_build/image_build_hpu.sh $REGISTRY $REPO $BUILDKITE_COMMIT
-    env:
-      DOCKER_BUILDKIT: "1"
-    retry:
-      automatic:
-        - exit_status: -1  # Agent was lost
-          limit: 2
-        - exit_status: -10  # Agent was lost
-          limit: 2
-  
-  - label: ":docker: Build CPU arm64 image"
-    key: cpu-arm64-image-build
-    depends_on: []
-    optional: true
-    commands:
-    - .buildkite/image_build/image_build_cpu_arm64.sh $REGISTRY $REPO $BUILDKITE_COMMIT
-    env:
-      DOCKER_BUILDKIT: "1"
-    retry:
-      automatic:
-        - exit_status: -1  # Agent was lost
-          limit: 2
-        - exit_status: -10  # Agent was lost
-          limit: 2
@@ -1,36 +0,0 @@
-#!/bin/bash
-set -e
-
-if [[ $# -lt 3 ]]; then
-  echo "Usage: $0 <registry> <repo> <commit>"
-  exit 1
-fi
-
-REGISTRY=$1
-REPO=$2
-BUILDKITE_COMMIT=$3
-
-# authenticate with AWS ECR
-aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin $REGISTRY
-
-# skip build if image already exists
-if [[ -z $(docker manifest inspect $REGISTRY/$REPO:$BUILDKITE_COMMIT-cpu) ]]; then
-  echo "Image not found, proceeding with build..."
-else
-  echo "Image found"
-  exit 0
-fi
-
-# build
-docker build --file docker/Dockerfile.cpu \
-  --build-arg max_jobs=16 \
-  --build-arg buildkite_commit=$BUILDKITE_COMMIT \
-  --build-arg VLLM_CPU_AVX512BF16=true \
-  --build-arg VLLM_CPU_AVX512VNNI=true \
-  --build-arg VLLM_CPU_AMXBF16=true \
-  --tag $REGISTRY/$REPO:$BUILDKITE_COMMIT-cpu \
-  --target vllm-test \
-  --progress plain .
-
-# push
-docker push $REGISTRY/$REPO:$BUILDKITE_COMMIT-cpu
@@ -1,33 +0,0 @@
-#!/bin/bash
-set -e
-
-if [[ $# -lt 3 ]]; then
-  echo "Usage: $0 <registry> <repo> <commit>"
-  exit 1
-fi
-
-REGISTRY=$1
-REPO=$2
-BUILDKITE_COMMIT=$3
-
-# authenticate with AWS ECR
-aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin $REGISTRY
-
-# skip build if image already exists
-if [[ -z $(docker manifest inspect $REGISTRY/$REPO:$BUILDKITE_COMMIT-cpu) ]]; then
-  echo "Image not found, proceeding with build..."
-else
-  echo "Image found"
-  exit 0
-fi
-
-# build
-docker build --file docker/Dockerfile.cpu \
-  --build-arg max_jobs=16 \
-  --build-arg buildkite_commit=$BUILDKITE_COMMIT \
-  --tag $REGISTRY/$REPO:$BUILDKITE_COMMIT-cpu \
-  --target vllm-test \
-  --progress plain .
-
-# push
-docker push $REGISTRY/$REPO:$BUILDKITE_COMMIT-cpu
@@ -1,34 +0,0 @@
-#!/bin/bash
-set -e
-
-if [[ $# -lt 3 ]]; then
-  echo "Usage: $0 <registry> <repo> <commit>"
-  exit 1
-fi
-
-REGISTRY=$1
-REPO=$2
-BUILDKITE_COMMIT=$3
-
-# authenticate with AWS ECR
-aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin $REGISTRY
-
-# skip build if image already exists
-if [[ -z $(docker manifest inspect $REGISTRY/$REPO:$BUILDKITE_COMMIT-hpu) ]]; then
-  echo "Image not found, proceeding with build..."
-else
-  echo "Image found"
-  exit 0
-fi
-
-# build
-docker build \
-  --file tests/pytorch_ci_hud_benchmark/Dockerfile.hpu \
-  --build-arg max_jobs=16 \
-  --build-arg buildkite_commit=$BUILDKITE_COMMIT \
-  --tag $REGISTRY/$REPO:$BUILDKITE_COMMIT-hpu \
-  --progress plain \
-  https://github.com/vllm-project/vllm-gaudi.git
-
-# push
-docker push $REGISTRY/$REPO:$BUILDKITE_COMMIT-hpu
@@ -8,4 +8,3 @@ tasks:
    value: 0.80
 limit: 250 # will run on 250 * 14 subjects = 3500 samples
 num_fewshot: 5
-rtol: 0.05
@@ -1 +0,0 @@
-Meta-Llama-4-Maverick-17B-128E-Instruct-FP8.yaml
@@ -9,40 +9,11 @@ pytest -s -v test_lm_eval_correctness.py \
    --tp-size=1
 """

-import os
-from contextlib import contextmanager
-
 import lm_eval
 import numpy as np
 import yaml

-DEFAULT_RTOL = 0.08
-
-
-@contextmanager
-def scoped_env_vars(new_env: dict[str, str]):
-    if not new_env:
-        # Fast path: nothing to do
-        yield
-        return
-
-    old_values = {}
-    new_keys = []
-
-    try:
-        for key, value in new_env.items():
-            if key in os.environ:
-                old_values[key] = os.environ[key]
-            else:
-                new_keys.append(key)
-            os.environ[key] = str(value)
-        yield
-    finally:
-        # Restore / clean up
-        for key, value in old_values.items():
-            os.environ[key] = value
-        for key in new_keys:
-            os.environ.pop(key, None)
+RTOL = 0.08


 def launch_lm_eval(eval_config, tp_size):
@@ -61,26 +32,23 @@ def launch_lm_eval(eval_config, tp_size):
        f"trust_remote_code={trust_remote_code},"
        f"max_model_len={max_model_len},"
    )
-
-    env_vars = eval_config.get("env_vars", None)
-    with scoped_env_vars(env_vars):
-        results = lm_eval.simple_evaluate(
-            model=backend,
-            model_args=model_args,
-            tasks=[task["name"] for task in eval_config["tasks"]],
-            num_fewshot=eval_config["num_fewshot"],
-            limit=eval_config["limit"],
-            # TODO(yeq): using chat template w/ fewshot_as_multiturn is supposed help
-            # text models. however, this is regressing measured strict-match for
-            # existing text models in CI, so only apply it for mm, or explicitly set
-            apply_chat_template=eval_config.get(
-                "apply_chat_template", backend == "vllm-vlm"
-            ),
-            fewshot_as_multiturn=eval_config.get("fewshot_as_multiturn", False),
-            # Forward decoding and early-stop controls (e.g., max_gen_toks, until=...)
-            gen_kwargs=eval_config.get("gen_kwargs"),
-            batch_size=batch_size,
-        )
+    results = lm_eval.simple_evaluate(
+        model=backend,
+        model_args=model_args,
+        tasks=[task["name"] for task in eval_config["tasks"]],
+        num_fewshot=eval_config["num_fewshot"],
+        limit=eval_config["limit"],
+        # TODO(yeq): using chat template w/ fewshot_as_multiturn is supposed help
+        # text models. however, this is regressing measured strict-match for
+        # existing text models in CI, so only apply it for mm, or explicitly set
+        apply_chat_template=eval_config.get(
+            "apply_chat_template", backend == "vllm-vlm"
+        ),
+        fewshot_as_multiturn=eval_config.get("fewshot_as_multiturn", False),
+        # Forward decoding and early-stop controls (e.g., max_gen_toks, until=...)
+        gen_kwargs=eval_config.get("gen_kwargs"),
+        batch_size=batch_size,
+    )
    return results


@@ -89,8 +57,6 @@ def test_lm_eval_correctness_param(config_filename, tp_size):

    results = launch_lm_eval(eval_config, tp_size)

-    rtol = eval_config.get("rtol", DEFAULT_RTOL)
-
    success = True
    for task in eval_config["tasks"]:
        for metric in task["metrics"]:
@@ -98,9 +64,8 @@ def test_lm_eval_correctness_param(config_filename, tp_size):
            measured_value = results["results"][task["name"]][metric["name"]]
            print(
                f"{task['name']} | {metric['name']}: "
-                f"ground_truth={ground_truth:.3f} | "
-                f"measured={measured_value:.3f} | rtol={rtol}"
+                f"ground_truth={ground_truth} | measured={measured_value}"
            )
-            success = success and np.isclose(ground_truth, measured_value, rtol=rtol)
+            success = success and np.isclose(ground_truth, measured_value, rtol=RTOL)

    assert success
@@ -108,65 +108,6 @@ The number of this test is less stable compared to the delay and latency benchma

 WARNING: The benchmarking script will save json results by itself, so please do not configure `--save-results` or other results-saving-related parameters in `serving-tests.json`.

-#### Default Parameters Field
-
-We can specify default parameters in a JSON field with key `defaults`. Parameters defined in the field are applied globally to all serving tests, and can be overridden in test case fields. Here is an example:
-
-<details>
-<summary> An Example of default parameters field </summary>
-
-```json
-{
-  "defaults": {
-    "qps_list": [
-      "inf"
-    ],
-    "server_environment_variables": {
-      "VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1
-    },
-    "server_parameters": {
-      "tensor_parallel_size": 1,
-      "dtype": "bfloat16",
-      "block_size": 128,
-      "disable_log_stats": "",
-      "load_format": "dummy"
-    },
-    "client_parameters": {
-      "backend": "vllm",
-      "dataset_name": "random",
-      "random-input-len": 128,
-      "random-output-len": 128,
-      "num_prompts": 200,
-      "ignore-eos": ""
-    }
-  },
-  "tests": [
-    {
-      "test_name": "serving_llama3B_tp2_random_128_128",
-      "server_parameters": {
-        "model": "meta-llama/Llama-3.2-3B-Instruct",
-        "tensor_parallel_size": 2,
-      },
-      "client_parameters": {
-        "model": "meta-llama/Llama-3.2-3B-Instruct",
-      }
-    },
-    {
-      "test_name": "serving_qwen3_tp4_random_128_128",
-      "server_parameters": {
-        "model": "Qwen/Qwen3-14B",
-        "tensor_parallel_size": 4,
-      },
-      "client_parameters": {
-        "model": "Qwen/Qwen3-14B",
-      }
-    },
-  ]
-}
-```
-
-</details>
-
 ### Visualizing the results

 The `convert-results-json-to-markdown.py` helps you put the benchmarking results inside a markdown table, by formatting [descriptions.md](performance-benchmarks-descriptions.md) with real benchmarking results.
@@ -110,8 +110,7 @@ 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 '
+  timeout 1200 bash -c '
    until curl -X POST localhost:8000/v1/completions; do
      sleep 1
    done' && return 0 || return 1
@@ -317,44 +316,12 @@ run_throughput_tests() {
 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
-      .[]
-    elif (type == "object" and has("tests")) then
-      # merge the default parameters into each test cases
-      . as $root
-      | ($root.defaults // {}) as $d
-      | ($root.tests // [])[]
-      # default qps / max_concurrency from defaults if missing
-      | .qps_list = (.qps_list // $d.qps_list)
-      | .max_concurrency_list = (.max_concurrency_list // $d.max_concurrency_list)
-      # merge envs / params: test overrides defaults
-      | .server_environment_variables =
-          (($d.server_environment_variables // {}) + (.server_environment_variables // {}))
-      | .server_parameters =
-          (($d.server_parameters // {}) + (.server_parameters // {}))
-      | .client_parameters =
-          (($d.client_parameters // {}) + (.client_parameters // {}))
-    else
-      error("Unsupported serving test file format: must be array or object with .tests")
-    end
-  ' "$serving_test_file" | while read -r params; do
+  jq -c '.[]' "$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
@@ -368,25 +335,20 @@ run_serving_tests() {
      continue
    fi

-    # get client and server arguments (after merged the default parameters)
+    # get client and server arguments
    server_params=$(echo "$params" | jq -r '.server_parameters')
    server_envs=$(echo "$params" | jq -r '.server_environment_variables')
    client_params=$(echo "$params" | jq -r '.client_parameters')
-
    server_args=$(json2args "$server_params")
    server_envs=$(json2envs "$server_envs")
    client_args=$(json2args "$client_params")
-
-    # qps_list
    qps_list=$(echo "$params" | jq -r '.qps_list')
    qps_list=$(echo "$qps_list" | jq -r '.[] | @sh')
    echo "Running over qps list $qps_list"
-
-    # max_concurrency_list (fallback to num_prompts if missing)
    max_concurrency_list=$(echo "$params" | jq -r '.max_concurrency_list')
    if [[ -z "$max_concurrency_list" || "$max_concurrency_list" == "null" ]]; then
-      num_prompts=$(echo "$client_params" | jq -r '.num_prompts')
-      max_concurrency_list="[$num_prompts]"
+        num_prompts=$(echo "$client_params" | jq -r '.num_prompts')
+        max_concurrency_list="[$num_prompts]"
    fi
    max_concurrency_list=$(echo "$max_concurrency_list" | jq -r '.[] | @sh')
    echo "Running over max concurrency list $max_concurrency_list"
@@ -0,0 +1,610 @@
+[
+    {
+        "test_name": "serving_llama8B_bf16_tp1_sharegpt",
+        "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": "",
+	    "enforce_eager": "",
+	    "max_num_batched_tokens": 2048,
+	    "max_num_seqs": 256,
+            "load_format": "dummy"
+        },
+        "client_parameters": {
+            "model": "meta-llama/Llama-3.1-8B-Instruct",
+            "backend": "vllm",
+            "dataset_name": "sharegpt",
+            "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json",
+            "num_prompts": 200
+        }
+    },
+    {
+        "test_name": "serving_llama8B_bf16_tp2_sharegpt",
+        "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": 2,
+	    "dtype": "bfloat16",
+	    "distributed_executor_backend": "mp",
+	    "block_size": 128,
+	    "trust_remote_code": "",
+            "disable_log_stats": "",
+	    "enforce_eager": "",
+	    "max_num_batched_tokens": 2048,
+	    "max_num_seqs": 256,
+            "load_format": "dummy"
+        },
+        "client_parameters": {
+            "model": "meta-llama/Llama-3.1-8B-Instruct",
+            "backend": "vllm",
+            "dataset_name": "sharegpt",
+            "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json",
+            "num_prompts": 200
+        }
+    },
+    {
+        "test_name": "serving_llama8B_bf16_tp4_sharegpt",
+        "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": 4,
+	    "dtype": "bfloat16",
+	    "distributed_executor_backend": "mp",
+	    "block_size": 128,
+	    "trust_remote_code": "",
+            "disable_log_stats": "",
+	    "enforce_eager": "",
+	    "max_num_batched_tokens": 2048,
+	    "max_num_seqs": 256,
+            "load_format": "dummy"
+        },
+        "client_parameters": {
+            "model": "meta-llama/Llama-3.1-8B-Instruct",
+            "backend": "vllm",
+            "dataset_name": "sharegpt",
+            "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json",
+            "num_prompts": 200
+        }
+    },
+    {
+        "test_name": "serving_llama8B_bf16_tp1_random_128_128",
+        "qps_list": ["inf"],
+        "max_concurrency_list": [12, 16, 24, 32, 64, 128, 200, 1000],
+        "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": "",
+	    "enable_chunked_prefill": "",
+            "disable_log_stats": "",
+	    "enforce_eager": "",
+	    "max_num_batched_tokens": 2048,
+	    "max_num_seqs": 256,
+            "load_format": "dummy"
+        },
+        "client_parameters": {
+            "model": "meta-llama/Llama-3.1-8B-Instruct",
+            "backend": "vllm",
+            "dataset_name": "random",
+	    "random-input-len": 128,
+	    "random-output-len": 128,
+	    "ignore-eos": "",
+            "num_prompts": 1000
+        }
+    },
+    {
+        "test_name": "serving_llama8B_bf16_tp2_random_128_128",
+        "qps_list": ["inf"],
+        "max_concurrency_list": [12, 16, 24, 32, 64, 128, 200, 1000],
+        "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": 2,
+	    "dtype": "bfloat16",
+	    "distributed_executor_backend": "mp",
+	    "block_size": 128,
+	    "trust_remote_code": "",
+	    "enable_chunked_prefill": "",
+            "disable_log_stats": "",
+	    "enforce_eager": "",
+	    "max_num_batched_tokens": 2048,
+	    "max_num_seqs": 256,
+            "load_format": "dummy"
+        },
+        "client_parameters": {
+            "model": "meta-llama/Llama-3.1-8B-Instruct",
+            "backend": "vllm",
+            "dataset_name": "random",
+	    "random-input-len": 128,
+	    "random-output-len": 128,
+	    "ignore-eos": "",
+            "num_prompts": 1000
+        }
+    },
+    {
+        "test_name": "serving_llama8B_bf16_tp4_random_128_128",
+        "qps_list": ["inf"],
+        "max_concurrency_list": [12, 16, 24, 32, 64, 128, 200, 1000],
+        "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": 4,
+	    "dtype": "bfloat16",
+	    "distributed_executor_backend": "mp",
+	    "block_size": 128,
+	    "trust_remote_code": "",
+	    "enable_chunked_prefill": "",
+            "disable_log_stats": "",
+	    "enforce_eager": "",
+	    "max_num_batched_tokens": 2048,
+	    "max_num_seqs": 256,
+            "load_format": "dummy"
+        },
+        "client_parameters": {
+            "model": "meta-llama/Llama-3.1-8B-Instruct",
+            "backend": "vllm",
+            "dataset_name": "random",
+	    "random-input-len": 128,
+	    "random-output-len": 128,
+            "num_prompts": 1000
+        }
+    },
+    {
+        "test_name": "serving_llama8B_int8_tp1_sharegpt",
+        "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": "RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8",
+            "tensor_parallel_size": 1,
+	    "dtype": "bfloat16",
+	    "distributed_executor_backend": "mp",
+	    "block_size": 128,
+	    "trust_remote_code": "",
+            "disable_log_stats": "",
+	    "enforce_eager": "",
+	    "max_num_batched_tokens": 2048,
+	    "max_num_seqs": 256,
+            "load_format": "dummy"
+        },
+        "client_parameters": {
+            "model": "RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8",
+            "backend": "vllm",
+            "dataset_name": "sharegpt",
+            "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json",
+            "num_prompts": 200
+        }
+    },
+    {
+        "test_name": "serving_llama8B_int8_tp2_sharegpt",
+        "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": "RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8",
+            "tensor_parallel_size": 2,
+	    "dtype": "bfloat16",
+	    "distributed_executor_backend": "mp",
+	    "block_size": 128,
+	    "trust_remote_code": "",
+            "disable_log_stats": "",
+	    "enforce_eager": "",
+	    "max_num_batched_tokens": 2048,
+	    "max_num_seqs": 256,
+            "load_format": "dummy"
+        },
+        "client_parameters": {
+            "model": "RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8",
+            "backend": "vllm",
+            "dataset_name": "sharegpt",
+            "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json",
+            "num_prompts": 200
+        }
+    },
+    {
+        "test_name": "serving_llama8B_int8_tp4_sharegpt",
+        "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": "RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8",
+            "tensor_parallel_size": 4,
+	    "dtype": "bfloat16",
+	    "distributed_executor_backend": "mp",
+	    "block_size": 128,
+	    "trust_remote_code": "",
+            "disable_log_stats": "",
+	    "enforce_eager": "",
+	    "max_num_batched_tokens": 2048,
+	    "max_num_seqs": 256,
+            "load_format": "dummy"
+        },
+        "client_parameters": {
+            "model": "RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8",
+            "backend": "vllm",
+            "dataset_name": "sharegpt",
+            "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json",
+            "num_prompts": 200
+        }
+    },
+    {
+        "test_name": "serving_llama8B_int8_tp1_random_128_128",
+        "qps_list": ["inf"],
+        "max_concurrency_list": [12, 16, 24, 32, 64, 128, 200, 1000],
+        "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": "RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8",
+            "tensor_parallel_size": 1,
+	    "dtype": "bfloat16",
+	    "distributed_executor_backend": "mp",
+	    "block_size": 128,
+	    "trust_remote_code": "",
+	    "enable_chunked_prefill": "",
+            "disable_log_stats": "",
+	    "enforce_eager": "",
+	    "max_num_batched_tokens": 2048,
+	    "max_num_seqs": 256,
+            "load_format": "dummy"
+        },
+        "client_parameters": {
+            "model": "RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8",
+            "backend": "vllm",
+            "dataset_name": "random",
+	    "random-input-len": 128,
+	    "random-output-len": 128,
+	    "ignore-eos": "",
+            "num_prompts": 1000
+        }
+    },
+    {
+        "test_name": "serving_llama8B_int8_tp2_random_128_128",
+        "qps_list": ["inf"],
+        "max_concurrency_list": [12, 16, 24, 32, 64, 128, 200, 1000],
+        "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": "RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8",
+            "tensor_parallel_size": 2,
+	    "dtype": "bfloat16",
+	    "distributed_executor_backend": "mp",
+	    "block_size": 128,
+	    "trust_remote_code": "",
+	    "enable_chunked_prefill": "",
+            "disable_log_stats": "",
+	    "enforce_eager": "",
+	    "max_num_batched_tokens": 2048,
+	    "max_num_seqs": 256,
+            "load_format": "dummy"
+        },
+        "client_parameters": {
+            "model": "RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8",
+            "backend": "vllm",
+            "dataset_name": "random",
+	    "random-input-len": 128,
+	    "random-output-len": 128,
+	    "ignore-eos": "",
+            "num_prompts": 1000
+        }
+    },
+    {
+        "test_name": "serving_llama8B_int8_tp4_random_128_128",
+        "qps_list": ["inf"],
+        "max_concurrency_list": [12, 16, 24, 32, 64, 128, 200, 1000],
+        "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": "RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8",
+            "tensor_parallel_size": 4,
+	    "dtype": "bfloat16",
+	    "distributed_executor_backend": "mp",
+	    "block_size": 128,
+	    "trust_remote_code": "",
+	    "enable_chunked_prefill": "",
+            "disable_log_stats": "",
+	    "enforce_eager": "",
+	    "max_num_batched_tokens": 2048,
+	    "max_num_seqs": 256,
+            "load_format": "dummy"
+        },
+        "client_parameters": {
+            "model": "RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8",
+            "backend": "vllm",
+            "dataset_name": "random",
+	    "random-input-len": 128,
+	    "random-output-len": 128,
+	    "ignore-eos": "",
+            "num_prompts": 1000
+        }
+    },
+    {
+        "test_name": "serving_llama8B_int4_tp1_sharegpt",
+        "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": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
+	    "quantization": "awq",
+            "tensor_parallel_size": 1,
+	    "dtype": "bfloat16",
+	    "distributed_executor_backend": "mp",
+	    "block_size": 128,
+	    "trust_remote_code": "",
+            "disable_log_stats": "",
+	    "enforce_eager": "",
+	    "max_num_batched_tokens": 2048,
+	    "max_num_seqs": 256,
+            "load_format": "dummy"
+        },
+        "client_parameters": {
+            "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
+            "backend": "vllm",
+            "dataset_name": "sharegpt",
+            "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json",
+            "num_prompts": 200
+        }
+    },
+    {
+        "test_name": "serving_llama8B_int4_tp2_sharegpt",
+        "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": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
+	    "quantization": "awq",
+            "tensor_parallel_size": 2,
+	    "dtype": "bfloat16",
+	    "distributed_executor_backend": "mp",
+	    "block_size": 128,
+	    "trust_remote_code": "",
+            "disable_log_stats": "",
+	    "enforce_eager": "",
+	    "max_num_batched_tokens": 2048,
+	    "max_num_seqs": 256,
+            "load_format": "dummy"
+        },
+        "client_parameters": {
+            "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
+            "backend": "vllm",
+            "dataset_name": "sharegpt",
+            "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json",
+            "num_prompts": 200
+        }
+    },
+    {
+        "test_name": "serving_llama8B_int4_tp4_sharegpt",
+        "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": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
+	    "quantization": "awq",
+            "tensor_parallel_size": 4,
+	    "dtype": "bfloat16",
+	    "distributed_executor_backend": "mp",
+	    "block_size": 128,
+	    "trust_remote_code": "",
+            "disable_log_stats": "",
+	    "enforce_eager": "",
+	    "max_num_batched_tokens": 2048,
+	    "max_num_seqs": 256,
+            "load_format": "dummy"
+        },
+        "client_parameters": {
+            "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
+            "backend": "vllm",
+            "dataset_name": "sharegpt",
+            "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json",
+            "num_prompts": 200
+        }
+    },
+    {
+        "test_name": "serving_llama8B_int4_tp1_random_128_128",
+        "qps_list": ["inf"],
+        "max_concurrency_list": [12, 16, 24, 32, 64, 128, 200, 1000],
+        "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": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
+	    "quantization": "awq",
+            "tensor_parallel_size": 1,
+	    "dtype": "bfloat16",
+	    "distributed_executor_backend": "mp",
+	    "block_size": 128,
+	    "trust_remote_code": "",
+	    "enable_chunked_prefill": "",
+            "disable_log_stats": "",
+	    "enforce_eager": "",
+	    "max_num_batched_tokens": 2048,
+	    "max_num_seqs": 256,
+            "load_format": "dummy"
+        },
+        "client_parameters": {
+            "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
+            "backend": "vllm",
+            "dataset_name": "random",
+	    "random-input-len": 128,
+	    "random-output-len": 128,
+	    "ignore-eos": "",
+            "num_prompts": 1000
+        }
+    },
+    {
+        "test_name": "serving_llama8B_int4_tp2_random_128_128",
+        "qps_list": ["inf"],
+        "max_concurrency_list": [12, 16, 24, 32, 64, 128, 200, 1000],
+        "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": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
+	    "quantization": "awq",
+            "tensor_parallel_size": 2,
+	    "dtype": "bfloat16",
+	    "distributed_executor_backend": "mp",
+	    "block_size": 128,
+	    "trust_remote_code": "",
+	    "enable_chunked_prefill": "",
+            "disable_log_stats": "",
+	    "enforce_eager": "",
+	    "max_num_batched_tokens": 2048,
+	    "max_num_seqs": 256,
+            "load_format": "dummy"
+        },
+        "client_parameters": {
+            "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
+            "backend": "vllm",
+            "dataset_name": "random",
+	    "random-input-len": 128,
+	    "random-output-len": 128,
+	    "ignore-eos": "",
+            "num_prompts": 1000
+        }
+    },
+    {
+        "test_name": "serving_llama8B_int4_tp4_random_128_128",
+        "qps_list": ["inf"],
+        "max_concurrency_list": [12, 16, 24, 32, 64, 128, 200, 1000],
+        "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": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
+	    "quantization": "awq",
+            "tensor_parallel_size": 4,
+	    "dtype": "bfloat16",
+	    "distributed_executor_backend": "mp",
+	    "block_size": 128,
+	    "trust_remote_code": "",
+	    "enable_chunked_prefill": "",
+            "disable_log_stats": "",
+	    "enforce_eager": "",
+	    "max_num_batched_tokens": 2048,
+	    "max_num_seqs": 256,
+            "load_format": "dummy"
+        },
+        "client_parameters": {
+            "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
+            "backend": "vllm",
+            "dataset_name": "random",
+	    "random-input-len": 128,
+	    "random-output-len": 128,
+	    "ignore-eos": "",
+            "num_prompts": 1000
+        }
+    }
+]
@@ -1,246 +1,276 @@
-{
-  "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
+[
+    {
+        "test_name": "serving_llama8B_tp1_sharegpt",
+        "qps_list": [1, 4, 16, "inf"],
+        "max_concurrency_list": [32],
+        "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": "",
+	    "enforce_eager": "",
+	    "max_num_batched_tokens": 2048,
+	    "max_num_seqs": 256,
+            "load_format": "dummy"
+        },
+        "client_parameters": {
+            "model": "meta-llama/Llama-3.1-8B-Instruct",
+            "backend": "vllm",
+            "dataset_name": "sharegpt",
+            "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json",
+            "num_prompts": 32
+        }
    },
-    "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": "",
-      "enforce_eager": "",
-      "max_num_batched_tokens": 2048,
-      "max_num_seqs": 256,
-      "load_format": "dummy"
+    {
+        "test_name": "serving_llama8B_tp2_sharegpt",
+        "qps_list": [1, 4, 16, "inf"],
+        "max_concurrency_list": [32],
+        "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": 2,
+	    "dtype": "bfloat16",
+	    "distributed_executor_backend": "mp",
+	    "block_size": 128,
+	    "trust_remote_code": "",
+            "disable_log_stats": "",
+	    "enforce_eager": "",
+	    "max_num_batched_tokens": 2048,
+	    "max_num_seqs": 256,
+            "load_format": "dummy"
+        },
+        "client_parameters": {
+            "model": "meta-llama/Llama-3.1-8B-Instruct",
+            "backend": "vllm",
+            "dataset_name": "sharegpt",
+            "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json",
+            "num_prompts": 32
+        }
    },
-    "client_parameters": {
-      "model": "meta-llama/Llama-3.1-8B-Instruct",
-      "backend": "vllm",
-      "ignore-eos": "",
-      "num_prompts": 200
+    {
+        "test_name": "serving_llama8B_tp1_random_128_128",
+        "qps_list": [1, 4, 16, "inf"],
+        "max_concurrency_list": [32],
+        "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": "",
+	    "enable_chunked_prefill": "",
+            "disable_log_stats": "",
+	    "enforce_eager": "",
+	    "max_num_batched_tokens": 2048,
+	    "max_num_seqs": 256,
+            "load_format": "dummy"
+        },
+        "client_parameters": {
+            "model": "meta-llama/Llama-3.1-8B-Instruct",
+            "backend": "vllm",
+            "dataset_name": "random",
+	    "random-input-len": 128,
+	    "random-output-len": 128,
+	    "ignore-eos": "",
+            "num_prompts": 32
+        }
+    },
+    {
+        "test_name": "serving_llama8B_tp2_random_128_128",
+        "qps_list": [1, 4, 16, "inf"],
+        "max_concurrency_list": [32],
+        "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": 2,
+	    "dtype": "bfloat16",
+	    "distributed_executor_backend": "mp",
+	    "block_size": 128,
+	    "trust_remote_code": "",
+	    "enable_chunked_prefill": "",
+            "disable_log_stats": "",
+	    "enforce_eager": "",
+	    "max_num_batched_tokens": 2048,
+	    "max_num_seqs": 256,
+            "load_format": "dummy"
+        },
+        "client_parameters": {
+            "model": "meta-llama/Llama-3.1-8B-Instruct",
+            "backend": "vllm",
+            "dataset_name": "random",
+	    "random-input-len": 128,
+	    "random-output-len": 128,
+	    "ignore-eos": "",
+            "num_prompts": 32
+        }
+    },
+    {
+        "test_name": "serving_llama8B_tp1_random_128_2048",
+        "qps_list": [1, 4, 16, "inf"],
+        "max_concurrency_list": [32],
+        "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": "",
+	    "enable_chunked_prefill": "",
+            "disable_log_stats": "",
+	    "enforce_eager": "",
+	    "max_num_batched_tokens": 2048,
+	    "max_num_seqs": 256,
+            "load_format": "dummy"
+        },
+        "client_parameters": {
+            "model": "meta-llama/Llama-3.1-8B-Instruct",
+            "backend": "vllm",
+            "dataset_name": "random",
+	    "random-input-len": 128,
+	    "random-output-len": 2048,
+	    "ignore-eos": "",
+            "num_prompts": 32
+        }
+    },
+    {
+        "test_name": "serving_llama8B_tp2_random_128_2048",
+        "qps_list": [1, 4, 16, "inf"],
+        "max_concurrency_list": [32],
+        "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": 2,
+	    "dtype": "bfloat16",
+	    "distributed_executor_backend": "mp",
+	    "block_size": 128,
+	    "trust_remote_code": "",
+	    "enable_chunked_prefill": "",
+            "disable_log_stats": "",
+	    "enforce_eager": "",
+	    "max_num_batched_tokens": 2048,
+	    "max_num_seqs": 256,
+            "load_format": "dummy"
+        },
+        "client_parameters": {
+            "model": "meta-llama/Llama-3.1-8B-Instruct",
+            "backend": "vllm",
+            "dataset_name": "random",
+	    "random-input-len": 128,
+	    "random-output-len": 2048,
+	    "ignore-eos": "",
+            "num_prompts": 32
+        }
+    },
+    {
+        "test_name": "serving_llama8B_tp1_random_2048_128",
+        "qps_list": [1, 4, 16, "inf"],
+        "max_concurrency_list": [32],
+        "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": "",
+	    "enable_chunked_prefill": "",
+            "disable_log_stats": "",
+	    "enforce_eager": "",
+	    "max_num_batched_tokens": 2048,
+	    "max_num_seqs": 256,
+            "load_format": "dummy"
+        },
+        "client_parameters": {
+            "model": "meta-llama/Llama-3.1-8B-Instruct",
+            "backend": "vllm",
+            "dataset_name": "random",
+	    "random-input-len": 2048,
+	    "random-output-len": 128,
+	    "ignore-eos": "",
+            "num_prompts": 32
+        }
+    },
+    {
+        "test_name": "serving_llama8B_tp2_random_2048_128",
+        "qps_list": [1, 4, 16, "inf"],
+        "max_concurrency_list": [32],
+        "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": 2,
+	    "dtype": "bfloat16",
+	    "distributed_executor_backend": "mp",
+	    "block_size": 128,
+	    "trust_remote_code": "",
+	    "enable_chunked_prefill": "",
+            "disable_log_stats": "",
+	    "enforce_eager": "",
+	    "max_num_batched_tokens": 2048,
+	    "max_num_seqs": 256,
+            "load_format": "dummy"
+        },
+        "client_parameters": {
+            "model": "meta-llama/Llama-3.1-8B-Instruct",
+            "backend": "vllm",
+            "dataset_name": "random",
+	    "random-input-len": 2048,
+	    "random-output-len": 128,
+	    "ignore-eos": "",
+            "num_prompts": 32
+        }
    }
-  },
-  "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_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
-      }
-    }
-  ]
-}
+]
@@ -8,28 +8,13 @@ steps:
    commands:
      # #NOTE: torch_cuda_arch_list is derived from upstream PyTorch build files here:
      # https://github.com/pytorch/pytorch/blob/main/.ci/aarch64_linux/aarch64_ci_build.sh#L7
-      - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
+      - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg VLLM_MAIN_CUDA_VERSION=12.9 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
      - "mkdir artifacts"
      - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
      - "bash .buildkite/scripts/upload-wheels.sh"
    env:
      DOCKER_BUILDKIT: "1"

-  - label: "Build arm64 wheel - CUDA 13.0"
-    depends_on: ~
-    id: build-wheel-arm64-cuda-13-0
-    agents:
-      queue: arm64_cpu_queue_postmerge
-    commands:
-      # #NOTE: torch_cuda_arch_list is derived from upstream PyTorch build files here:
-      # https://github.com/pytorch/pytorch/blob/main/.ci/aarch64_linux/aarch64_ci_build.sh#L7
-      - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04  --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
-      - "mkdir artifacts"
-      - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
-      - "bash .buildkite/scripts/upload-wheels.sh manylinux_2_35"
-    env:
-      DOCKER_BUILDKIT: "1"
-
  # aarch64 build
  - label: "Build arm64 CPU wheel"
    depends_on: ~
@@ -40,11 +25,24 @@ steps:
      - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_BUILD_ACL=ON --tag vllm-ci:build-image --target vllm-build --progress plain -f docker/Dockerfile.cpu ."
      - "mkdir artifacts"
      - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
-      - "bash .buildkite/scripts/upload-wheels.sh manylinux_2_35"
+      - "bash .buildkite/scripts/upload-wheels.sh"
    env:
      DOCKER_BUILDKIT: "1"

  # x86 + CUDA builds
+  - label: "Build wheel - CUDA 12.8"
+    depends_on: ~
+    id: build-wheel-cuda-12-8
+    agents:
+      queue: cpu_queue_postmerge
+    commands:
+      - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.8.1 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
+      - "mkdir artifacts"
+      - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
+      - "bash .buildkite/scripts/upload-wheels.sh"
+    env:
+      DOCKER_BUILDKIT: "1"
+
  - label: "Build wheel - CUDA 12.9"
    depends_on: ~
    id: build-wheel-cuda-12-9
@@ -54,7 +52,7 @@ steps:
      - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
      - "mkdir artifacts"
      - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
-      - "bash .buildkite/scripts/upload-wheels.sh manylinux_2_31"
+      - "bash .buildkite/scripts/upload-wheels.sh"
    env:
      DOCKER_BUILDKIT: "1"

@@ -67,21 +65,7 @@ steps:
      - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
      - "mkdir artifacts"
      - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
-      - "bash .buildkite/scripts/upload-wheels.sh manylinux_2_35"
-    env:
-      DOCKER_BUILDKIT: "1"
-
-  # x86 CPU wheel build
-  - label: "Build x86 CPU wheel"
-    depends_on: ~
-    id: build-wheel-x86-cpu
-    agents:
-      queue: cpu_queue_postmerge
-    commands:
-      - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_CPU_AVX512BF16=true --build-arg VLLM_CPU_AVX512VNNI=true --build-arg VLLM_CPU_AMXBF16=true --tag vllm-ci:build-image --target vllm-build --progress plain -f docker/Dockerfile.cpu ."
-      - "mkdir artifacts"
-      - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
-      - "bash .buildkite/scripts/upload-wheels.sh manylinux_2_35"
+      - "bash .buildkite/scripts/upload-wheels.sh"
    env:
      DOCKER_BUILDKIT: "1"

@@ -125,6 +109,7 @@ steps:
  - label: "Annotate release workflow"
    depends_on:
      - create-multi-arch-manifest
+      - build-wheel-cuda-12-8
    id: annotate-release-workflow
    agents:
      queue: cpu_queue_postmerge
@@ -1,400 +0,0 @@
-#!/usr/bin/env python3
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-# do not complain about line length (for docstring)
-# ruff: noqa: E501
-
-import argparse
-import json
-import sys
-from dataclasses import asdict, dataclass
-from datetime import datetime
-from pathlib import Path
-from typing import Any
-from urllib.parse import quote
-
-import regex as re
-
-if not sys.version_info >= (3, 12):
-    raise RuntimeError("This script requires Python 3.12 or higher.")
-
-INDEX_HTML_TEMPLATE = """<!DOCTYPE html>
-<html>
-  <!-- {comment} -->
-  <meta name="pypi:repository-version" content="1.0">
-  <body>
-{items}
-  </body>
-</html>
-"""
-
-
-@dataclass
-class WheelFileInfo:
-    package_name: str
-    version: str
-    build_tag: str | None
-    python_tag: str
-    abi_tag: str
-    platform_tag: str
-    variant: str | None
-    filename: str
-
-
-def parse_from_filename(file: str) -> WheelFileInfo:
-    """
-    Parse wheel file name to extract metadata.
-
-    The format of wheel names:
-        {package_name}-{version}(-{build_tag})?-{python_tag}-{abi_tag}-{platform_tag}.whl
-    All versions could contain a variant like '+cu129' or '.cpu' or `.rocm` (or not).
-    Example:
-        vllm-0.11.0-cp38-abi3-manylinux1_x86_64.whl
-        vllm-0.10.2rc2+cu129-cp38-abi3-manylinux2014_aarch64.whl
-        vllm-0.11.1rc8.dev14+gaa384b3c0-cp38-abi3-manylinux2014_aarch64.whl
-        vllm-0.11.1rc8.dev14+gaa384b3c0.cu130-cp38-abi3-manylinux1_x86_64.whl
-    """
-    wheel_file_re = re.compile(
-        r"^(?P<package_name>.+)-(?P<version>[^-]+?)(-(?P<build_tag>[^-]+))?-(?P<python_tag>[^-]+)-(?P<abi_tag>[^-]+)-(?P<platform_tag>[^-]+)\.whl$"
-    )
-    match = wheel_file_re.match(file)
-    if not match:
-        raise ValueError(f"Invalid wheel file name: {file}")
-
-    package_name = match.group("package_name")
-    version = match.group("version")
-    build_tag = match.group("build_tag")
-    python_tag = match.group("python_tag")
-    abi_tag = match.group("abi_tag")
-    platform_tag = match.group("platform_tag")
-
-    # extract variant from version
-    variant = None
-    if "dev" in version:
-        ver_after_dev = version.split("dev")[-1]
-        if "." in ver_after_dev:
-            variant = ver_after_dev.split(".")[-1]
-            version = version.removesuffix("." + variant)
-    else:
-        if "+" in version:
-            version, variant = version.split("+")
-
-    return WheelFileInfo(
-        package_name=package_name,
-        version=version,
-        build_tag=build_tag,
-        python_tag=python_tag,
-        abi_tag=abi_tag,
-        platform_tag=platform_tag,
-        variant=variant,
-        filename=file,
-    )
-
-
-def generate_project_list(subdir_names: list[str], comment: str = "") -> str:
-    """
-    Generate project list HTML content linking to each project & variant sub-directory.
-    """
-    href_tags = []
-    for name in sorted(subdir_names):
-        name = name.strip("/").strip(".")
-        href_tags.append(f'    <a href="{name}/">{name}/</a><br/>')
-    return INDEX_HTML_TEMPLATE.format(items="\n".join(href_tags), comment=comment)
-
-
-def generate_package_index_and_metadata(
-    wheel_files: list[WheelFileInfo],
-    wheel_base_dir: Path,
-    index_base_dir: Path,
-    comment: str = "",
-) -> tuple[str, str]:
-    """
-    Generate package index HTML content for a specific package, linking to actual wheel files.
-    """
-    href_tags = []
-    metadata = []
-    for file in sorted(wheel_files, key=lambda x: x.filename):
-        relative_path = (
-            wheel_base_dir.relative_to(index_base_dir, walk_up=True) / file.filename
-        )
-        # handle with '+' in URL, and avoid double-encoding '/' and already-encoded '%2B'
-        # NOTE: this is AWS S3 specific behavior!
-        file_path_quoted = quote(relative_path.as_posix(), safe=":%/")
-        href_tags.append(f'    <a href="{file_path_quoted}">{file.filename}</a><br/>')
-        file_meta = asdict(file)
-        file_meta["path"] = file_path_quoted
-        metadata.append(file_meta)
-    index_str = INDEX_HTML_TEMPLATE.format(items="\n".join(href_tags), comment=comment)
-    metadata_str = json.dumps(metadata, indent=2)
-    return index_str, metadata_str
-
-
-def generate_index_and_metadata(
-    whl_files: list[str],
-    wheel_base_dir: Path,
-    index_base_dir: Path,
-    default_variant: str | None = None,
-    alias_to_default: str | None = None,
-    comment: str = "",
-):
-    """
-    Generate index for all wheel files.
-
-    Args:
-        whl_files (list[str]): List of wheel files (must be directly under `wheel_base_dir`).
-        wheel_base_dir (Path): Base directory for wheel files.
-        index_base_dir (Path): Base directory to store index files.
-        default_variant (str | None): The default variant name, if any.
-        alias_to_default (str | None): Alias variant name for the default variant, if any.
-        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 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 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 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 sub-directory
-                index.html
-                vllm/
-                    index.html
-                    metadata.json
-            cu129/ # cu129 is actually the alias to default variant
-                index.html
-                vllm/
-                    index.html
-                    metadata.json
-            cu130/ # cu130 variant sub-directory
-                index.html
-                vllm/
-                    index.html
-                    metadata.json
-            ...
-
-    metadata.json stores a dump of all wheel files' metadata in a machine-readable format:
-        [
-            {
-                "package_name": "vllm",
-                "version": "0.10.2rc2",
-                "build_tag": null,
-                "python_tag": "cp38",
-                "abi_tag": "abi3",
-                "platform_tag": "manylinux2014_aarch64",
-                "variant": "cu129",
-                "filename": "vllm-0.10.2rc2+cu129-cp38-abi3-manylinux2014_aarch64.whl",
-                "path": "../vllm-0.10.2rc2%2Bcu129-cp38-abi3-manylinux2014_aarch64.whl" # to be concatenated with the directory URL and URL-encoded
-            },
-            ...
-        ]
-    """
-
-    parsed_files = [parse_from_filename(f) for f in whl_files]
-
-    if not parsed_files:
-        print("No wheel files found, skipping index generation.")
-        return
-
-    # Group by variant
-    variant_to_files: dict[str, list[WheelFileInfo]] = {}
-    for file in parsed_files:
-        variant = file.variant or "default"
-        if variant not in variant_to_files:
-            variant_to_files[variant] = []
-        variant_to_files[variant].append(file)
-
-    print(f"Found variants: {list(variant_to_files.keys())}")
-
-    # sanity check for default variant
-    if default_variant:
-        if "default" in variant_to_files:
-            raise ValueError(
-                "All wheel files must have variant suffixes when `default_variant` is specified."
-            )
-        if default_variant not in variant_to_files:
-            raise ValueError(
-                f"Default variant '{default_variant}' not found among wheel files."
-            )
-
-    if alias_to_default:
-        if "default" not in variant_to_files:
-            # e.g. only some wheels are uploaded to S3 currently
-            print(
-                "[WARN] Alias to default variant specified, but no default variant found."
-            )
-        elif alias_to_default in variant_to_files:
-            raise ValueError(
-                f"Alias variant name '{alias_to_default}' already exists among wheel files."
-            )
-        else:
-            variant_to_files[alias_to_default] = variant_to_files["default"].copy()
-            print(f"Alias variant '{alias_to_default}' created for default variant.")
-
-    # Generate comment in HTML header
-    comment_str = f" ({comment})" if comment else ""
-    comment_tmpl = f"Generated on {datetime.now().isoformat()}{comment_str}"
-
-    # Generate index for each variant
-    subdir_names = set()
-    for variant, files in variant_to_files.items():
-        if variant == "default":
-            variant_dir = index_base_dir
-        else:
-            variant_dir = index_base_dir / variant
-            subdir_names.add(variant)
-
-        variant_dir.mkdir(parents=True, exist_ok=True)
-
-        # gather all package names in this variant
-        packages = set(f.package_name for f in files)
-        if variant == "default":
-            # these packages should also appear in the "project list"
-            # generate after all variants are processed
-            subdir_names = subdir_names.union(packages)
-        else:
-            # generate project list for this variant directly
-            project_list_str = generate_project_list(sorted(packages), comment_tmpl)
-            with open(variant_dir / "index.html", "w") as f:
-                f.write(project_list_str)
-
-        for package in packages:
-            # filter files belonging to this package only
-            package_files = [f for f in files if f.package_name == package]
-            package_dir = variant_dir / package
-            package_dir.mkdir(parents=True, exist_ok=True)
-            index_str, metadata_str = generate_package_index_and_metadata(
-                package_files, wheel_base_dir, package_dir, comment
-            )
-            with open(package_dir / "index.html", "w") as f:
-                f.write(index_str)
-            with open(package_dir / "metadata.json", "w") as f:
-                f.write(metadata_str)
-
-    # Generate top-level project list index
-    project_list_str = generate_project_list(sorted(subdir_names), comment_tmpl)
-    with open(index_base_dir / "index.html", "w") as f:
-        f.write(project_list_str)
-
-
-if __name__ == "__main__":
-    """
-    Arguments:
-        --version <version> : version string for the current build (e.g., commit hash)
-        --current-objects <path_to_json> : path to JSON file containing current S3 objects listing in this version directory
-        --output-dir <output_directory> : directory to store generated index files
-        --alias-to-default <alias_variant_name> : (optional) alias variant name for the default variant
-        --comment <comment_string> : (optional) comment string to include in generated HTML files
-    """
-
-    parser = argparse.ArgumentParser(
-        description="Process nightly build wheel files to generate indices."
-    )
-    parser.add_argument(
-        "--version",
-        type=str,
-        required=True,
-        help="Version string for the current build (e.g., commit hash)",
-    )
-    parser.add_argument(
-        "--current-objects",
-        type=str,
-        required=True,
-        help="Path to JSON file containing current S3 objects listing in this version directory",
-    )
-    parser.add_argument(
-        "--output-dir",
-        type=str,
-        required=True,
-        help="Directory to store generated index files",
-    )
-    parser.add_argument(
-        "--alias-to-default",
-        type=str,
-        default=None,
-        help="Alias variant name for the default variant",
-    )
-    parser.add_argument(
-        "--comment",
-        type=str,
-        default="",
-        help="Optional comment string to include in generated HTML files",
-    )
-
-    args = parser.parse_args()
-
-    version = args.version
-    if "/" in version or "\\" in version:
-        raise ValueError("Version string must not contain slashes.")
-    current_objects_path = Path(args.current_objects)
-    output_dir = Path(args.output_dir)
-    if not output_dir.exists():
-        output_dir.mkdir(parents=True, exist_ok=True)
-
-    # Read current objects JSON
-    with open(current_objects_path) as f:
-        current_objects: dict[str, list[dict[str, Any]]] = json.load(f)
-
-    # current_objects looks like from list_objects_v2 S3 API:
-    """
-    "Contents": [
-        {
-            "Key": "e2f56c309d2a28899c68975a7e104502d56deb8f/vllm-0.11.2.dev363+ge2f56c309-cp38-abi3-manylinux1_x86_64.whl",
-            "LastModified": "2025-11-28T14:00:32+00:00",
-            "ETag": "\"37a38339c7cdb61ca737021b968075df-52\"",
-            "ChecksumAlgorithm": [
-                "CRC64NVME"
-            ],
-            "ChecksumType": "FULL_OBJECT",
-            "Size": 435649349,
-            "StorageClass": "STANDARD"
-        },
-        ...
-    ]
-    """
-
-    # Extract wheel file keys
-    wheel_files = []
-    for item in current_objects.get("Contents", []):
-        key: str = item["Key"]
-        if key.endswith(".whl"):
-            wheel_files.append(key.split("/")[-1])  # only the filename is used
-
-    print(f"Found {len(wheel_files)} wheel files for version {version}: {wheel_files}")
-
-    # keep only "official" files for a non-nightly version (specifed by cli args)
-    PY_VERSION_RE = re.compile(r"^\d+\.\d+\.\d+([a-zA-Z0-9.+-]*)?$")
-    if PY_VERSION_RE.match(version):
-        # upload-wheels.sh ensures no "dev" is in args.version
-        wheel_files = list(
-            filter(lambda x: version in x and "dev" not in x, wheel_files)
-        )
-        print(f"Non-nightly version detected, wheel files used: {wheel_files}")
-    else:
-        print("Nightly version detected, keeping all wheel files.")
-
-    # Generate index and metadata, assuming wheels and indices are stored as:
-    # s3://vllm-wheels/{version}/<wheel files>
-    # s3://vllm-wheels/<anything>/<index files>
-    wheel_base_dir = Path(output_dir).parent / version
-    index_base_dir = Path(output_dir)
-
-    generate_index_and_metadata(
-        whl_files=wheel_files,
-        wheel_base_dir=wheel_base_dir,
-        index_base_dir=index_base_dir,
-        default_variant=None,
-        alias_to_default=args.alias_to_default,
-        comment=args.comment.strip(),
-    )
-    print(f"Successfully generated index and metadata in {output_dir}")
@@ -7,57 +7,53 @@ set -ex
 # allow to bind to different cores
 CORE_RANGE=${CORE_RANGE:-0-16}
 OMP_CORE_RANGE=${OMP_CORE_RANGE:-0-16}
+NUMA_NODE=${NUMA_NODE:-0}

-export CMAKE_BUILD_PARALLEL_LEVEL=16
+export CMAKE_BUILD_PARALLEL_LEVEL=32

 # Setup cleanup
 remove_docker_container() {
    set -e;
-    docker rm -f cpu-test || true;
+    docker rm -f cpu-test-"$NUMA_NODE" || true;
 }
 trap remove_docker_container EXIT
 remove_docker_container

 # Try building the docker image
-docker build --tag cpu-test --target vllm-test -f docker/Dockerfile.cpu .
+numactl -C "$CORE_RANGE" -N "$NUMA_NODE" docker build --tag cpu-test-"$NUMA_NODE" --target vllm-test -f docker/Dockerfile.cpu .

-# Run the image
-docker run -itd --cpuset-cpus="$CORE_RANGE" --entrypoint /bin/bash -v ~/.cache/huggingface:/root/.cache/huggingface -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 cpu-test
+# Run the image, setting --shm-size=4g for tensor parallel.
+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"

 function cpu_tests() {
  set -e
+  export NUMA_NODE=$2

-  docker exec cpu-test bash -c "
+  docker exec cpu-test-"$NUMA_NODE" bash -c "
    set -e
    pip list"

  # offline inference
-  docker exec cpu-test bash -c "
+  docker exec cpu-test-"$NUMA_NODE" bash -c "
    set -e
    python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m"

-  # Run model tests
-  docker exec cpu-test bash -c "
-    set -e
-    pytest -x -v -s tests/models/multimodal/generation/test_whisper.py -m cpu_model"
-
  # Run kernel tests
-  docker exec cpu-test bash -c "
+  docker exec cpu-test-"$NUMA_NODE" bash -c "
    set -e
    pytest -x -v -s tests/kernels/test_onednn.py
-    pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
-    pytest -x -v -s tests/kernels/moe/test_moe.py -k test_cpu_fused_moe_basic"
+    pytest -x -v -s tests/kernels/attention/test_cpu_attn.py"

  # basic online serving
-  docker exec cpu-test bash -c '
+  docker exec cpu-test-"$NUMA_NODE" bash -c '
    set -e
-    VLLM_CPU_OMP_THREADS_BIND=$E2E_OMP_THREADS vllm serve Qwen/Qwen3-0.6B --max-model-len 2048 &
+    VLLM_CPU_OMP_THREADS_BIND=$E2E_OMP_THREADS vllm serve meta-llama/Llama-3.2-3B-Instruct --max-model-len 2048 &
    server_pid=$!
    timeout 600 bash -c "until curl localhost:8000/v1/models; do sleep 1; done" || exit 1
    vllm bench serve \
      --backend vllm \
      --dataset-name random \
-      --model Qwen/Qwen3-0.6B \
+      --model meta-llama/Llama-3.2-3B-Instruct \
      --num-prompts 20 \
      --endpoint /v1/completions
    kill -s SIGTERM $server_pid &'
@@ -65,4 +61,4 @@ function cpu_tests() {

 # All of CPU tests are expected to be finished less than 40 mins.
 export -f cpu_tests
-timeout 2h bash -c cpu_tests
+timeout 2h bash -c "cpu_tests $CORE_RANGE $NUMA_NODE"
@@ -25,22 +25,20 @@ function cpu_tests() {

  # offline inference
  podman exec -it "$container_id" bash -c "
-    export TORCH_COMPILE_DISABLE=1
    set -xve
    python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m" >> $HOME/test_basic.log

  # Run basic model test
  podman exec -it "$container_id" bash -c "
-    export TORCH_COMPILE_DISABLE=1
    set -evx
    pip install pytest pytest-asyncio einops peft Pillow soundfile transformers_stream_generator matplotlib
-    pip install sentence-transformers datamodel_code_generator tblib 
+    pip install sentence-transformers datamodel_code_generator

    # Note: disable Bart until supports V1
    # pytest -v -s tests/models/language/generation/test_bart.py -m cpu_model
-    pytest -v -s tests/models/language/generation/test_common.py::test_models[False-False-5-32-openai-community/gpt2]
-    pytest -v -s tests/models/language/generation/test_common.py::test_models[False-False-5-32-facebook/opt-125m]
-    pytest -v -s tests/models/language/generation/test_common.py::test_models[False-False-5-32-google/gemma-1.1-2b-it]
+    pytest -v -s tests/models/language/generation/test_common.py::test_models[False-5-32-openai-community/gpt2]
+    pytest -v -s tests/models/language/generation/test_common.py::test_models[False-5-32-facebook/opt-125m]
+    pytest -v -s tests/models/language/generation/test_common.py::test_models[False-5-32-google/gemma-1.1-2b-it]
    pytest -v -s tests/models/language/pooling/test_classification.py::test_models[float-jason9693/Qwen2.5-1.5B-apeach]
    # TODO: Below test case tests/models/language/pooling/test_embedding.py::test_models[True-ssmits/Qwen2-7B-Instruct-embed-base] fails on ppc64le. Disabling it for time being.
    # pytest -v -s tests/models/language/pooling/test_embedding.py -m cpu_model" >> $HOME/test_rest.log
@@ -21,8 +21,8 @@ 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 .
+numactl -C "$CORE_RANGE" -N "$NUMA_NODE" docker build --tag cpu-test-"$NUMA_NODE" --target vllm-test -f docker/Dockerfile.cpu .
+numactl -C "$CORE_RANGE" -N "$NUMA_NODE" docker build --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 -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"
@@ -74,7 +74,6 @@ FROM ${BASE_IMAGE_NAME}

 # Define environments
 ENV DEBIAN_FRONTEND=noninteractive
-ENV SOC_VERSION="ascend910b1"

 RUN pip config set global.index-url http://cache-service-vllm.nginx-pypi-cache.svc.cluster.local:${PYPI_CACHE_PORT}/pypi/simple && \
    pip config set global.trusted-host cache-service-vllm.nginx-pypi-cache.svc.cluster.local && \
@@ -35,10 +35,9 @@ docker run \
    echo $ZE_AFFINITY_MASK
    pip install tblib==3.1.0
    python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager
-    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 -O3 -O.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
    VLLM_ATTENTION_BACKEND=TRITON_ATTN python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager
    cd tests
    pytest -v -s v1/core
@@ -47,6 +46,6 @@ docker run \
    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
-    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/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_shared_storage_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py
    pytest -v -s v1/test_serial_utils.py
 '
@@ -12,11 +12,6 @@ REPO_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)"
 PRIME_RL_REPO="https://github.com/PrimeIntellect-ai/prime-rl.git"
 PRIME_RL_DIR="${REPO_ROOT}/prime-rl"

-if command -v rocm-smi &> /dev/null || command -v rocminfo &> /dev/null; then
-    echo "AMD GPU detected. Prime-RL currently only supports NVIDIA. Skipping..."
-    exit 0
-fi
-
 echo "Setting up Prime-RL integration test environment..."

 # Clean up any existing Prime-RL directory
@@ -17,17 +17,7 @@ wait_for_server() {
 }

 MODEL="deepseek-ai/DeepSeek-V2-lite"
-
-# Set BACKENDS based on platform
-if command -v rocm-smi &> /dev/null || [[ -d /opt/rocm ]] || [[ -n "${ROCM_PATH:-}" ]]; then
-  # ROCm platform
-  BACKENDS=("allgather_reducescatter")
-  # Disable MOE padding for ROCm since it is causing eplb to fail
-  export VLLM_ROCM_MOE_PADDING=0
-else
-  # Non-ROCm platform (CUDA/other)
-  BACKENDS=("deepep_high_throughput" "deepep_low_latency")
-fi
+BACKENDS=("deepep_high_throughput" "deepep_low_latency")

 cleanup() {
  if [[ -n "${SERVER_PID:-}" ]] && kill -0 "${SERVER_PID}" 2>/dev/null; then
@@ -2,9 +2,9 @@
 set -euxo pipefail

 # args: [THRESHOLD] [NUM_QUESTIONS] [START_PORT]
-THRESHOLD=${1:-0.25}
+THRESHOLD=${1:-0.8}
 NUM_Q=${2:-1319}
-PORT=${3:-8040}
+PORT=${3:-8020}
 OUT_DIR=${OUT_DIR:-/tmp/vllm-scheduled}
 mkdir -p "${OUT_DIR}"

@@ -16,18 +16,8 @@ wait_for_server() {
    done'
 }

-MODEL="Qwen/Qwen3-Next-80B-A3B-Instruct"
-
-# Set BACKENDS based on platform
-if command -v rocm-smi &> /dev/null || [[ -d /opt/rocm ]] || [[ -n "${ROCM_PATH:-}" ]]; then
-  # ROCm platform
-  BACKENDS=("allgather_reducescatter")
-  # Disable MOE padding for ROCm since it is causing eplb to fail
-  export VLLM_ROCM_MOE_PADDING=0
-else
-  # Non-ROCm platform (CUDA/other)
-  BACKENDS=("deepep_high_throughput" "deepep_low_latency")
-fi
+MODEL="QWen/Qwen3-30B-A3B-FP8"
+BACKENDS=("deepep_high_throughput" "deepep_low_latency")

 cleanup() {
  if [[ -n "${SERVER_PID:-}" ]] && kill -0 "${SERVER_PID}" 2>/dev/null; then
@@ -46,14 +36,11 @@ for BACK in "${BACKENDS[@]}"; do
  VLLM_ALL2ALL_BACKEND=$BACK \
  vllm serve "$MODEL" \
    --enforce-eager \
-    --tensor-parallel-size 4 \
+    --tensor-parallel-size 2 \
+    --data-parallel-size 2 \
    --enable-expert-parallel \
-    --enable-eplb \
-    --eplb-config '{"window_size":200,"step_interval":600,"use_async":true}' \
-    --speculative-config '{"method":"qwen3_next_mtp","num_speculative_tokens":1}' \
    --trust-remote-code \
    --max-model-len 2048 \
-    --gpu-memory-utilization 0.9 \
    --port $PORT &
  SERVER_PID=$!
  wait_for_server $PORT
@@ -1,74 +0,0 @@
-#!/usr/bin/env bash
-set -euxo pipefail
-
-# args: [THRESHOLD] [NUM_QUESTIONS] [START_PORT] [DATA_PARALLEL_SIZE] [TENSOR_PARALLEL_SIZE]
-THRESHOLD=${1:-0.8}
-NUM_Q=${2:-1319}
-PORT=${3:-8020}
-DATA_PARALLEL_SIZE=${4:-2}
-TENSOR_PARALLEL_SIZE=${5:-2}
-OUT_DIR=${OUT_DIR:-/tmp/vllm-scheduled}
-mkdir -p "${OUT_DIR}"
-
-wait_for_server() {
-  local port=$1
-  timeout 600 bash -c '
-    until curl -sf "http://127.0.0.1:'"$port"'/health" > /dev/null; do
-      sleep 1
-    done'
-}
-
-MODEL="QWen/Qwen3-30B-A3B-FP8"
-# Set BACKENDS based on platform
-if command -v rocm-smi &> /dev/null || [[ -d /opt/rocm ]] || [[ -n "${ROCM_PATH:-}" ]]; then
-  # ROCm platform
-  BACKENDS=("allgather_reducescatter")
-  # Disable MOE padding for ROCm since it is causing eplb to fail
-  export VLLM_ROCM_MOE_PADDING=0
-else
-  # Non-ROCm platform (CUDA/other)
-  BACKENDS=("deepep_high_throughput" "deepep_low_latency")
-fi
-
-cleanup() {
-  if [[ -n "${SERVER_PID:-}" ]] && kill -0 "${SERVER_PID}" 2>/dev/null; then
-    kill "${SERVER_PID}" 2>/dev/null || true
-    for _ in {1..20}; do
-      kill -0 "${SERVER_PID}" 2>/dev/null || break
-      sleep 0.5
-    done
-    kill -9 "${SERVER_PID}" 2>/dev/null || true
-  fi
-}
-trap cleanup EXIT
-
-for BACK in "${BACKENDS[@]}"; do
-  VLLM_DEEP_GEMM_WARMUP=skip \
-  VLLM_ALL2ALL_BACKEND=$BACK \
-  vllm serve "$MODEL" \
-    --enforce-eager \
-    --enable-eplb \
-    --eplb-config '{"window_size":10, "step_interval":100, "num_redundant_experts":0, "log_balancedness":true}' \
-    --tensor-parallel-size ${TENSOR_PARALLEL_SIZE} \
-    --data-parallel-size ${DATA_PARALLEL_SIZE} \
-    --enable-expert-parallel \
-    --trust-remote-code \
-    --max-model-len 2048 \
-    --port $PORT &
-  SERVER_PID=$!
-  wait_for_server $PORT
-
-  TAG=$(echo "$MODEL" | tr '/: \\n' '_____')
-  OUT="${OUT_DIR}/${TAG}_${BACK}.json"
-  python3 tests/evals/gsm8k/gsm8k_eval.py --host http://127.0.0.1 --port $PORT --num-questions ${NUM_Q} --save-results ${OUT}
-  python3 - <<PY
-import json; acc=json.load(open('${OUT}'))['accuracy']
-print(f"${MODEL} ${BACK}: accuracy {acc:.3f}")
-assert acc >= ${THRESHOLD}, f"${MODEL} ${BACK} accuracy {acc}"
-PY
-
-  cleanup
-  SERVER_PID=
-  sleep 1
-  PORT=$((PORT+1))
-done
@@ -2,28 +2,6 @@

 set -ex

-# ======== part 0: setup ========
-
-BUCKET="vllm-wheels"
-INDICES_OUTPUT_DIR="indices"
-DEFAULT_VARIANT_ALIAS="cu129" # align with vLLM_MAIN_CUDA_VERSION in vllm/envs.py
-PYTHON=${PYTHON_PROG:=python3} # try to read from env var, otherwise use python3
-SUBPATH=$BUILDKITE_COMMIT
-S3_COMMIT_PREFIX="s3://$BUCKET/$SUBPATH/"
-
-# detect if python3.10+ is available
-has_new_python=$($PYTHON -c "print(1 if __import__('sys').version_info >= (3,12) else 0)")
-if [[ "$has_new_python" -eq 0 ]]; then
-    # use new python from docker
-    docker pull python:3-slim
-    PYTHON="docker run --rm -v $(pwd):/app -w /app python:3-slim python3"
-fi
-
-echo "Using python interpreter: $PYTHON"
-echo "Python version: $($PYTHON --version)"
-
-# ========= part 1: collect, rename & upload the wheel ==========
-
 # Assume wheels are in artifacts/dist/*.whl
 wheel_files=(artifacts/dist/*.whl)

@@ -32,76 +10,74 @@ if [[ ${#wheel_files[@]} -ne 1 ]]; then
  echo "Error: Expected exactly one wheel file in artifacts/dist/, but found ${#wheel_files[@]}"
  exit 1
 fi
+
+# Get the single wheel file
 wheel="${wheel_files[0]}"

-# default build image uses ubuntu 20.04, which corresponds to manylinux_2_31
-# we also accept params as manylinux tag
-# refer to https://github.com/mayeut/pep600_compliance?tab=readme-ov-file#acceptable-distros-to-build-wheels
-manylinux_version="${1:-manylinux_2_31}"
+# Detect architecture and rename 'linux' to appropriate manylinux version
+arch=$(uname -m)
+if [[ $arch == "x86_64" ]]; then
+    manylinux_version="manylinux1"
+elif [[ $arch == "aarch64" ]]; then
+    manylinux_version="manylinux2014"
+else
+    echo "Warning: Unknown architecture $arch, using manylinux1 as default"
+    manylinux_version="manylinux1"
+fi

 # Rename 'linux' to the appropriate manylinux version in the wheel filename
-if [[ "$wheel" != *"linux"* ]]; then
-  echo "Error: Wheel filename does not contain 'linux': $wheel"
-  exit 1
-fi
 new_wheel="${wheel/linux/$manylinux_version}"
 mv -- "$wheel" "$new_wheel"
 wheel="$new_wheel"
-echo "Renamed wheel to: $wheel"

 # Extract the version from the wheel
 version=$(unzip -p "$wheel" '**/METADATA' | grep '^Version: ' | cut -d' ' -f2)
-echo "Version in wheel: $version"
-pure_version="${version%%+*}"
-echo "Pure version (without variant): $pure_version"
+echo "Version: $version"

-# copy wheel to its own bucket
-aws s3 cp "$wheel" "$S3_COMMIT_PREFIX"
+normal_wheel="$wheel" # Save the original wheel filename

-# ========= part 2: generate and upload indices ==========
-# generate indices for all existing wheels in the commit directory
-# this script might be run multiple times if there are multiple variants being built
-# so we need to guarantee there is little chance for "TOCTOU" issues
-# i.e., one process is generating indices while another is uploading a new wheel
-# so we need to ensure no time-consuming operations happen below
+# If the version contains "dev", rename it to v1.0.0.dev for consistency
+if [[ $version == *dev* ]]; then
+    suffix="${version##*.}"
+    if [[ $suffix == cu* ]]; then
+        new_version="1.0.0.dev+${suffix}"
+    else
+        new_version="1.0.0.dev"
+    fi
+    new_wheel="${wheel/$version/$new_version}"
+    # use cp to keep both files in the artifacts directory
+    cp -- "$wheel" "$new_wheel"
+    wheel="$new_wheel"
+    version="$new_version"
+fi

-# list all wheels in the commit directory
-echo "Existing wheels on S3:"
-aws s3 ls "$S3_COMMIT_PREFIX"
-obj_json="objects.json"
-aws s3api list-objects-v2 --bucket "$BUCKET" --prefix "$SUBPATH/" --delimiter / --output json > "$obj_json"
-mkdir -p "$INDICES_OUTPUT_DIR"
+# Upload the wheel to S3
+python3 .buildkite/generate_index.py --wheel "$normal_wheel"

-# call script to generate indicies for all existing wheels
-# this indices have relative paths that could work as long as it is next to the wheel directory in s3
-# i.e., the wheels are always in s3://vllm-wheels/<commit>/
-# and indices can be placed in /<commit>/, or /nightly/, or /<version>/
-if [[ ! -z "$DEFAULT_VARIANT_ALIAS" ]]; then
-    alias_arg="--alias-to-default $DEFAULT_VARIANT_ALIAS"
+# generate index for this commit
+aws s3 cp "$wheel" "s3://vllm-wheels/$BUILDKITE_COMMIT/"
+aws s3 cp "$normal_wheel" "s3://vllm-wheels/$BUILDKITE_COMMIT/"
+
+if [[ $normal_wheel == *"cu129"* ]]; then
+    # only upload index.html for cu129 wheels (default wheels) as it
+    # is available on both x86 and arm64
+    aws s3 cp index.html "s3://vllm-wheels/$BUILDKITE_COMMIT/vllm/index.html"
+    aws s3 cp "s3://vllm-wheels/nightly/index.html" "s3://vllm-wheels/$BUILDKITE_COMMIT/index.html"
 else
-    alias_arg=""
+    echo "Skipping index files for non-cu129 wheels"
 fi

-# HACK: we do not need regex module here, but it is required by pre-commit hook
-# To avoid any external dependency, we simply replace it back to the stdlib re module
-sed -i 's/import regex as re/import re/g' .buildkite/scripts/generate-nightly-index.py
-$PYTHON .buildkite/scripts/generate-nightly-index.py --version "$SUBPATH" --current-objects "$obj_json" --output-dir "$INDICES_OUTPUT_DIR" --comment "commit $BUILDKITE_COMMIT" $alias_arg
+# generate index for nightly
+aws s3 cp "$wheel" "s3://vllm-wheels/nightly/"
+aws s3 cp "$normal_wheel" "s3://vllm-wheels/nightly/"

-# copy indices to /<commit>/ unconditionally
-echo "Uploading indices to $S3_COMMIT_PREFIX"
-aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "$S3_COMMIT_PREFIX"
-
-# copy to /nightly/ only if it is on the main branch and not a PR 
-if [[ "$BUILDKITE_BRANCH" == "main" && "$BUILDKITE_PULL_REQUEST" == "false" ]]; then
-    echo "Uploading indices to overwrite /nightly/"
-    aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/nightly/"
+if [[ $normal_wheel == *"cu129"* ]]; then
+    # only upload index.html for cu129 wheels (default wheels) as it
+    # is available on both x86 and arm64
+    aws s3 cp index.html "s3://vllm-wheels/nightly/vllm/index.html"
+else
+    echo "Skipping index files for non-cu129 wheels"
 fi

-# re-generate and copy to /<pure_version>/ only if it does not have "dev" in the version
-if [[ "$version" != *"dev"* ]]; then
-    echo "Re-generating indices for /$pure_version/"
-    rm -rf "$INDICES_OUTPUT_DIR/*"
-    mkdir -p "$INDICES_OUTPUT_DIR"
-    $PYTHON .buildkite/scripts/generate-nightly-index.py --version "$pure_version" --current-objects "$obj_json" --output-dir "$INDICES_OUTPUT_DIR" --comment "version $pure_version" $alias_arg
-    aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/$pure_version/"
-fi
+aws s3 cp "$wheel" "s3://vllm-wheels/$version/"
+aws s3 cp index.html "s3://vllm-wheels/$version/vllm/index.html"
@@ -39,9 +39,9 @@ steps:
  # if this test fails, it means the nightly torch version is not compatible with some
  # of the dependencies. Please check the error message and add the package to whitelist
  # in /vllm/tools/pre_commit/generate_nightly_torch_test.py
-  mirror_hardwares: [amdexperimental, amdproduction, amdtentative]
+  mirror_hardwares: [amdexperimental]
  agent_pool: mi325_1
-  grade: Blocking
+  # grade: Blocking
  soft_fail: true
  source_file_dependencies:
  - requirements/nightly_torch_test.txt
@@ -50,9 +50,9 @@ steps:

 - label: Async Engine, Inputs, Utils, Worker Test # 10min
  timeout_in_minutes: 15
-  mirror_hardwares: [amdexperimental, amdproduction, amdtentative]
+  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi325_1
-  grade: Blocking
+  # grade: Blocking
  source_file_dependencies:
  - vllm/
  - tests/multimodal
@@ -61,19 +61,17 @@ steps:
  - pytest -v -s -m 'not cpu_test' multimodal
  - pytest -v -s utils_

- label: Async Engine, Inputs, Utils, Worker, Config Test (CPU) # 20min
-  timeout_in_minutes: 30
-  mirror_hardwares: [amdexperimental, amdproduction, amdtentative]
+- label: Async Engine, Inputs, Utils, Worker, Config Test (CPU) # 4 mins
+  timeout_in_minutes: 10
+  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi325_1
-  grade: Blocking
+  # grade: Blocking
  source_file_dependencies:
  - vllm/
  - tests/test_inputs.py
  - tests/test_outputs.py
  - tests/multimodal
  - tests/standalone_tests/lazy_imports.py
-  - tests/tokenizers_
-  - tests/tool_parsers
  - tests/transformers_utils
  - tests/config
  no_gpu: true
@@ -82,8 +80,6 @@ steps:
  - pytest -v -s test_inputs.py
  - pytest -v -s test_outputs.py
  - pytest -v -s -m 'cpu_test' multimodal
-  - pytest -v -s tokenizers_
-  - pytest -v -s tool_parsers
  - pytest -v -s transformers_utils
  - pytest -v -s config

@@ -117,9 +113,9 @@ steps:
  - pytest -v -s basic_correctness/test_cpu_offload.py

 - label: Entrypoints Unit Tests # 5min
-  mirror_hardwares: [amdexperimental, amdproduction, amdtentative]
+  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi325_1
-  grade: Blocking
+  # grade: Blocking
  timeout_in_minutes: 10
  working_dir: "/vllm-workspace/tests"
  fast_check: true
@@ -216,7 +212,6 @@ steps:
  # test with internal dp
  - python3 ../examples/offline_inference/data_parallel.py --enforce-eager
  - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py
-  - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
  - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
  - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py
  - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py
@@ -255,9 +250,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: EPLB Algorithm Test # 5min
-  mirror_hardwares: [amdexperimental, amdproduction, amdtentative]
+  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi325_1
-  grade: Blocking
+  # grade: Blocking
  timeout_in_minutes: 15
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -313,25 +308,28 @@ steps:
  - pytest -v -s test_regression.py
  working_dir: "/vllm-workspace/tests" # optional

- label: Engine Test # 9min
-  timeout_in_minutes: 15
+- label: Engine Test # 25min
+  timeout_in_minutes: 40
  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi325_1
  # grade: Blocking
  source_file_dependencies:
  - vllm/
  - tests/engine
+  - tests/tokenization
  - tests/test_sequence
  - tests/test_config
  - tests/test_logger
  - tests/test_vllm_port
  commands:
  - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py
+  # OOM in the CI unless we run this separately
+  - pytest -v -s tokenization

- label: V1 Test e2e + engine # 65min
-  timeout_in_minutes: 90
+- label: V1 Test e2e + engine # 30min
+  timeout_in_minutes: 45
  mirror_hardwares: [amdexperimental]
-  agent_pool: mi325_4
+  agent_pool: mi325_1
  # grade: Blocking
  source_file_dependencies:
    - vllm/
@@ -344,9 +342,9 @@ steps:

 - label: V1 Test entrypoints # 35min
  timeout_in_minutes: 50
-  mirror_hardwares: [amdexperimental, amdproduction, amdtentative]
+  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi325_1
-  grade: Blocking
+  # grade: Blocking
  source_file_dependencies:
    - vllm/
    - tests/v1
@@ -394,21 +392,6 @@ steps:
  commands:
    - pytest -v -s v1/attention

- label: Batch Invariance Tests (H100) # 10min
-  mirror_hardwares: [amdexperimental]
-  agent_pool: mi325_1
-  timeout_in_minutes: 25
-  gpu: h100
-  source_file_dependencies:
-    - vllm/v1/attention
-    - vllm/model_executor/layers
-    - tests/v1/determinism/
-  commands:
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - pip install pytest-timeout pytest-forked
-    - pytest -v -s v1/determinism/test_batch_invariance.py
-    - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
-
 - label: V1 Test attention (B200) # 10min
  timeout_in_minutes: 30
  gpu: b200
@@ -419,9 +402,9 @@ steps:
    - VLLM_DISABLE_FLASHINFER_PREFILL=1 pytest -v -s v1/attention # TODO: FI prefill is bugged and causes incorrectness, fix this

 - label: V1 Test others (CPU) # 5 mins
-  mirror_hardwares: [amdexperimental, amdproduction, amdtentative]
+  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi325_1
-  grade: Blocking
+  # grade: Blocking
  source_file_dependencies:
    - vllm/
    - tests/v1
@@ -437,34 +420,29 @@ steps:

 - label: Examples Test # 30min
  timeout_in_minutes: 45
-  mirror_hardwares: [amdexperimental, amdproduction]
+  mirror_hardwares: [amdexperimental]
  agent_pool: mi325_1
  # grade: Blocking
  working_dir: "/vllm-workspace/examples"
  source_file_dependencies:
  - vllm/entrypoints
-  - vllm/multimodal
  - examples/
  commands:
    - pip install tensorizer # for tensorizer test
-    # for basic
-    - python3 offline_inference/basic/chat.py
    - python3 offline_inference/basic/generate.py --model facebook/opt-125m
    - python3 offline_inference/basic/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10
+    - python3 offline_inference/basic/chat.py
+    - python3 offline_inference/prefix_caching.py
+    - python3 offline_inference/llm_engine_example.py
+    - python3 offline_inference/audio_language.py --seed 0
+    - python3 offline_inference/vision_language.py --seed 0
+    - python3 offline_inference/vision_language_pooling.py --seed 0
+    - python3 offline_inference/vision_language_multi_image.py --seed 0
+    - python3 others/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 others/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
+    - python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
    - python3 offline_inference/basic/classify.py
    - python3 offline_inference/basic/embed.py
    - python3 offline_inference/basic/score.py
-    # for multi-modal models
-    - python3 offline_inference/audio_language.py --seed 0
-    - python3 offline_inference/vision_language.py --seed 0
-    - 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/pooling/vision_language_pooling.py --seed 0
-    # for features demo
-    - python3 offline_inference/prefix_caching.py
-    - python3 offline_inference/llm_engine_example.py
-    - python3 others/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 others/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
    - python3 offline_inference/spec_decode.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
    # https://github.com/vllm-project/vllm/pull/26682 uses slightly more memory in PyTorch 2.9+ causing this test to OOM in 1xL4 GPU
    - python3 offline_inference/spec_decode.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
@@ -518,7 +496,7 @@ steps:

 - label: PyTorch Compilation Unit Tests # 15min
  timeout_in_minutes: 30
-  mirror_hardwares: [amdexperimental, amdproduction]
+  mirror_hardwares: [amdexperimental]
  agent_pool: mi325_1
  # grade: Blocking
  torch_nightly: true
@@ -535,7 +513,7 @@ steps:

 - label: PyTorch Fullgraph Smoke Test # 15min
  timeout_in_minutes: 30
-  mirror_hardwares: [amdexperimental, amdproduction]
+  mirror_hardwares: [amdexperimental]
  agent_pool: mi325_1
  # grade: Blocking
  torch_nightly: true
@@ -591,7 +569,7 @@ steps:

 - label: Kernels Attention Test %N # 23min
  timeout_in_minutes: 35
-  mirror_hardwares: [amdexperimental, amdproduction]
+  mirror_hardwares: [amdexperimental]
  agent_pool: mi325_8
  # grade: Blocking
  source_file_dependencies:
@@ -618,7 +596,7 @@ steps:

 - label: Kernels MoE Test %N # 40min
  timeout_in_minutes: 60
-  mirror_hardwares: [amdexperimental, amdproduction]
+  mirror_hardwares: [amdexperimental]
  agent_pool: mi325_8
  # grade: Blocking
  source_file_dependencies:
@@ -645,26 +623,6 @@ steps:
  commands:
    - pytest -v -s kernels/mamba

- label: Kernels DeepGEMM Test (H100) # Nvidia-centric
-# Not replicating for CUTLAS & CuTe
-  timeout_in_minutes: 45
-  gpu: h100
-  num_gpus: 1
-  source_file_dependencies:
-  - tools/install_deepgemm.sh
-  - vllm/utils/deep_gemm.py
-  - vllm/model_executor/layers/fused_moe
-  - vllm/model_executor/layers/quantization
-  - tests/kernels/quantization/test_block_fp8.py
-  - tests/kernels/moe/test_deepgemm.py
-  - tests/kernels/moe/test_batched_deepgemm.py
-  - tests/kernels/attention/test_deepgemm_attention.py
-  commands:
-    - pytest -v -s kernels/quantization/test_block_fp8.py -k deep_gemm
-    - pytest -v -s kernels/moe/test_deepgemm.py
-    - pytest -v -s kernels/moe/test_batched_deepgemm.py
-    - pytest -v -s kernels/attention/test_deepgemm_attention.py
-
 - label: Model Executor Test # 23min
  timeout_in_minutes: 35
  torch_nightly: true
@@ -723,18 +681,16 @@ steps:
  # we can only upgrade after this is resolved
  # TODO(jerryzh168): resolve the above comment
  - uv pip install --system torchao==0.13.0
-  - uv pip install --system conch-triton-kernels
  - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py

- label: LM Eval Small Models # 53min
-  timeout_in_minutes: 75
-  mirror_hardwares: [amdexperimental]
+- label: LM Eval Small Models # 15min
+  timeout_in_minutes: 20
+  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi325_1
  # grade: Blocking
  source_file_dependencies:
  - csrc/
  - vllm/model_executor/layers/quantization
-  autorun_on_main: true
  commands:
  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.txt --tp-size=1

@@ -747,7 +703,7 @@ steps:
  - csrc/
  - vllm/entrypoints/openai/
  - vllm/model_executor/models/whisper.py
-  commands: # LMEval+Transcription WER check
+  commands: # LMEval
  # Transcription WER check is skipped because encoder-decoder models are not supported on ROCm, see https://github.com/vllm-project/vllm/issues/27442
  - pytest -s entrypoints/openai/correctness/

@@ -761,7 +717,19 @@ steps:
    - vllm/
    - tests/tool_use
  commands:
-    - pytest -v -s tool_use
+    - pytest -v -s -m 'not cpu_test' tool_use
+
+- label: OpenAI-Compatible Tool Use (CPU) # 5 mins
+  mirror_hardwares: [amdexperimental, amdproduction]
+  agent_pool: mi325_1
+  # grade: Blocking
+  timeout_in_minutes: 10
+  source_file_dependencies:
+    - vllm/
+    - tests/tool_use
+  no_gpu: true
+  commands:
+    - pytest -v -s -m 'cpu_test' tool_use

 #####  models test  #####

@@ -786,7 +754,6 @@ steps:
  torch_nightly: true
  source_file_dependencies:
  - vllm/model_executor/models/
-  - vllm/transformers_utils/
  - tests/models/test_initialization.py
  commands:
    # Only when vLLM model source is modified - test initialization of a large
@@ -932,18 +899,6 @@ steps:
  commands:
    - pytest -v -s models/language/pooling_mteb_test

- label: Multi-Modal Processor Test (CPU)
-  timeout_in_minutes: 60
-  mirror_hardwares: [amdexperimental]
-  agent_pool: mi325_1
-  source_file_dependencies:
-  - vllm/
-  - tests/models/multimodal
-  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
-
 - label: Multi-Modal Processor Test # 44min
  timeout_in_minutes: 60
  mirror_hardwares: [amdexperimental]
@@ -971,8 +926,8 @@ steps:
    - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/processing
    - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model  # Otherwise, mp_method="spawn" doesn't work

- label: Multi-Modal Accuracy Eval (Small Models) # 150min - 180min
-  timeout_in_minutes: 180
+- label: Multi-Modal Accuracy Eval (Small Models) # 10min
+  timeout_in_minutes: 70
  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi325_1
  # grade: Blocking
@@ -984,8 +939,7 @@ steps:
  commands:
  - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-mm-small.txt --tp-size=1

- label: Multi-Modal Models Test (Extended) 1 # 60min
-  timeout_in_minutes: 120
+- label: Multi-Modal Models Test (Extended) 1
  mirror_hardwares: [amdexperimental]
  agent_pool: mi325_1
  # grade: Blocking
@@ -1009,8 +963,7 @@ steps:
    - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
    - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=0) and not core_model'

- label: Multi-Modal Models Test (Extended) 3 # 75min
-  timeout_in_minutes: 150
+- label: Multi-Modal Models Test (Extended) 3
  mirror_hardwares: [amdexperimental]
  agent_pool: mi325_1
  # grade: Blocking
@@ -1102,7 +1055,6 @@ steps:
    - pytest -v -s tests/kernels/moe/test_nvfp4_moe.py
    - pytest -v -s tests/kernels/moe/test_ocp_mx_moe.py
    - pytest -v -s tests/kernels/moe/test_flashinfer.py
-    - pytest -v -s tests/kernels/moe/test_cutedsl_moe.py

 - label: Blackwell Fusion and Compile Tests # 30 min
  timeout_in_minutes: 40
@@ -1112,18 +1064,11 @@ steps:
  - csrc/quantization/fp4/
  - 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/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:
    - nvidia-smi
    - pytest -v -s tests/compile/test_fusion_attn.py
@@ -1134,7 +1079,7 @@ steps:
    # 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
+    - pytest -v -s tests/compile/distributed/test_full_graph.py::test_fp8_kv_scale_compile

 - label: Blackwell Fusion E2E Tests # 30 min
  timeout_in_minutes: 40
@@ -1152,15 +1097,17 @@ steps:
  - vllm/model_executor/layers/activation.py
  - vllm/model_executor/layers/quantization/input_quant_fp8.py
  - tests/compile/distributed/test_fusions_e2e.py
+  - tests/compile/fullgraph/test_full_graph.py
  commands:
    - nvidia-smi
    # Run all e2e fusion tests
-    - pytest -v -s tests/compile/distributed/test_fusions_e2e.py
+    - pytest -v -s tests/compile/test_fusions_e2e.py

- label: Blackwell GPT-OSS Eval
+- label: ROCm GPT-OSS Eval
  timeout_in_minutes: 60
  working_dir: "/vllm-workspace/"
-  gpu: b200
+  agent_pool: mi325_1
+  mirror_hardwares: [amdexperimental, amdproduction]
  optional: true # run on nightlies
  source_file_dependencies:
  - tests/evals/gpt_oss
@@ -1169,7 +1116,7 @@ steps:
  - vllm/v1/attention/backends/flashinfer.py
  commands:
    - uv pip install --system 'gpt-oss[eval]==0.0.5'
-    - pytest -s -v tests/evals/gpt_oss/test_gpqa_correctness.py --model openai/gpt-oss-20b --metric 0.58
+    - VLLM_ROCM_USE_AITER_MHA=0 VLLM_ROCM_USE_AITER=1 VLLM_USE_AITER_UNIFIED_ATTENTION=1 pytest -s -v tests/evals/gpt_oss/test_gpqa_correctness.py --model openai/gpt-oss-20b --metric 0.58

 - label: Blackwell Quantized MoE Test
  timeout_in_minutes: 60
@@ -1269,7 +1216,6 @@ steps:
  - tests/v1/worker/test_worker_memory_snapshot.py
  commands:
  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py
-  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
  - DP_SIZE=2 pytest -v -s v1/entrypoints/openai/test_multi_api_servers.py
  - pytest -v -s entrypoints/llm/test_collective_rpc.py
@@ -1305,7 +1251,7 @@ steps:

 - label: Plugin Tests (2 GPUs) # 40min
  timeout_in_minutes: 60
-  mirror_hardwares: [amdexperimental, amdproduction]
+  mirror_hardwares: [amdexperimental]
  agent_pool: mi325_2
  # grade: Blocking
  working_dir: "/vllm-workspace/tests"
@@ -1373,15 +1319,12 @@ steps:
    - pytest -v -s -x lora/test_llama_tp.py
    - pytest -v -s -x lora/test_llm_with_multi_loras.py
    - pytest -v -s -x lora/test_olmoe_tp.py
-
-    # Disabled for now because MXFP4 backend on non-cuda platform
-    # doesn't support LoRA yet
-    #- pytest -v -s -x lora/test_gptoss_tp.py
+    - pytest -v -s -x lora/test_gptoss_tp.py


 - label: Weight Loading Multiple GPU Test  # 33min
  timeout_in_minutes: 45
-  mirror_hardwares: [amdexperimental, amdproduction]
+  mirror_hardwares: [amdexperimental]
  agent_pool: mi325_2
  # grade: Blocking
  working_dir: "/vllm-workspace/tests"
@@ -1440,13 +1383,12 @@ steps:
  - TARGET_TEST_SUITE=A100 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
  - pytest -v -s -x lora/test_mixtral.py

-
 - label: LM Eval Large Models # optional
-  gpu: a100
-  optional: true
-  mirror_hardwares: [amdexperimental]
+  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi325_4
  # grade: Blocking
+  gpu: a100
+  optional: true
  num_gpus: 4
  working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
  source_file_dependencies:
@@ -1458,11 +1400,11 @@ steps:

 ##### H100 test #####
 - label: LM Eval Large Models (H100) # optional
-  gpu: h100
-  optional: true
-  mirror_hardwares: [amdexperimental]
+  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi325_4
  # grade: Blocking
+  gpu: h100
+  optional: true
  num_gpus: 4
  working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
  source_file_dependencies:
@@ -1472,7 +1414,6 @@ steps:
    - export VLLM_USE_DEEP_GEMM=0  # We found Triton is faster than DeepGEMM for H100
    - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-hopper.txt --tp-size=4

-
 ##### H200 test #####
 - label: Distributed Tests (H200) # optional
  mirror_hardwares: [amdexperimental]
@@ -1483,14 +1424,14 @@ steps:
  working_dir: "/vllm-workspace/"
  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_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'"
-    - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/distributed/test_sequence_parallel.py
+    - "pytest -v -s tests/compile/distributed/test_fusions_e2e.py -k 'not Llama-4'"
+    - pytest -v -s tests/compile/distributed/test_sequence_parallel.py
    - pytest -v -s tests/distributed/test_context_parallel.py
-    - HIP_VISIBLE_DEVICES=0,1 VLLM_ALL2ALL_BACKEND=deepep_high_throughput VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model Qwen/Qwen1.5-MoE-A2.7B --tp-size=1  --dp-size=2 --max-model-len 2048
+    - CUDA_VISIBLE_DEVICES=1,2 VLLM_ALL2ALL_BACKEND=deepep_high_throughput VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model Qwen/Qwen1.5-MoE-A2.7B --tp-size=1  --dp-size=2 --max-model-len 2048
    - pytest -v -s tests/v1/distributed/test_dbo.py

 ##### B200 test #####
@@ -1504,57 +1445,6 @@ steps:
    - pytest -v -s tests/distributed/test_nccl_symm_mem_allreduce.py
    - pytest -v -s tests/v1/distributed/test_dbo.py

-##### E2E Eval Tests #####
- label: LM Eval Small Models (1 Card) # 15min
-  timeout_in_minutes: 20
-  mirror_hardwares: [amdexperimental, amdproduction]
-  agent_pool: mi325_1
-  # grade: Blocking
-  source_file_dependencies:
-  - csrc/
-  - vllm/model_executor/layers/quantization
-  commands:
-  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.txt --tp-size=1
-
- label: LM Eval Large Models (4 Card)
-  mirror_hardwares: [amdexperimental, amdproduction]
-  agent_pool: mi325_4
-  # grade: Blocking
-  gpu: a100
-  optional: true
-  num_gpus: 4
-  working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
-  source_file_dependencies:
-  - csrc/
-  - vllm/model_executor/layers/quantization
-  commands:
-  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large.txt --tp-size=4
-
- label: ROCm LM Eval Large Models (8 Card)
-  mirror_hardwares: [amdproduction]
-  agent_pool: mi325_8
-  num_gpus: 8
-  working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
-  commands:
-  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm.txt --tp-size=8
-
- label: ROCm GPT-OSS Eval
-  timeout_in_minutes: 60
-  working_dir: "/vllm-workspace/"
-  agent_pool: mi325_1
-  mirror_hardwares: [amdexperimental, amdproduction]
-  optional: true # run on nightlies
-  source_file_dependencies:
-  - tests/evals/gpt_oss
-  - vllm/model_executor/models/gpt_oss.py
-  - vllm/model_executor/layers/quantization/mxfp4.py
-  - vllm/v1/attention/backends/flashinfer.py
-  commands:
-    - uv pip install --system 'gpt-oss[eval]==0.0.5'
-    - VLLM_ROCM_USE_AITER_MHA=0 VLLM_ROCM_USE_AITER=1 VLLM_USE_AITER_UNIFIED_ATTENTION=1 pytest -s -v tests/evals/gpt_oss/test_gpqa_correctness.py --model openai/gpt-oss-20b --metric 0.58
-
 ##### RL Integration Tests #####
 - label: Prime-RL Integration Test # 15min
  mirror_hardwares: [amdexperimental]
@@ -1569,8 +1459,9 @@ steps:
  - .buildkite/scripts/run-prime-rl-test.sh
  commands:
    - bash .buildkite/scripts/run-prime-rl-test.sh
+
 - label: DeepSeek V2-Lite Accuracy
-  mirror_hardwares: [amdexperimental, amdproduction]
+  mirror_hardwares: [amdexperimental]
  agent_pool: mi325_4
  # grade: Blocking
  timeout_in_minutes: 60
@@ -1581,8 +1472,8 @@ steps:
  commands:
  - bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh 0.25 200 8010

- label: Qwen3-30B-A3B-FP8-block Accuracy (H100)
-  mirror_hardwares: [amdexperimental, amdproduction]
+- label: Qwen3-30B-A3B-FP8-block Accuracy
+  mirror_hardwares: [amdexperimental]
  agent_pool: mi325_4
  # grade: Blocking
  timeout_in_minutes: 60
@@ -1591,36 +1482,4 @@ steps:
  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
-  gpu: b200
-  optional: true
-  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
-
- label: DeepSeek V2-Lite Async EPLB Accuracy
-  timeout_in_minutes: 60
-  mirror_hardwares: [amdexperimental]
-  agent_pool: mi325_4
-  # grade: Blocking
-  gpu: h100
-  optional: true
-  num_gpus: 4
-  working_dir: "/vllm-workspace"
-  commands:
-  - bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_async_eplb.sh 0.25 1319 8030
-
- label: Qwen3-Next-80B-A3B-Instruct MTP Async EPLB Accuracy
-  timeout_in_minutes: 60
-  mirror_hardwares: [amdexperimental]
-  agent_pool: mi325_4
-  # grade: Blocking
-  optional: true
-  num_gpus: 4
-  working_dir: "/vllm-workspace"
-  commands:
-  - bash .buildkite/scripts/scheduled_integration_test/qwen3_next_mtp_async_eplb.sh 0.8 1319 8040
+  - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep.sh 0.8 200 8020
@@ -57,16 +57,14 @@ steps:
  - pytest -v -s -m 'not cpu_test' multimodal
  - pytest -v -s utils_

- label: Async Engine, Inputs, Utils, Worker, Config Test (CPU) # 20min
-  timeout_in_minutes: 30
+- label: Async Engine, Inputs, Utils, Worker, Config Test (CPU) # 4 mins
+  timeout_in_minutes: 10
  source_file_dependencies:
  - vllm/
  - tests/test_inputs.py
  - tests/test_outputs.py
  - tests/multimodal
  - tests/standalone_tests/lazy_imports.py
-  - tests/tokenizers_
-  - tests/tool_parsers
  - tests/transformers_utils
  - tests/config
  no_gpu: true
@@ -75,8 +73,6 @@ steps:
  - pytest -v -s test_inputs.py
  - pytest -v -s test_outputs.py
  - pytest -v -s -m 'cpu_test' multimodal
-  - pytest -v -s tokenizers_
-  - pytest -v -s tool_parsers
  - pytest -v -s transformers_utils
  - pytest -v -s config

@@ -196,7 +192,6 @@ steps:
  # test with internal dp
  - python3 ../examples/offline_inference/data_parallel.py --enforce-eager
  - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py
-  - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
  - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
  - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py
  - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py
@@ -280,18 +275,21 @@ steps:
  - pytest -v -s test_regression.py
  working_dir: "/vllm-workspace/tests" # optional

- label: Engine Test # 9min
-  timeout_in_minutes: 15
+- label: Engine Test # 25min
+  timeout_in_minutes: 40
  mirror_hardwares: [amdexperimental]
  source_file_dependencies:
  - vllm/
  - tests/engine
+  - tests/tokenization
  - tests/test_sequence
  - tests/test_config
  - tests/test_logger
  - tests/test_vllm_port
  commands:
  - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py
+  # OOM in the CI unless we run this separately
+  - pytest -v -s tokenization

 - label: V1 Test e2e + engine # 30min
  timeout_in_minutes: 45
@@ -348,19 +346,6 @@ steps:
  commands:
    - pytest -v -s v1/attention

- label: Batch Invariance Tests (H100) # 10min
-  timeout_in_minutes: 25
-  gpu: h100
-  source_file_dependencies:
-    - vllm/v1/attention
-    - vllm/model_executor/layers
-    - tests/v1/determinism/
-  commands:
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - pip install pytest-timeout pytest-forked
-    - pytest -v -s v1/determinism/test_batch_invariance.py
-    - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
-
 - label: V1 Test attention (B200) # 10min
  timeout_in_minutes: 30
  gpu: b200
@@ -390,28 +375,23 @@ steps:
  working_dir: "/vllm-workspace/examples"
  source_file_dependencies:
  - vllm/entrypoints
-  - vllm/multimodal
  - examples/
  commands:
    - pip install tensorizer # for tensorizer test
-    # for basic
-    - python3 offline_inference/basic/chat.py
    - python3 offline_inference/basic/generate.py --model facebook/opt-125m
    - python3 offline_inference/basic/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10
+    - python3 offline_inference/basic/chat.py
+    - python3 offline_inference/prefix_caching.py
+    - python3 offline_inference/llm_engine_example.py
+    - python3 offline_inference/audio_language.py --seed 0
+    - python3 offline_inference/vision_language.py --seed 0
+    - python3 offline_inference/vision_language_pooling.py --seed 0
+    - python3 offline_inference/vision_language_multi_image.py --seed 0
+    - python3 others/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 others/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
+    - python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
    - python3 offline_inference/basic/classify.py
    - python3 offline_inference/basic/embed.py
    - python3 offline_inference/basic/score.py
-    # for multi-modal models
-    - python3 offline_inference/audio_language.py --seed 0
-    - python3 offline_inference/vision_language.py --seed 0
-    - 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/pooling/vision_language_pooling.py --seed 0
-    # for features demo
-    - python3 offline_inference/prefix_caching.py
-    - python3 offline_inference/llm_engine_example.py
-    - python3 others/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 others/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
    - python3 offline_inference/spec_decode.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
    # https://github.com/vllm-project/vllm/pull/26682 uses slightly more memory in PyTorch 2.9+ causing this test to OOM in 1xL4 GPU
    - python3 offline_inference/spec_decode.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
@@ -470,9 +450,7 @@ steps:
  # tests covered elsewhere.
  # Use `find` to launch multiple instances of pytest so that
  # they do not suffer from https://github.com/vllm-project/vllm/issues/28965
-  # 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 '{}'"
+  - "find compile/ -maxdepth 1 -name 'test_*.py' -exec pytest -s -v {} \\\\;"

 - label: PyTorch Fullgraph Smoke Test # 15min
  timeout_in_minutes: 30
@@ -486,9 +464,7 @@ steps:
  # as it is a heavy test that is covered in other steps.
  # Use `find` to launch multiple instances of pytest so that
  # they do not suffer from https://github.com/vllm-project/vllm/issues/28965
-  # However, find does not normally propagate error codes, so we combine it with xargs
-  # (using -0 for proper path handling)
-  - "find compile/fullgraph -maxdepth 1 -name 'test_*.py' -not -name 'test_full_graph.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
+  - "find compile/fullgraph/ -name 'test_*.py' -not -name 'test_full_graph.py' -exec pytest -s -v {} \\\\;"

 - label: PyTorch Fullgraph Test # 27min
  timeout_in_minutes: 40
@@ -578,6 +554,7 @@ steps:
  timeout_in_minutes: 45
  gpu: h100
  num_gpus: 1
+  optional: true
  source_file_dependencies:
  - tools/install_deepgemm.sh
  - vllm/utils/deep_gemm.py
@@ -588,10 +565,10 @@ steps:
  - tests/kernels/moe/test_batched_deepgemm.py
  - tests/kernels/attention/test_deepgemm_attention.py
  commands:
-    - pytest -v -s kernels/quantization/test_block_fp8.py -k deep_gemm
-    - pytest -v -s kernels/moe/test_deepgemm.py
-    - pytest -v -s kernels/moe/test_batched_deepgemm.py
-    - pytest -v -s kernels/attention/test_deepgemm_attention.py
+    - pytest -v -s tests/kernels/quantization/test_block_fp8.py -k deep_gemm
+    - pytest -v -s tests/kernels/moe/test_deepgemm.py
+    - pytest -v -s tests/kernels/moe/test_batched_deepgemm.py
+    - pytest -v -s tests/kernels/attention/test_deepgemm_attention.py

 - label: Model Executor Test # 23min
  timeout_in_minutes: 35
@@ -643,7 +620,6 @@ steps:
  # we can only upgrade after this is resolved
  # TODO(jerryzh168): resolve the above comment
  - 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: LM Eval Small Models # 53min
@@ -674,7 +650,16 @@ steps:
    - vllm/
    - tests/tool_use
  commands:
-    - pytest -v -s tool_use
+    - pytest -v -s -m 'not cpu_test' tool_use
+
+- label: OpenAI-Compatible Tool Use (CPU) # 5 mins
+  timeout_in_minutes: 10
+  source_file_dependencies:
+    - vllm/
+    - tests/tool_use
+  no_gpu: true
+  commands:
+    - pytest -v -s -m 'cpu_test' tool_use

 #####  models test  #####

@@ -685,7 +670,6 @@ steps:
  source_file_dependencies:
  - vllm/
  - tests/models/test_initialization.py
-  - tests/models/registry.py
  commands:
    # Run a subset of model initialization tests
    - pytest -v -s models/test_initialization.py::test_can_initialize_small_subset
@@ -696,9 +680,7 @@ steps:
  torch_nightly: true
  source_file_dependencies:
  - vllm/model_executor/models/
-  - vllm/transformers_utils/
  - tests/models/test_initialization.py
-  - tests/models/registry.py
  commands:
    # Only when vLLM model source is modified - test initialization of a large
    # subset of supported models (the complement of the small subset in the above
@@ -824,24 +806,14 @@ steps:
  commands:
    - pytest -v -s models/language/pooling_mteb_test

- label: Multi-Modal Processor Test (CPU)
-  timeout_in_minutes: 60
-  source_file_dependencies:
-  - vllm/
-  - tests/models/multimodal
-  no_gpu: true
-  commands:
-    - "pip install git+https://github.com/TIGER-AI-Lab/Mantis.git || echo 'Mantis installation skipped (decord not available on CPU-only environment)'"
-    - pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py
-
- label: Multi-Modal Processor Test
+- label: Multi-Modal Processor Test # 44min
  timeout_in_minutes: 60
  source_file_dependencies:
  - vllm/
  - tests/models/multimodal
  commands:
    - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
-    - pytest -v -s models/multimodal/processing/test_tensor_schema.py
+    - pytest -v -s models/multimodal/processing

 - label: Multi-Modal Models Test (Standard) # 60min
  timeout_in_minutes: 80
@@ -918,12 +890,11 @@ steps:
 - label: Transformers Nightly Models Test
  working_dir: "/vllm-workspace/"
  optional: true
-  soft_fail: true
  commands:
    - pip install --upgrade git+https://github.com/huggingface/transformers
-    - pytest -v -s tests/models/test_initialization.py
+    - pytest -v -s tests/models/test_initialization.py -k 'not (Ultravox or Phi4Multimodal or MiniCPMO or Lfm2Moe or RobertaForSequenceClassification or Ovis2_5 or DeepseekOCR or KimiVL)'
    - pytest -v -s tests/models/test_transformers.py
-    - pytest -v -s tests/models/multimodal/processing/
+    # - pytest -v -s tests/models/multimodal/processing/
    - pytest -v -s tests/models/multimodal/test_mapping.py
    - python3 examples/offline_inference/basic/chat.py
    - python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
@@ -1133,7 +1104,6 @@ steps:
  # https://github.com/NVIDIA/nccl/issues/1838
  - export NCCL_CUMEM_HOST_ENABLE=0
  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py
-  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
  - DP_SIZE=2 pytest -v -s v1/entrypoints/openai/test_multi_api_servers.py
  - pytest -v -s entrypoints/llm/test_collective_rpc.py
@@ -1223,8 +1193,6 @@ steps:
    # FIXIT: find out which code initialize cuda before running the test
    # before the fix, we need to use spawn to test it
    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    # Alot of these tests are on the edge of OOMing
-    - export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
    # There is some Tensor Parallelism related processing logic in LoRA that
    # requires multi-GPU testing for validation.
    - pytest -v -s -x lora/test_chatglm3_tp.py
@@ -1319,11 +1287,11 @@ steps:
  working_dir: "/vllm-workspace/"
  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_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/compile/distributed/test_fusions_e2e.py -k 'not Llama-4'"
+    - pytest -v -s tests/distributed/test_sequence_parallel.py
    - pytest -v -s tests/distributed/test_context_parallel.py
    - CUDA_VISIBLE_DEVICES=1,2 VLLM_ALL2ALL_BACKEND=deepep_high_throughput VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model Qwen/Qwen1.5-MoE-A2.7B --tp-size=1  --dp-size=2 --max-model-len 2048
    - pytest -v -s tests/v1/distributed/test_dbo.py
@@ -1343,7 +1311,6 @@ steps:
 - label: Prime-RL Integration Test # 15min
  timeout_in_minutes: 30
  optional: true
-  soft_fail: true
  num_gpus: 2
  working_dir: "/vllm-workspace"
  source_file_dependencies:
@@ -1361,20 +1328,11 @@ steps:
  commands:
  - bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh 0.25 200 8010

- label: Qwen3-30B-A3B-FP8-block Accuracy (H100)
+- label: Qwen3-30B-A3B-FP8-block Accuracy
  timeout_in_minutes: 60
  gpu: h100
  optional: true
  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
-  gpu: b200
-  optional: true
-  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
+  - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep.sh 0.8 200 8020
@@ -1,21 +0,0 @@
-group: Attention
-depends_on: 
-  - image-build
-steps:
- label: V1 attention (H100)
-  timeout_in_minutes: 30
-  gpu: h100
-  source_file_dependencies:
-    - vllm/v1/attention
-    - tests/v1/attention
-  commands:
-    - pytest -v -s v1/attention
-
- label: V1 attention (B200)
-  timeout_in_minutes: 30
-  gpu: b200
-  source_file_dependencies:
-    - vllm/v1/attention
-    - tests/v1/attention
-  commands:
-    - VLLM_DISABLE_FLASHINFER_PREFILL=1 pytest -v -s v1/attention # TODO: FI prefill is bugged and causes incorrectness, fix this
@@ -1,16 +0,0 @@
-group: Basic Correctness
-depends_on: 
-  - image-build
-steps:
- label: Basic Correctness
-  timeout_in_minutes: 30
-  source_file_dependencies:
-  - vllm/
-  - tests/basic_correctness/test_basic_correctness
-  - tests/basic_correctness/test_cpu_offload
-  - tests/basic_correctness/test_cumem.py
-  commands:
-  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - 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
@@ -1,19 +0,0 @@
-group: Benchmarks
-depends_on: 
-  - image-build
-steps:
- label: Benchmarks
-  timeout_in_minutes: 20
-  working_dir: "/vllm-workspace/.buildkite"
-  source_file_dependencies:
-  - benchmarks/
-  commands:
-  - bash scripts/run-benchmarks.sh
-
- label: Benchmarks CLI Test
-  timeout_in_minutes: 20
-  source_file_dependencies:
-  - vllm/
-  - tests/benchmarks/
-  commands:
-  - pytest -v -s benchmarks/
@@ -1,57 +0,0 @@
-group: Compile
-depends_on: 
-  - image-build
-steps:
- label: Fusion and Compile Tests (B200)
-  timeout_in_minutes: 40
-  working_dir: "/vllm-workspace/"
-  gpu: b200
-  source_file_dependencies:
-  - csrc/quantization/fp4/
-  - 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/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:
-    - 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/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: Fusion E2E (2 GPUs)(B200)
-  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
-
@@ -1,22 +0,0 @@
-group: CUDA
-depends_on: 
-  - image-build
-steps:
- label: Platform Tests (CUDA)
-  timeout_in_minutes: 15
-  source_file_dependencies:
-  - vllm/
-  - tests/cuda
-  commands:
-    - pytest -v -s cuda/test_cuda_context.py
-
- label: Cudagraph
-  timeout_in_minutes: 20
-  source_file_dependencies:
-  - tests/v1/cudagraph
-  - vllm/v1/cudagraph_dispatcher.py
-  - vllm/config/compilation.py
-  - vllm/compilation
-  commands:
-    - pytest -v -s v1/cudagraph/test_cudagraph_dispatch.py
-    - pytest -v -s v1/cudagraph/test_cudagraph_mode.py
@@ -1,199 +0,0 @@
-group: Distributed
-depends_on: 
-  - image-build
-steps:
- label: Distributed Comm Ops
-  timeout_in_minutes: 20
-  working_dir: "/vllm-workspace/tests"
-  num_gpus: 2
-  source_file_dependencies:
-  - vllm/distributed
-  - tests/distributed
-  commands:
-  - pytest -v -s distributed/test_comm_ops.py
-  - pytest -v -s distributed/test_shm_broadcast.py
-  - pytest -v -s distributed/test_shm_buffer.py
-  - pytest -v -s distributed/test_shm_storage.py
-
- label: Distributed (2 GPUs)
-  timeout_in_minutes: 90
-  working_dir: "/vllm-workspace/tests"
-  num_gpus: 2
-  source_file_dependencies:
-  - vllm/compilation/
-  - vllm/distributed/
-  - vllm/engine/
-  - vllm/executor/
-  - vllm/worker/worker_base.py
-  - vllm/v1/engine/
-  - vllm/v1/worker/
-  - tests/compile/fullgraph/test_basic_correctness.py
-  - tests/compile/test_wrapper.py
-  - tests/distributed/
-  - tests/entrypoints/llm/test_collective_rpc.py
-  - tests/v1/distributed
-  - tests/v1/entrypoints/openai/test_multi_api_servers.py
-  - tests/v1/shutdown
-  - tests/v1/worker/test_worker_memory_snapshot.py
-  commands:
-  # https://github.com/NVIDIA/nccl/issues/1838
-  - export NCCL_CUMEM_HOST_ENABLE=0
-  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py
-  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
-  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
-  - DP_SIZE=2 pytest -v -s v1/entrypoints/openai/test_multi_api_servers.py
-  - pytest -v -s entrypoints/llm/test_collective_rpc.py
-  - pytest -v -s ./compile/fullgraph/test_basic_correctness.py
-  - pytest -v -s ./compile/test_wrapper.py
-  - 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_gpus: 4
-  source_file_dependencies:
-  - vllm/distributed/
-  - tests/distributed/test_utils
-  - tests/distributed/test_pynccl
-  - tests/distributed/test_events
-  - tests/compile/fullgraph/test_basic_correctness.py
-  - examples/offline_inference/rlhf.py
-  - examples/offline_inference/rlhf_colocate.py
-  - tests/examples/offline_inference/data_parallel.py
-  - tests/v1/distributed
-  - tests/v1/engine/test_engine_core_client.py
-  - tests/distributed/test_symm_mem_allreduce.py
-  commands:
-  # https://github.com/NVIDIA/nccl/issues/1838
-  - export NCCL_CUMEM_HOST_ENABLE=0
-  # test with torchrun tp=2 and external_dp=2
-  - torchrun --nproc-per-node=4 distributed/test_torchrun_example.py
-  # test with torchrun tp=2 and pp=2
-  - PP_SIZE=2 torchrun --nproc-per-node=4 distributed/test_torchrun_example.py
-  # test with torchrun tp=4 and dp=1
-  - TP_SIZE=4 torchrun --nproc-per-node=4 distributed/test_torchrun_example_moe.py
-  # test with torchrun tp=2, pp=2 and dp=1
-  - PP_SIZE=2 TP_SIZE=2 torchrun --nproc-per-node=4 distributed/test_torchrun_example_moe.py
-  # test with torchrun tp=1 and dp=4 with ep
-  - DP_SIZE=4 ENABLE_EP=1 torchrun --nproc-per-node=4 distributed/test_torchrun_example_moe.py
-  # test with torchrun tp=2 and dp=2 with ep
-  - TP_SIZE=2 DP_SIZE=2 ENABLE_EP=1 torchrun --nproc-per-node=4 distributed/test_torchrun_example_moe.py
-  # test with internal dp
-  - python3 ../examples/offline_inference/data_parallel.py --enforce-eager
-  - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py
-  - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
-  - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
-  - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py
-  - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py
-  - pytest -v -s v1/engine/test_engine_core_client.py::test_kv_cache_events_dp
-  - pytest -v -s distributed/test_utils.py
-  - pytest -v -s compile/fullgraph/test_basic_correctness.py
-  - pytest -v -s distributed/test_pynccl.py
-  - pytest -v -s distributed/test_events.py
-  - 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
-  - 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
-
- label: Distributed Tests (8 GPUs)(H100)
-  timeout_in_minutes: 10
-  gpu: h100
-  num_gpus: 8
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - examples/offline_inference/torchrun_dp_example.py
-  - vllm/config/parallel.py
-  - vllm/distributed/
-  - vllm/v1/engine/llm_engine.py
-  - vllm/v1/executor/uniproc_executor.py
-  - vllm/v1/worker/gpu_worker.py
-  commands:
-  # https://github.com/NVIDIA/nccl/issues/1838
-  - export NCCL_CUMEM_HOST_ENABLE=0
-  # test with torchrun tp=2 and dp=4 with ep
-  - 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)
-  gpu: a100
-  optional: true
-  num_gpus: 4
-  source_file_dependencies:
-  - vllm/
-  commands:
-  # 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
-  - torchrun --nproc_per_node=2 distributed/test_ca_buffer_sharing.py
-  - TARGET_TEST_SUITE=A100 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
-  - pytest -v -s -x lora/test_mixtral.py
-
- label: Distributed Tests (2 GPUs)(H200)
-  gpu: h200
-  optional: true
-  working_dir: "/vllm-workspace/"
-  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
-    - CUDA_VISIBLE_DEVICES=1,2 VLLM_ALL2ALL_BACKEND=deepep_high_throughput VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model Qwen/Qwen1.5-MoE-A2.7B --tp-size=1  --dp-size=2 --max-model-len 2048
-    - pytest -v -s tests/v1/distributed/test_dbo.py
-
- label: Distributed Tests (2 GPUs)(B200)
-  gpu: b200
-  optional: true
-  working_dir: "/vllm-workspace/"
-  num_gpus: 2
-  commands:
-    - pytest -v -s tests/distributed/test_context_parallel.py
-    - pytest -v -s tests/distributed/test_nccl_symm_mem_allreduce.py
-    - pytest -v -s tests/v1/distributed/test_dbo.py
-
- label: 2 Node Test (4 GPUs)
-  timeout_in_minutes: 30
-  working_dir: "/vllm-workspace/tests"
-  num_gpus: 2
-  num_nodes: 2
-  source_file_dependencies:
-  - vllm/distributed/
-  - vllm/engine/
-  - vllm/executor/
-  - vllm/model_executor/models/
-  - tests/distributed/
-  - tests/examples/offline_inference/data_parallel.py
-  commands:
-    - ./.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-size=2 --tp-size=1 --node-size=2 --node-rank=0 --master-addr=192.168.10.10 --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-size=2 --tp-size=1 --node-size=2 --node-rank=1 --master-addr=192.168.10.10 --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_gpus: 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
-    - bash v1/kv_connector/nixl_integration/tp_config_sweep_accuracy_test.sh
-
- label: Pipeline + Context Parallelism (4 GPUs))
-  timeout_in_minutes: 60
-  working_dir: "/vllm-workspace/tests"
-  num_gpus: 4
-  source_file_dependencies:
-  - vllm/distributed/
-  - vllm/engine/
-  - vllm/executor/
-  - vllm/model_executor/models/
-  - tests/distributed/
-  commands:
-  - pytest -v -s distributed/test_pp_cudagraph.py
-  - pytest -v -s distributed/test_pipeline_parallel.py
@@ -1,59 +0,0 @@
-group: E2E Integration
-depends_on: 
-  - image-build
-steps:
- label: DeepSeek V2-Lite Accuracy
-  timeout_in_minutes: 60
-  gpu: h100
-  optional: true
-  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
-  gpu: h100
-  optional: true
-  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
-  gpu: b200
-  optional: true
-  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
-
- label: Prime-RL Integration (2 GPUs)
-  timeout_in_minutes: 30
-  optional: true
-  num_gpus: 2
-  working_dir: "/vllm-workspace"
-  source_file_dependencies:
-  - vllm/
-  - .buildkite/scripts/run-prime-rl-test.sh
-  commands:
-    - bash .buildkite/scripts/run-prime-rl-test.sh
-
- label: DeepSeek V2-Lite Async EPLB Accuracy
-  timeout_in_minutes: 60
-  gpu: h100
-  optional: true
-  num_gpus: 4
-  working_dir: "/vllm-workspace"
-  commands:
-  - bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_async_eplb.sh 0.25 1319 8030
-
- label: Qwen3-Next-80B-A3B-Instruct MTP Async EPLB Accuracy
-  timeout_in_minutes: 60
-  gpu: h100
-  optional: true
-  num_gpus: 4
-  working_dir: "/vllm-workspace"
-  commands:
-  - bash .buildkite/scripts/scheduled_integration_test/qwen3_next_mtp_async_eplb.sh 0.8 1319 8040
@@ -1,26 +0,0 @@
-group: Engine
-depends_on: 
-  - image-build
-steps:
- label: Engine
-  timeout_in_minutes: 15
-  source_file_dependencies:
-  - vllm/
-  - tests/engine
-  - tests/test_sequence
-  - tests/test_config
-  - tests/test_logger
-  - tests/test_vllm_port
-  commands:
-  - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py
-
- label: V1 e2e + engine
-  timeout_in_minutes: 45
-  source_file_dependencies:
-    - vllm/
-    - tests/v1
-  commands:
-    # TODO: accuracy does not match, whether setting
-    # VLLM_USE_FLASHINFER_SAMPLER or not on H100.
-    - pytest -v -s v1/e2e
-    - pytest -v -s v1/engine
@@ -1,68 +0,0 @@
-group: Entrypoints
-depends_on: 
-  - image-build
-steps:
- label: Entrypoints Unit Tests  
-  timeout_in_minutes: 10
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/entrypoints
-  - tests/entrypoints/
-  commands:
-  - pytest -v -s entrypoints/openai/tool_parsers
-  - pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/openai --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py  --ignore=entrypoints/pooling
-
- label: Entrypoints Integration (LLM)
-  timeout_in_minutes: 40
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/
-  - tests/entrypoints/llm
-  - tests/entrypoints/offline_mode
-  commands:
-  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - 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
-
- label: Entrypoints Integration (API Server)
-  timeout_in_minutes: 130
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/
-  - tests/entrypoints/openai
-  - tests/entrypoints/test_chat_utils
-  commands:
-  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/openai/test_collective_rpc.py # PYTHONPATH is needed to import custom Worker extension
-  - 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/test_collective_rpc.py --ignore=entrypoints/openai/tool_parsers/
-  - pytest -v -s entrypoints/test_chat_utils.py
-
-
- label: Entrypoints Integration (Pooling)
-  timeout_in_minutes: 50
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/
-  - tests/entrypoints/pooling
-  commands:
-  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/pooling
-
-
- label: Entrypoints V1
-  timeout_in_minutes: 50
-  source_file_dependencies:
-    - vllm/
-    - tests/v1
-  commands:
-    - pytest -v -s v1/entrypoints
-
- label: OpenAI API Correctness
-  timeout_in_minutes: 30
-  source_file_dependencies:
-  - csrc/
-  - vllm/entrypoints/openai/
-  - vllm/model_executor/models/whisper.py
-  commands: # LMEval+Transcription WER check
-  - pytest -s entrypoints/openai/correctness/
@@ -1,23 +0,0 @@
-group: Expert Parallelism
-depends_on: 
-  - image-build
-steps:
- label: EPLB Algorithm
-  timeout_in_minutes: 15
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/distributed/eplb
-  - tests/distributed/test_eplb_algo.py
-  commands:
-  - pytest -v -s distributed/test_eplb_algo.py
-
- label: EPLB Execution
-  timeout_in_minutes: 20
-  working_dir: "/vllm-workspace/tests"
-  num_gpus: 4
-  source_file_dependencies:
-  - vllm/distributed/eplb
-  - tests/distributed/test_eplb_execute.py
-  commands:
-  - pytest -v -s distributed/test_eplb_execute.py
-  - pytest -v -s distributed/test_eplb_spec_decode.py
@@ -1,117 +0,0 @@
-group: Kernels
-depends_on: 
-  - image-build
-steps:
- label: Kernels Core Operation Test
-  timeout_in_minutes: 75
-  source_file_dependencies:
-  - csrc/
-  - tests/kernels/core
-  - tests/kernels/test_top_k_per_row.py
-  commands:
-    - pytest -v -s kernels/core kernels/test_top_k_per_row.py
-
- label: Kernels Attention Test %N
-  timeout_in_minutes: 35
-  source_file_dependencies:
-  - csrc/attention/
-  - vllm/attention
-  - vllm/v1/attention
-  - tests/kernels/attention
-  commands:
-    - pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
-  parallelism: 2
-
- label: Kernels Quantization Test %N
-  timeout_in_minutes: 90
-  source_file_dependencies:
-  - csrc/quantization/
-  - vllm/model_executor/layers/quantization
-  - tests/kernels/quantization
-  commands:
-    - pytest -v -s kernels/quantization --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
-  parallelism: 2
-
- label: Kernels MoE Test %N
-  timeout_in_minutes: 60
-  source_file_dependencies:
-  - csrc/quantization/cutlass_w8a8/moe/
-  - csrc/moe/
-  - tests/kernels/moe
-  - vllm/model_executor/layers/fused_moe/
-  - vllm/distributed/device_communicators/
-  - vllm/envs.py
-  - vllm/config
-  commands:
-    - pytest -v -s kernels/moe --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
-  parallelism: 2
-
- label: Kernels Mamba Test
-  timeout_in_minutes: 45
-  source_file_dependencies:
-  - csrc/mamba/
-  - tests/kernels/mamba
-  - vllm/model_executor/layers/mamba/ops
-  commands:
-    - pytest -v -s kernels/mamba
-
- label: Kernels DeepGEMM Test (H100)
-  timeout_in_minutes: 45
-  gpu: h100
-  num_gpus: 1
-  source_file_dependencies:
-  - tools/install_deepgemm.sh
-  - vllm/utils/deep_gemm.py
-  - vllm/model_executor/layers/fused_moe
-  - vllm/model_executor/layers/quantization
-  - tests/kernels/quantization/test_block_fp8.py
-  - tests/kernels/moe/test_deepgemm.py
-  - tests/kernels/moe/test_batched_deepgemm.py
-  - tests/kernels/attention/test_deepgemm_attention.py
-  commands:
-    - pytest -v -s kernels/quantization/test_block_fp8.py -k deep_gemm
-    - pytest -v -s kernels/moe/test_deepgemm.py
-    - pytest -v -s kernels/moe/test_batched_deepgemm.py
-    - pytest -v -s kernels/attention/test_deepgemm_attention.py
-
- label: Kernels (B200)
-  timeout_in_minutes: 30
-  working_dir: "/vllm-workspace/"
-  gpu: b200
-  # optional: true
-  source_file_dependencies:
-  - csrc/quantization/fp4/
-  - csrc/attention/mla/
-  - 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_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
-  - vllm/v1/attention/backends/mla/flashinfer_mla.py
-  - vllm/platforms/cuda.py
-  - vllm/attention/selector.py
-  commands:
-    - nvidia-smi
-    - python3 examples/offline_inference/basic/chat.py
-    # Attention
-    # num_heads2 broken by https://github.com/flashinfer-ai/flashinfer/issues/1353
-    - pytest -v -s tests/kernels/attention/test_attention_selector.py
-    - pytest -v -s tests/kernels/attention/test_flashinfer.py -k 'not num_heads2'
-    - pytest -v -s tests/kernels/attention/test_flashinfer_trtllm_attention.py
-    - pytest -v -s tests/kernels/attention/test_cutlass_mla_decode.py
-    - pytest -v -s tests/kernels/attention/test_flashinfer_mla_decode.py
-    # Quantization
-    - pytest -v -s tests/kernels/quantization/test_cutlass_scaled_mm.py -k 'fp8'
-    - pytest -v -s tests/kernels/quantization/test_nvfp4_quant.py
-    - pytest -v -s tests/kernels/quantization/test_silu_mul_nvfp4_quant.py
-    - pytest -v -s tests/kernels/quantization/test_nvfp4_scaled_mm.py
-    - pytest -v -s tests/kernels/quantization/test_flashinfer_scaled_mm.py
-    - pytest -v -s tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py
-    - pytest -v -s tests/kernels/quantization/test_nvfp4_qutlass.py
-    - pytest -v -s tests/kernels/quantization/test_mxfp4_qutlass.py
-    - 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
@@ -1,46 +0,0 @@
-group: LM Eval
-depends_on: 
-  - image-build
-steps:
- label: LM Eval Small Models
-  timeout_in_minutes: 75
-  source_file_dependencies:
-  - csrc/
-  - vllm/model_executor/layers/quantization
-  autorun_on_main: true
-  commands:
-  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.txt --tp-size=1
-
- label: LM Eval Large Models (4 GPUs)(A100)
-  gpu: a100
-  optional: true
-  num_gpus: 4
-  working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
-  source_file_dependencies:
-  - csrc/
-  - vllm/model_executor/layers/quantization
-  commands:
-  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large.txt --tp-size=4
-
- label: LM Eval Large Models (4 GPUs)(H100)
-  gpu: h100
-  optional: true
-  num_gpus: 4
-  working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
-  source_file_dependencies:
-  - csrc/
-  - vllm/model_executor/layers/quantization
-  commands:
-    - export VLLM_USE_DEEP_GEMM=0  # We found Triton is faster than DeepGEMM for H100
-    - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-hopper.txt --tp-size=4
-
- label: LM Eval Small Models (B200)
-  timeout_in_minutes: 120
-  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 --tp-size=1
@@ -1,31 +0,0 @@
-group: LoRA
-depends_on: 
-  - image-build
-steps:
- label: LoRA %N
-  timeout_in_minutes: 30
-  source_file_dependencies:
-  - vllm/lora
-  - tests/lora
-  commands:
-    - pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_llm_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py
-  parallelism: 4
-
-
- label: LoRA TP (Distributed)
-  timeout_in_minutes: 30
-  num_gpus: 4
-  source_file_dependencies:
-  - vllm/lora
-  - tests/lora
-  commands:
-    # FIXIT: find out which code initialize cuda before running the test
-    # before the fix, we need to use spawn to test it
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    # There is some Tensor Parallelism related processing logic in LoRA that
-    # requires multi-GPU testing for validation.
-    - pytest -v -s -x lora/test_chatglm3_tp.py
-    - pytest -v -s -x lora/test_llama_tp.py
-    - pytest -v -s -x lora/test_llm_with_multi_loras.py
-    - pytest -v -s -x lora/test_olmoe_tp.py
-    - pytest -v -s -x lora/test_gptoss_tp.py
@@ -1,165 +0,0 @@
-group: Miscellaneous
-depends_on: 
-  - image-build
-steps:
- label: V1 Others
-  timeout_in_minutes: 60
-  source_file_dependencies:
-    - vllm/
-    - tests/v1
-  commands:
-    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
-    # split the test to avoid interference
-    - pytest -v -s -m 'not cpu_test' v1/core
-    - pytest -v -s v1/executor
-    - pytest -v -s v1/kv_offload
-    - pytest -v -s v1/sample
-    - pytest -v -s v1/logits_processors
-    - pytest -v -s v1/worker
-    - 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
-    - pytest -v -s v1/test_request.py
-    - pytest -v -s v1/test_outputs.py
-    # Integration test for streaming correctness (requires special branch).
-    - pip install -U git+https://github.com/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api
-    - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
-
- label: V1 Others (CPU)
-  depends_on: ~
-  source_file_dependencies:
-    - vllm/
-    - tests/v1
-  no_gpu: true
-  commands:
-    # split the test to avoid interference
-    - pytest -v -s -m 'cpu_test' v1/core
-    - pytest -v -s v1/structured_output
-    - pytest -v -s v1/test_serial_utils.py
-    - pytest -v -s -m 'cpu_test' v1/kv_connector/unit
-    - pytest -v -s -m 'cpu_test' v1/metrics
-
- label: Regression
-  timeout_in_minutes: 20
-  source_file_dependencies:
-  - vllm/
-  - tests/test_regression
-  commands:
-  - pip install modelscope
-  - pytest -v -s test_regression.py
-  working_dir: "/vllm-workspace/tests" # optional
-
- label: Examples
-  timeout_in_minutes: 45
-  working_dir: "/vllm-workspace/examples"
-  source_file_dependencies:
-  - vllm/entrypoints
-  - vllm/multimodal
-  - examples/
-  commands:
-    - pip install tensorizer # for tensorizer test
-    - python3 offline_inference/basic/chat.py # for basic
-    - python3 offline_inference/basic/generate.py --model facebook/opt-125m
-    - python3 offline_inference/basic/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10
-    - python3 offline_inference/basic/classify.py
-    - python3 offline_inference/basic/embed.py
-    - python3 offline_inference/basic/score.py
-    # for multi-modal models
-    - python3 offline_inference/audio_language.py --seed 0
-    - python3 offline_inference/vision_language.py --seed 0
-    - 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/pooling/vision_language_pooling.py --seed 0
-    # for features demo
-    - python3 offline_inference/prefix_caching.py
-    - python3 offline_inference/llm_engine_example.py
-    - python3 others/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 others/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
-    - python3 offline_inference/spec_decode.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
-    # https://github.com/vllm-project/vllm/pull/26682 uses slightly more memory in PyTorch 2.9+ causing this test to OOM in 1xL4 GPU
-    - python3 offline_inference/spec_decode.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
-
- label: Metrics, Tracing (2 GPUs)
-  timeout_in_minutes: 20
-  num_gpus: 2
-  source_file_dependencies:
-  - vllm/
-  - tests/v1/tracing
-  commands:
-  - "pip install \
-      'opentelemetry-sdk>=1.26.0' \
-      'opentelemetry-api>=1.26.0' \
-      'opentelemetry-exporter-otlp>=1.26.0' \
-      'opentelemetry-semantic-conventions-ai>=0.4.1'"
-  - pytest -v -s v1/tracing
-
- label: Python-only Installation
-  depends_on: ~
-  timeout_in_minutes: 20
-  source_file_dependencies:
-  - tests/standalone_tests/python_only_compile.sh
-  - setup.py
-  commands:
-  - bash standalone_tests/python_only_compile.sh
-
- label: Async Engine, Inputs, Utils, Worker
-  timeout_in_minutes: 50
-  source_file_dependencies:
-  - vllm/
-  - tests/multimodal
-  - tests/utils_
-  commands:
-  - pytest -v -s -m 'not cpu_test' multimodal
-  - pytest -v -s utils_
-
- label: Async Engine, Inputs, Utils, Worker, Config (CPU)
-  depends_on: ~
-  timeout_in_minutes: 30
-  source_file_dependencies:
-  - vllm/
-  - tests/test_inputs.py
-  - tests/test_outputs.py
-  - tests/multimodal
-  - tests/standalone_tests/lazy_imports.py
-  - tests/tokenizers_
-  - tests/tool_parsers
-  - tests/transformers_utils
-  - tests/config
-  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 -m 'cpu_test' multimodal
-  - pytest -v -s tokenizers_
-  - pytest -v -s tool_parsers
-  - pytest -v -s transformers_utils
-  - pytest -v -s config
-
- label: GPT-OSS Eval (B200)
-  timeout_in_minutes: 60
-  working_dir: "/vllm-workspace/"
-  gpu: b200
-  optional: true
-  source_file_dependencies:
-  - tests/evals/gpt_oss
-  - vllm/model_executor/models/gpt_oss.py
-  - vllm/model_executor/layers/quantization/mxfp4.py
-  - vllm/v1/attention/backends/flashinfer.py
-  commands:
-    - uv pip install --system 'gpt-oss[eval]==0.0.5'
-    - pytest -s -v tests/evals/gpt_oss/test_gpqa_correctness.py --model openai/gpt-oss-20b --metric 0.58
-
- label: Batch Invariance (H100)
-  timeout_in_minutes: 25
-  gpu: h100
-  source_file_dependencies:
-    - vllm/v1/attention
-    - vllm/model_executor/layers
-    - tests/v1/determinism/
-  commands:
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - pip install pytest-timeout pytest-forked
-    - pytest -v -s v1/determinism/test_batch_invariance.py
-    - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
@@ -1,17 +0,0 @@
-group: Model Executor
-depends_on: 
-  - image-build
-steps:
- label: Model Executor
-  timeout_in_minutes: 35
-  source_file_dependencies:
-  - vllm/engine/arg_utils.py
-  - vllm/config/model.py
-  - vllm/model_executor
-  - tests/model_executor
-  - tests/entrypoints/openai/test_tensorizer_entrypoint.py
-  commands:
-    - apt-get update && apt-get install -y curl libsodium23
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - pytest -v -s model_executor
-    - pytest -v -s entrypoints/openai/test_tensorizer_entrypoint.py
@@ -1,62 +0,0 @@
-group: Models - Basic
-depends_on: 
-  - image-build
-steps:
- label: Basic Models Tests (Initialization)
-  timeout_in_minutes: 45
-  mirror_hardwares: [amdexperimental]
-  torch_nightly: true
-  source_file_dependencies:
-  - vllm/
-  - tests/models/test_initialization.py
-  commands:
-    # Run a subset of model initialization tests
-    - pytest -v -s models/test_initialization.py::test_can_initialize_small_subset
-
- label: Basic Models Tests (Extra Initialization) %N
-  timeout_in_minutes: 45
-  mirror_hardwares: [amdexperimental]
-  torch_nightly: true
-  source_file_dependencies:
-  - vllm/model_executor/models/
-  - tests/models/test_initialization.py
-  commands:
-    # Only when vLLM model source is modified - test initialization of a large
-    # subset of supported models (the complement of the small subset in the above
-    # test.) Also run if model initialization test file is modified
-    - pytest -v -s models/test_initialization.py -k 'not test_can_initialize_small_subset' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
-  parallelism: 2
-
- label: Basic Models Tests (Other)
-  timeout_in_minutes: 45
-  source_file_dependencies:
-  - vllm/
-  - tests/models/test_transformers.py
-  - tests/models/test_registry.py
-  commands:
-    - pytest -v -s models/test_transformers.py models/test_registry.py
-
- label: Basic Models Test (Other CPU) # 5min
-  timeout_in_minutes: 10
-  source_file_dependencies:
-  - vllm/
-  - tests/models/test_utils.py
-  - tests/models/test_vision.py
-  no_gpu: true
-  commands:
-    - pytest -v -s models/test_utils.py models/test_vision.py
-
- label: Transformers Nightly Models
-  working_dir: "/vllm-workspace/"
-  optional: true
-  soft_fail: true
-  commands:
-    - pip install --upgrade git+https://github.com/huggingface/transformers
-    - pytest -v -s tests/models/test_initialization.py
-    - pytest -v -s tests/models/test_transformers.py
-    - pytest -v -s tests/models/multimodal/processing/
-    - pytest -v -s tests/models/multimodal/test_mapping.py
-    - python3 examples/offline_inference/basic/chat.py
-    - python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
-    # Whisper needs spawn method to avoid deadlock
-    - VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper
@@ -1,22 +0,0 @@
-group: Models - Distributed
-depends_on: 
-  - image-build
-steps:
- label: Distributed Model Tests (2 GPUs)
-  timeout_in_minutes: 50
-  working_dir: "/vllm-workspace/tests"
-  num_gpus: 2
-  source_file_dependencies:
-  - vllm/model_executor/model_loader/sharded_state_loader.py
-  - vllm/model_executor/models/
-  - tests/basic_correctness/
-  - tests/model_executor/model_loader/test_sharded_state_loader.py
-  - tests/models/
-  commands:
-  - TARGET_TEST_SUITE=L4 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
-  - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py
-  # Avoid importing model tests that cause CUDA reinitialization error
-  - pytest models/test_transformers.py -v -s -m 'distributed(num_gpus=2)'
-  - pytest models/language -v -s -m 'distributed(num_gpus=2)'
-  - pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py
-  - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest models/multimodal/generation/test_whisper.py -v -s -m 'distributed(num_gpus=2)'
@@ -1,91 +0,0 @@
-group: Models - Language
-depends_on: 
-  - image-build
-steps:
- label: Language Models Tests (Standard)
-  timeout_in_minutes: 25
-  mirror_hardwares: [amdexperimental]
-  torch_nightly: true
-  source_file_dependencies:
-  - vllm/
-  - tests/models/language
-  commands:
-    # Test standard language models, excluding a subset of slow tests
-    - pip freeze | grep -E 'torch'
-    - pytest -v -s models/language -m 'core_model and (not slow_test)'
-
- label: Language Models Tests (Extra Standard) %N
-  timeout_in_minutes: 45
-  mirror_hardwares: [amdexperimental]
-  torch_nightly: true
-  source_file_dependencies:
-  - vllm/model_executor/models/
-  - tests/models/language/pooling/test_embedding.py
-  - tests/models/language/generation/test_common.py
-  - tests/models/language/pooling/test_classification.py
-  commands:
-    # Shard slow subset of standard language models tests. Only run when model
-    # source is modified, or when specified test files are modified
-    - pip freeze | grep -E 'torch'
-    - pytest -v -s models/language -m 'core_model and slow_test' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
-  parallelism: 2
-
- label: Language Models Tests (Hybrid) %N
-  timeout_in_minutes: 75
-  mirror_hardwares: [amdexperimental]
-  torch_nightly: true
-  source_file_dependencies:
-  - vllm/
-  - tests/models/language/generation
-  commands:
-    # Install fast path packages for testing against transformers
-    # Note: also needed to run plamo2 model in vLLM
-    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.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
-  parallelism: 2
-
- label: Language Models Test (Extended Generation) # 80min
-  timeout_in_minutes: 110
-  mirror_hardwares: [amdexperimental]
-  optional: true
-  source_file_dependencies:
-  - vllm/
-  - tests/models/language/generation
-  commands:
-    # Install fast path packages for testing against transformers
-    # Note: also needed to run plamo2 model in vLLM
-    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.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/
-  - tests/models/language/generation_ppl_test
-  commands:
-    - pytest -v -s models/language/generation_ppl_test
-
- label: Language Models Test (Extended Pooling)  # 36min
-  timeout_in_minutes: 50
-  mirror_hardwares: [amdexperimental]
-  optional: true
-  source_file_dependencies:
-  - vllm/
-  - tests/models/language/pooling
-  commands:
-    - pytest -v -s models/language/pooling -m 'not core_model'
-
- label: Language Models Test (MTEB)
-  timeout_in_minutes: 110
-  mirror_hardwares: [amdexperimental]
-  optional: true
-  source_file_dependencies:
-  - vllm/
-  - tests/models/language/pooling_mteb_test
-  commands:
-    - pytest -v -s models/language/pooling_mteb_test
@@ -1,79 +0,0 @@
-group: Models - Multimodal
-depends_on: 
-  - image-build
-steps:
- label: Multi-Modal Models (Standard) # 60min
-  timeout_in_minutes: 80
-  source_file_dependencies:
-  - vllm/
-  - tests/models/multimodal
-  commands:
-    - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
-    - pip freeze | grep -E 'torch'
-    - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/processing
-    - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model  # Otherwise, mp_method="spawn" doesn't work
-
- label: Multi-Modal Processor Test (CPU)
-  timeout_in_minutes: 60
-  source_file_dependencies:
-  - vllm/
-  - tests/models/multimodal
-  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
-
- label: Multi-Modal Processor # 44min
-  timeout_in_minutes: 60
-  source_file_dependencies:
-  - vllm/
-  - tests/models/multimodal
-  commands:
-    - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
-    - pytest -v -s models/multimodal/processing/test_tensor_schema.py
-
- label: Multi-Modal Accuracy Eval (Small Models) # 50min
-  timeout_in_minutes: 70
-  working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
-  source_file_dependencies:
-  - vllm/multimodal/
-  - vllm/inputs/
-  - vllm/v1/core/
-  commands:
-  - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-mm-small.txt --tp-size=1
-
- label: Multi-Modal Models (Extended) 1
-  optional: true
-  source_file_dependencies:
-  - vllm/
-  - tests/models/multimodal
-  commands:
-    - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
-    - pytest -v -s models/multimodal -m 'not core_model' --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/processing
-
- label: Multi-Modal Models (Extended) 2
-  optional: true
-  source_file_dependencies:
-  - vllm/
-  - tests/models/multimodal
-  commands:
-    - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
-    - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=0) and not core_model'
-
- label: Multi-Modal Models (Extended) 3
-  optional: true
-  source_file_dependencies:
-  - vllm/
-  - tests/models/multimodal
-  commands:
-    - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
-    - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=1) and not core_model'
-
-# This test is used only in PR development phase to test individual models and should never run on main
- label: Custom Models
-  optional: true
-  commands:
-    - echo 'Testing custom models...'
-    # PR authors can temporarily add commands below to test individual models
-    # e.g. pytest -v -s models/encoder_decoder/vision_language/test_mllama.py
-    # *To avoid merge conflicts, remember to REMOVE (not just comment out) them before merging the PR*
@@ -1,34 +0,0 @@
-group: Plugins
-depends_on: 
-  - image-build
-steps:
- label: Plugin Tests (2 GPUs)
-  timeout_in_minutes: 60
-  working_dir: "/vllm-workspace/tests"
-  num_gpus: 2
-  source_file_dependencies:
-  - vllm/plugins/
-  - tests/plugins/
-  commands:
-  # begin platform plugin and general plugin tests, all the code in-between runs on dummy platform
-  - pip install -e ./plugins/vllm_add_dummy_platform
-  - pytest -v -s plugins_tests/test_platform_plugins.py
-  - pip uninstall vllm_add_dummy_platform -y
-  # end platform plugin tests
-  # begin io_processor plugins test, all the code in between uses the prithvi_io_processor plugin
-  - pip install -e ./plugins/prithvi_io_processor_plugin
-  - pytest -v -s plugins_tests/test_io_processor_plugins.py
-  - pip uninstall prithvi_io_processor_plugin -y
-  # end io_processor plugins test
-  # begin stat_logger plugins test
-  - pip install -e ./plugins/vllm_add_dummy_stat_logger
-  - pytest -v -s plugins_tests/test_stats_logger_plugins.py
-  - pip uninstall dummy_stat_logger -y
-  # end stat_logger plugins test
-  # other tests continue here:
-  - pytest -v -s plugins_tests/test_scheduler_plugins.py
-  - pip install -e ./plugins/vllm_add_dummy_model
-  - 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 in-tree lora resolver plugins
@@ -1,50 +0,0 @@
-group: PyTorch
-depends_on: 
-  - image-build
-steps:
- label: PyTorch Compilation Unit Tests
-  timeout_in_minutes: 30
-  source_file_dependencies:
-    - vllm/
-    - tests/compile
-  commands:
-  # Run unit tests defined directly under compile/,
-  # not including subdirectories, which are usually heavier
-  # tests covered elsewhere.
-  # Use `find` to launch multiple instances of pytest so that
-  # they do not suffer from https://github.com/vllm-project/vllm/issues/28965
-  - "find compile/ -maxdepth 1 -name 'test_*.py' -exec pytest -s -v {} \\;"
-
- label: PyTorch Fullgraph Smoke Test
-  timeout_in_minutes: 30
-  source_file_dependencies:
-  - vllm/
-  - tests/compile
-  commands:
-  # Run smoke tests under fullgraph directory, except test_full_graph.py
-  # as it is a heavy test that is covered in other steps.
-  # Use `find` to launch multiple instances of pytest so that
-  # they do not suffer from https://github.com/vllm-project/vllm/issues/28965
-  - "find compile/fullgraph/ -name 'test_*.py' -not -name 'test_full_graph.py' -exec pytest -s -v {} \\;"
-
- label: PyTorch Fullgraph
-  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
-  # of the dependencies. Please check the error message and add the package to whitelist
-  # in /vllm/tools/pre_commit/generate_nightly_torch_test.py
-  soft_fail: true
-  source_file_dependencies:
-  - requirements/nightly_torch_test.txt
-  commands:
-  - bash standalone_tests/pytorch_nightly_dependency.sh
@@ -1,46 +0,0 @@
-group: Quantization
-depends_on: 
-  - image-build
-steps:
- label: Quantization
-  timeout_in_minutes: 90
-  source_file_dependencies:
-  - csrc/
-  - vllm/model_executor/layers/quantization
-  - tests/quantization
-  commands:
-  # temporary install here since we need nightly, will move to requirements/test.in
-  # after torchao 0.12 release, and pin a working version of torchao nightly here
-
-  # since torchao nightly is only compatible with torch nightly currently
-  # 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.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/"
-  gpu: b200
-  source_file_dependencies:
-  - tests/quantization/test_blackwell_moe.py
-  - vllm/model_executor/models/deepseek_v2.py
-  - vllm/model_executor/models/gpt_oss.py
-  - vllm/model_executor/models/llama4.py
-  - vllm/model_executor/layers/fused_moe
-  - vllm/model_executor/layers/quantization/compressed_tensors
-  - vllm/model_executor/layers/quantization/modelopt.py
-  - vllm/model_executor/layers/quantization/mxfp4.py
-  - vllm/v1/attention/backends/flashinfer.py
-  commands:
-    - pytest -s -v tests/quantization/test_blackwell_moe.py
-
- label: Quantized Models Test
-  timeout_in_minutes: 60
-  source_file_dependencies:
-  - vllm/model_executor/layers/quantization
-  - tests/models/quantization
-  commands:
-    - pytest -v -s models/quantization
@@ -1,14 +0,0 @@
-group: Samplers
-depends_on: 
-  - image-build
-steps:
- label: Samplers Test
-  timeout_in_minutes: 75
-  source_file_dependencies:
-  - vllm/model_executor/layers
-  - vllm/sampling_metadata.py
-  - tests/samplers
-  - tests/conftest.py
-  commands:
-    - pytest -v -s samplers
-    - VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers
@@ -1,13 +0,0 @@
-group: Tool use
-depends_on: 
-  - image-build
-steps:
- label: OpenAI-Compatible Tool Use
-  timeout_in_minutes: 35
-  mirror_hardwares: [amdexperimental]
-  fast_check: false
-  source_file_dependencies:
-    - vllm/
-    - tests/tool_use
-  commands:
-    - pytest -v -s tool_use
@@ -1,25 +0,0 @@
-group: Weight Loading
-depends_on: 
-  - image-build
-steps:
- label: Weight Loading Multiple GPU  # 33min
-  timeout_in_minutes: 45
-  working_dir: "/vllm-workspace/tests"
-  num_gpus: 2
-  optional: true
-  source_file_dependencies:
-  - vllm/
-  - tests/weight_loading
-  commands:
-    - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models.txt
-
- label: Weight Loading Multiple GPU - Large Models # optional
-  working_dir: "/vllm-workspace/tests"
-  num_gpus: 2
-  gpu: a100
-  optional: true
-  source_file_dependencies:
-  - vllm/
-  - tests/weight_loading
-  commands:
-    - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-large.txt
@@ -9,7 +9,6 @@
 /vllm/model_executor/layers/quantization @mgoin @robertgshaw2-redhat @tlrmchlsmth @yewentao256 @pavanimajety
 /vllm/model_executor/layers/mamba @tdoublep
 /vllm/model_executor/model_loader @22quinn
-/vllm/model_executor/layers/batch_invariant.py @yewentao256 
 /vllm/multimodal @DarkLight1337 @ywang96 @NickLucche @tjtanaa
 /vllm/vllm_flash_attn @LucasWilkinson
 /vllm/lora @jeejeelee
@@ -36,9 +35,6 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
 /vllm/v1/kv_cache_interface.py @heheda12345
 /vllm/v1/offloading @ApostaC

-# Model runner V2
-/vllm/v1/worker/gpu @WoosukKwon
-
 # Test ownership
 /.buildkite/lm-eval-harness @mgoin 
 /tests/distributed/test_multi_node_assignment.py @youkaichao
@@ -60,7 +56,6 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
 /tests/v1/kv_connector/nixl_integration @NickLucche
 /tests/v1/kv_connector @ApostaC
 /tests/v1/offloading @ApostaC
-/tests/v1/determinism @yewentao256 

 # Transformers modeling backend
 /vllm/model_executor/models/transformers @hmellor
@@ -146,10 +141,9 @@ mkdocs.yaml @hmellor
 /requirements/kv_connectors.txt @NickLucche

 # Pooling models
-/examples/pooling @noooop
+/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.py @noooop
@@ -14,52 +14,6 @@ pull_request_rules:
    comment:
      message: "Documentation preview: https://vllm--{{number}}.org.readthedocs.build/en/{{number}}/"

- name: comment-pre-commit-failure
-  description: Comment on PR when pre-commit check fails
-  conditions:
-    - status-failure=pre-commit
-    - -closed
-    - -draft
-  actions:
-    comment:
-      message: |
-        Hi @{{author}}, the pre-commit checks have failed. Please run:
-
-        ```bash 
-        uv pip install pre-commit
-        pre-commit install
-        pre-commit run --all-files
-        ```
-
-        Then, commit the changes and push to your branch.
-
-        For future commits, `pre-commit` will run automatically on changed files before each commit.
-
-        > [!TIP]
-        > <details>
-        > <summary>Is <code>mypy</code> or <code>markdownlint</code> failing?</summary>
-        > <br/>
-        > <code>mypy</code> and <code>markdownlint</code> are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:
-        >
-        > ```bash
-        > # For mypy (substitute "3.10" with the failing version if needed)
-        > pre-commit run --hook-stage manual mypy-3.10
-        > # For markdownlint
-        > pre-commit run --hook-stage manual markdownlint
-        > ```
-        > </details>
-
- name: comment-dco-failure
-  description: Comment on PR when DCO check fails
-  conditions:
-    - status-failure=dco
-    - -closed
-    - -draft
-  actions:
-    comment:
-      message: |
-        Hi @{{author}}, the DCO check has failed. Please click on DCO in the Checks section for instructions on how to resolve this.
-
 - name: label-ci-build
  description: Automatically apply ci/build label
  conditions:
@@ -186,7 +140,7 @@ pull_request_rules:
      - files~=^tests/entrypoints/test_context.py
      - files~=^vllm/model_executor/models/.*gpt[-_]?oss.*\.py
      - files~=^vllm/model_executor/layers/.*gpt[-_]?oss.*\.py
-      - files~=^vllm/entrypoints/openai/parser/harmony_utils.py
+      - files~=^vllm/entrypoints/harmony_utils.py
      - files~=^vllm/entrypoints/tool_server.py
      - files~=^vllm/entrypoints/tool.py
      - files~=^vllm/entrypoints/context.py
@@ -404,4 +358,4 @@ pull_request_rules:
  actions:
    label:
      add:
-        - kv-connector
+        - kv-connector
@@ -13,10 +13,10 @@ jobs:

    steps:
      - name: Checkout repository
-        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1
+        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2

      - name: Set up Python
-        uses: actions/setup-python@83679a892e2d95755f2dac6acb0bfd1e9ac5d548 # v6.1.0
+        uses: actions/setup-python@e797f83bcb11b83ae66e0230d6156d7c80228e7c # v6.0.0
        with:
          python-version: '3.12'

@@ -105,31 +105,6 @@ jobs:
                  }
                ],
              },
-              cpu: {
-                // Keyword search - matches whole words only (with word boundaries)
-                keywords: [
-                  {
-                    term: "CPU Backend",
-                    searchIn: "title"
-                  },
-                  {
-                    term: "x86",
-                    searchIn: "title"
-                  },
-                  {
-                    term: "ARM",
-                    searchIn: "title"
-                  },
-                  {
-                    term: "Apple Silicon",
-                    searchIn: "title"
-                  },
-                  {
-                    term: "IBM Z",
-                    searchIn: "title"
-                  },
-                ],
-              },
              // Add more label configurations here as needed
              // example: {
              //   keywords: [...],
@@ -12,7 +12,7 @@ jobs:
    timeout-minutes: 30

    steps:
-      - uses: actions/checkout@v6.0.1
+      - uses: actions/checkout@v4

      - uses: astral-sh/setup-uv@v7
        with:
@@ -16,8 +16,8 @@ jobs:
  pre-commit:
    runs-on: ubuntu-latest
    steps:
-    - uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1
-    - uses: actions/setup-python@83679a892e2d95755f2dac6acb0bfd1e9ac5d548 # v6.1.0
+    - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+    - uses: actions/setup-python@e797f83bcb11b83ae66e0230d6156d7c80228e7c # v6.0.0
      with:
        python-version: "3.12"
    - run: echo "::add-matcher::.github/workflows/matchers/actionlint.json"
@@ -7,15 +7,13 @@ on:

 jobs:
  close-issues-and-pull-requests:
-    # Prevents triggering on forks or other repos
-    if: github.repository == 'vllm-project/vllm'
    permissions:
      issues: write
      pull-requests: write
      actions: write
    runs-on: ubuntu-latest
    steps:
-      - uses: actions/stale@997185467fa4f803885201cee163a9f38240193d # v10.1.1
+      - uses: actions/stale@5f858e3efba33a5ca4407a664cc011ad407f2008 # v10.1.0
        with:
          # Increasing this value ensures that changes to this workflow
          # propagate to all issues and PRs in days rather than months
@@ -136,7 +136,7 @@ elseif(HIP_FOUND)

  # ROCm 5.X and 6.X
  if (ROCM_VERSION_DEV_MAJOR GREATER_EQUAL 5 AND
-      Torch_VERSION VERSION_LESS ${TORCH_SUPPORTED_VERSION_ROCM})
+      NOT Torch_VERSION VERSION_EQUAL ${TORCH_SUPPORTED_VERSION_ROCM})
    message(WARNING "Pytorch version >= ${TORCH_SUPPORTED_VERSION_ROCM} "
      "expected for ROCm build, saw ${Torch_VERSION} instead.")
  endif()
@@ -354,17 +354,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
  # Only build Marlin kernels if we are building for at least some compatible archs.
  # Keep building Marlin for 9.0 as there are some group sizes and shapes that
  # are not supported by Machete yet.
-
-  # marlin arches for fp16 output
-  cuda_archs_loose_intersection(MARLIN_ARCHS "8.0+PTX" "${CUDA_ARCHS}")
-  # marlin arches for bf16 output (we need 9.0 for bf16 atomicAdd PTX)
-  cuda_archs_loose_intersection(MARLIN_BF16_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}")
-  # marlin arches for fp8 input
-  # - sm80 doesn't support fp8 computation
-  # - sm90 and sm100 don't support QMMA.16832.F32.E4M3.E4M3 SAAS instruction
-  # so we only enable fp8 computation for SM89 (e.g. RTX 40x0)  and 12.0 (e.g. RTX 50x0)
-  cuda_archs_loose_intersection(MARLIN_FP8_ARCHS "8.9;12.0" "${CUDA_ARCHS}")
-
+  # 9.0 for latest bf16 atomicAdd PTX
+  cuda_archs_loose_intersection(MARLIN_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}")
  if (MARLIN_ARCHS)

    #
@@ -374,18 +365,16 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
    set(MARLIN_GEN_SCRIPT
      ${CMAKE_CURRENT_SOURCE_DIR}/csrc/quantization/gptq_marlin/generate_kernels.py)
    file(MD5 ${MARLIN_GEN_SCRIPT} MARLIN_GEN_SCRIPT_HASH)
-    list(JOIN CUDA_ARCHS "," CUDA_ARCHS_STR)
-    set(MARLIN_GEN_SCRIPT_HASH_AND_ARCH "${MARLIN_GEN_SCRIPT_HASH}(ARCH:${CUDA_ARCHS_STR})")

-    message(STATUS "Marlin generation script hash: ${MARLIN_GEN_SCRIPT_HASH_AND_ARCH}")
-    message(STATUS "Last run Marlin generate script hash: $CACHE{MARLIN_GEN_SCRIPT_HASH_AND_ARCH}")
+    message(STATUS "Marlin generation script hash: ${MARLIN_GEN_SCRIPT_HASH}")
+    message(STATUS "Last run Marlin generate script hash: $CACHE{MARLIN_GEN_SCRIPT_HASH}")

-    if (NOT DEFINED CACHE{MARLIN_GEN_SCRIPT_HASH_AND_ARCH}
-        OR NOT $CACHE{MARLIN_GEN_SCRIPT_HASH_AND_ARCH} STREQUAL ${MARLIN_GEN_SCRIPT_HASH_AND_ARCH})
+    if (NOT DEFINED CACHE{MARLIN_GEN_SCRIPT_HASH}
+        OR NOT $CACHE{MARLIN_GEN_SCRIPT_HASH} STREQUAL ${MARLIN_GEN_SCRIPT_HASH})
      execute_process(
        COMMAND ${CMAKE_COMMAND} -E env
-        PYTHONPATH=$ENV{PYTHONPATH}
-          ${Python_EXECUTABLE} ${MARLIN_GEN_SCRIPT} ${CUDA_ARCHS_STR}
+        PYTHONPATH=$PYTHONPATH
+          ${Python_EXECUTABLE} ${MARLIN_GEN_SCRIPT}
        RESULT_VARIABLE marlin_generation_result
        OUTPUT_VARIABLE marlin_generation_result
        OUTPUT_FILE ${CMAKE_CURRENT_BINARY_DIR}/marlin_generation.log
@@ -398,15 +387,15 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
                            "\nCheck the log for details: "
                            "${CMAKE_CURRENT_BINARY_DIR}/marlin_generation.log")
      else()
-        set(MARLIN_GEN_SCRIPT_HASH_AND_ARCH ${MARLIN_GEN_SCRIPT_HASH_AND_ARCH}
-            CACHE STRING "Last run Marlin generate script hash and arch" FORCE)
+        set(MARLIN_GEN_SCRIPT_HASH ${MARLIN_GEN_SCRIPT_HASH}
+            CACHE STRING "Last run Marlin generate script hash" FORCE)
        message(STATUS "Marlin generation completed successfully.")
      endif()
    else()
      message(STATUS "Marlin generation script has not changed, skipping generation.")
    endif()

-    file(GLOB MARLIN_TEMPLATE_KERNEL_SRC "csrc/quantization/gptq_marlin/sm80_kernel_*_float16.cu")
+    file(GLOB MARLIN_TEMPLATE_KERNEL_SRC "csrc/quantization/gptq_marlin/kernel_*.cu")
    set_gencode_flags_for_srcs(
      SRCS "${MARLIN_TEMPLATE_KERNEL_SRC}"
      CUDA_ARCHS "${MARLIN_ARCHS}")
@@ -414,34 +403,12 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
      set_source_files_properties(${MARLIN_TEMPLATE_KERNEL_SRC}
        PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
    endif()
+
    list(APPEND VLLM_EXT_SRC ${MARLIN_TEMPLATE_KERNEL_SRC})

-    file(GLOB MARLIN_TEMPLATE_BF16_KERNEL_SRC "csrc/quantization/gptq_marlin/sm80_kernel_*_bfloat16.cu")
-    set_gencode_flags_for_srcs(
-      SRCS "${MARLIN_TEMPLATE_BF16_KERNEL_SRC}"
-      CUDA_ARCHS "${MARLIN_BF16_ARCHS}")
-    if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
-      set_source_files_properties(${MARLIN_TEMPLATE_BF16_KERNEL_SRC}
-        PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
-    endif()
-    list(APPEND VLLM_EXT_SRC ${MARLIN_TEMPLATE_BF16_KERNEL_SRC})
-
-    if (MARLIN_FP8_ARCHS) 
-      file(GLOB MARLIN_TEMPLATE_FP8_KERNEL_SRC "csrc/quantization/gptq_marlin/sm89_kernel_*.cu")
-      set_gencode_flags_for_srcs(
-        SRCS "${MARLIN_TEMPLATE_FP8_KERNEL_SRC}"
-        CUDA_ARCHS "${MARLIN_FP8_ARCHS}")
-      if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
-        set_source_files_properties(${MARLIN_TEMPLATE_FP8_KERNEL_SRC}
-          PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
-      endif()
-      list(APPEND VLLM_EXT_SRC ${MARLIN_TEMPLATE_FP8_KERNEL_SRC})
-    endif()
-
    set(MARLIN_SRCS
       "csrc/quantization/marlin/sparse/marlin_24_cuda_kernel.cu"
       "csrc/quantization/gptq_marlin/gptq_marlin.cu"
-       "csrc/quantization/gptq_marlin/marlin_int4_fp8_preprocess.cu"
       "csrc/quantization/gptq_marlin/gptq_marlin_repack.cu"
       "csrc/quantization/gptq_marlin/awq_marlin_repack.cu")
    set_gencode_flags_for_srcs(
@@ -637,15 +604,12 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
    set(SRCS
      "csrc/quantization/fp4/nvfp4_quant_kernels.cu"
      "csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu"
-      "csrc/quantization/fp4/nvfp4_experts_quant.cu"
-      "csrc/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu"
-      "csrc/quantization/fp4/nvfp4_blockwise_moe_kernel.cu")
+      "csrc/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu")
    set_gencode_flags_for_srcs(
      SRCS "${SRCS}"
      CUDA_ARCHS "${FP4_ARCHS}")
    list(APPEND VLLM_EXT_SRC "${SRCS}")
    list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM120=1")
-    list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM120=1")
    message(STATUS "Building NVFP4 for archs: ${FP4_ARCHS}")
  else()
    message(STATUS "Not building NVFP4 as no compatible archs were found.")
@@ -822,7 +786,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
        OR NOT $CACHE{MACHETE_GEN_SCRIPT_HASH} STREQUAL ${MACHETE_GEN_SCRIPT_HASH})
      execute_process(
        COMMAND ${CMAKE_COMMAND} -E env
-        PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/csrc/cutlass_extensions/:${CUTLASS_DIR}/python/:${VLLM_PYTHON_PATH}:$ENV{PYTHONPATH}
+        PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/csrc/cutlass_extensions/:${CUTLASS_DIR}/python/:${VLLM_PYTHON_PATH}:$PYTHONPATH
          ${Python_EXECUTABLE} ${MACHETE_GEN_SCRIPT}
        RESULT_VARIABLE machete_generation_result
        OUTPUT_VARIABLE machete_generation_output
@@ -874,10 +838,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
  cuda_archs_loose_intersection(W4A8_ARCHS "9.0a" "${CUDA_ARCHS}")
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND W4A8_ARCHS)
    set(SRCS
-       "csrc/quantization/cutlass_w4a8/w4a8_mm_entry.cu"
-       "csrc/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu"
-       "csrc/quantization/cutlass_w4a8/w4a8_utils.cu"
-       )
+       "csrc/quantization/cutlass_w4a8/w4a8_mm_entry.cu")

    set_gencode_flags_for_srcs(
      SRCS "${SRCS}"
@@ -947,7 +908,9 @@ target_compile_definitions(_C PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
 set(VLLM_MOE_EXT_SRC
  "csrc/moe/torch_bindings.cpp"
  "csrc/moe/moe_align_sum_kernels.cu"
-  "csrc/moe/topk_softmax_kernels.cu")
+  "csrc/moe/moe_lora_align_sum_kernels.cu"
+  "csrc/moe/topk_softmax_kernels.cu"
+  "csrc/moe/moe_fused_gate.cu")

 if(VLLM_GPU_LANG STREQUAL "CUDA")
  list(APPEND VLLM_MOE_EXT_SRC
@@ -976,15 +939,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
    CUDA_ARCHS "${CUDA_ARCHS}")

  list(APPEND VLLM_MOE_EXT_SRC "${VLLM_MOE_WNA16_SRC}")
-  # moe marlin arches
-  # note that we always set `use_atomic_add=False` for moe marlin now,
-  # so we don't need 9.0 for bf16 atomicAdd PTX
-  cuda_archs_loose_intersection(MARLIN_MOE_ARCHS "8.0+PTX" "${CUDA_ARCHS}")
-  # moe marlin arches for fp8 input
-  # - sm80 doesn't support fp8 computation
-  # - sm90 and sm100 don't support QMMA.16832.F32.E4M3.E4M3 SAAS instruction
-  # so we only enable fp8 computation for SM89 (e.g. RTX 40x0)  and 12.0 (e.g. RTX 50x0)
-  cuda_archs_loose_intersection(MARLIN_MOE_FP8_ARCHS "8.9;12.0" "${CUDA_ARCHS}")
+  # 9.0 for latest bf16 atomicAdd PTX
+  cuda_archs_loose_intersection(MARLIN_MOE_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}")
  if (MARLIN_MOE_ARCHS)

    #
@@ -994,18 +950,16 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
    set(MOE_MARLIN_GEN_SCRIPT
      ${CMAKE_CURRENT_SOURCE_DIR}/csrc/moe/marlin_moe_wna16/generate_kernels.py)
    file(MD5 ${MOE_MARLIN_GEN_SCRIPT} MOE_MARLIN_GEN_SCRIPT_HASH)
-    list(JOIN CUDA_ARCHS "," CUDA_ARCHS_STR)
-    set(MOE_MARLIN_GEN_SCRIPT_HASH_AND_ARCH "${MOE_MARLIN_GEN_SCRIPT_HASH}(ARCH:${CUDA_ARCHS_STR})")

-    message(STATUS "Marlin MOE generation script hash with arch: ${MOE_MARLIN_GEN_SCRIPT_HASH_AND_ARCH}")
-    message(STATUS "Last run Marlin MOE generate script hash with arch: $CACHE{MOE_MARLIN_GEN_SCRIPT_HASH_AND_ARCH}")
+    message(STATUS "Marlin MOE generation script hash: ${MOE_MARLIN_GEN_SCRIPT_HASH}")
+    message(STATUS "Last run Marlin MOE generate script hash: $CACHE{MOE_MARLIN_GEN_SCRIPT_HASH}")

-    if (NOT DEFINED CACHE{MOE_MARLIN_GEN_SCRIPT_HASH_AND_ARCH}
-        OR NOT $CACHE{MOE_MARLIN_GEN_SCRIPT_HASH_AND_ARCH} STREQUAL ${MOE_MARLIN_GEN_SCRIPT_HASH_AND_ARCH})
+    if (NOT DEFINED CACHE{MOE_MARLIN_GEN_SCRIPT_HASH}
+        OR NOT $CACHE{MOE_MARLIN_GEN_SCRIPT_HASH} STREQUAL ${MOE_MARLIN_GEN_SCRIPT_HASH})
      execute_process(
        COMMAND ${CMAKE_COMMAND} -E env
-        PYTHONPATH=$ENV{PYTHONPATH}
-          ${Python_EXECUTABLE} ${MOE_MARLIN_GEN_SCRIPT} ${CUDA_ARCHS_STR}
+        PYTHONPATH=$PYTHONPATH
+          ${Python_EXECUTABLE} ${MOE_MARLIN_GEN_SCRIPT}
        RESULT_VARIABLE moe_marlin_generation_result
        OUTPUT_VARIABLE moe_marlin_generation_output
        OUTPUT_FILE ${CMAKE_CURRENT_BINARY_DIR}/moe_marlin_generation.log
@@ -1018,7 +972,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
                            "\nCheck the log for details: "
                            "${CMAKE_CURRENT_BINARY_DIR}/moe_marlin_generation.log")
      else()
-        set(MOE_MARLIN_GEN_SCRIPT_HASH_AND_ARCH ${MOE_MARLIN_GEN_SCRIPT_HASH_AND_ARCH}
+        set(MOE_MARLIN_GEN_SCRIPT_HASH ${MOE_MARLIN_GEN_SCRIPT_HASH}
            CACHE STRING "Last run Marlin MOE generate script hash" FORCE)
        message(STATUS "Marlin MOE generation completed successfully.")
      endif()
@@ -1026,28 +980,16 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
      message(STATUS "Marlin MOE generation script has not changed, skipping generation.")
    endif()

-    file(GLOB MARLIN_MOE_SRC "csrc/moe/marlin_moe_wna16/sm80_kernel_*.cu")
-    list(APPEND MARLIN_MOE_SRC "csrc/moe/marlin_moe_wna16/ops.cu")
+    file(GLOB MOE_WNAA16_MARLIN_SRC "csrc/moe/marlin_moe_wna16/*.cu")
    set_gencode_flags_for_srcs(
-      SRCS "${MARLIN_MOE_SRC}"
+      SRCS "${MOE_WNAA16_MARLIN_SRC}"
      CUDA_ARCHS "${MARLIN_MOE_ARCHS}")
    if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
-      set_source_files_properties(${MARLIN_MOE_SRC}
+      set_source_files_properties(${MOE_WNAA16_MARLIN_SRC}
        PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
    endif()
-    list(APPEND VLLM_MOE_EXT_SRC ${MARLIN_MOE_SRC})

-    if (MARLIN_MOE_FP8_ARCHS)
-      file(GLOB MARLIN_MOE_FP8_SRC "csrc/moe/marlin_moe_wna16/sm89_kernel_*.cu")
-      set_gencode_flags_for_srcs(
-        SRCS "${MARLIN_MOE_FP8_SRC}"
-        CUDA_ARCHS "${MARLIN_MOE_FP8_ARCHS}")
-      if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
-        set_source_files_properties(${MARLIN_MOE_FP8_SRC}
-          PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
-      endif()
-      list(APPEND VLLM_MOE_EXT_SRC ${MARLIN_MOE_FP8_SRC})
-    endif()
+    list(APPEND VLLM_MOE_EXT_SRC ${MOE_WNAA16_MARLIN_SRC})

    message(STATUS "Building Marlin MOE kernels for archs: ${MARLIN_MOE_ARCHS}")
  else()
@@ -21,7 +21,6 @@ Join us at the [PyTorch Conference, October 22-23](https://events.linuxfoundatio

 *Latest News* 🔥

- [2025/11] We hosted [vLLM Bangkok Meetup](https://luma.com/v0f647nv). We explored vLLM and LMCache inference and low-resource language adaptation with speakers from Embedded LLM, AMD, and Red Hat. Please find the meetup slides [here](https://drive.google.com/drive/folders/1H0DS57F8HQ5q3kSOSoRmucPJWL3E0A_X?usp=sharing).
 - [2025/11] We hosted [the first vLLM Europe Meetup in Zurich](https://luma.com/0gls27kb) focused on quantization, distributed inference, and reinforcement learning at scale with speakers from Mistral, IBM, and Red Hat. Please find the meetup slides [here](https://docs.google.com/presentation/d/1UC9PTLCHYXQpOmJDSFg6Sljra3iVXzc09DeEI7dnxMc/edit?usp=sharing) and recording [here](https://www.youtube.com/watch?v=6m6ZE6yVEDI)
 - [2025/11] We hosted [vLLM Beijing Meetup](https://mp.weixin.qq.com/s/xSrYXjNgr1HbCP4ExYNG1w) focusing on distributed inference and diverse accelerator support with vLLM! Please find the meetup slides [here](https://drive.google.com/drive/folders/1nQJ8ZkLSjKxvu36sSHaceVXtttbLvvu-?usp=drive_link).
 - [2025/10] We hosted [vLLM Shanghai Meetup](https://mp.weixin.qq.com/s/__xb4OyOsImz-9eAVrdlcg) focused on hands-on vLLM inference optimization! Please find the meetup slides [here](https://drive.google.com/drive/folders/1KqwjsFJLfEsC8wlDugnrR61zsWHt94Q6).
@@ -137,19 +136,16 @@ Compute Resources:
 - Alibaba Cloud
 - AMD
 - Anyscale
- Arm
 - AWS
 - Crusoe Cloud
 - Databricks
 - DeepInfra
 - Google Cloud
- IBM
 - Intel
 - Lambda Lab
 - Nebius
 - Novita AI
 - NVIDIA
- Red Hat
 - Replicate
 - Roblox
 - RunPod
@@ -83,7 +83,7 @@ MIN_CACHE_HIT_PCT=0
 MAX_LATENCY_ALLOWED_MS=100000000000 # A very large number
 ```

-### 2. Maximize Throughput with a Latency Requirement
+#### 2. Maximize Throughput with a Latency Requirement

 - **Goal**: Find the best server parameters when P99 end-to-end latency must be below 500ms.
 - **Configuration**:
@@ -96,7 +96,7 @@ MIN_CACHE_HIT_PCT=0
 MAX_LATENCY_ALLOWED_MS=500
 ```

-### 3. Maximize Throughput with Prefix Caching and Latency Requirements
+#### 3. Maximize Throughput with Prefix Caching and Latency Requirements

 - **Goal**: Find the best server parameters assuming a 60% prefix cache hit rate and a latency requirement of 500ms.
 - **Configuration**:
@@ -18,11 +18,6 @@ MIN_CACHE_HIT_PCT=${MIN_CACHE_HIT_PCT:-0}
 MAX_LATENCY_ALLOWED_MS=${MAX_LATENCY_ALLOWED_MS:-100000000000}
 NUM_SEQS_LIST=${NUM_SEQS_LIST:-"128 256"}
 NUM_BATCHED_TOKENS_LIST=${NUM_BATCHED_TOKENS_LIST:-"512 1024 2048 4096"}
-HOSTNAME=$(hostname)
-if [[ -z "$HOSTNAME" ]]; then
-    echo "Error: Failed to determine hostname." >&2
-    exit 1
-fi

 LOG_FOLDER="$BASE/auto-benchmark/$TAG"
 RESULT="$LOG_FOLDER/result.txt"
@@ -87,7 +82,6 @@ start_server() {
        "$MODEL"
        "--disable-log-requests"
        "--port" "8004"
-        "--host" "$HOSTNAME"
        "--gpu-memory-utilization" "$gpu_memory_utilization"
        "--max-num-seqs" "$max_num_seqs"
        "--max-num-batched-tokens" "$max_num_batched_tokens"
@@ -102,9 +96,8 @@ start_server() {
    # This correctly passes each element as a separate argument.
    if [[ -n "$profile_dir" ]]; then
        # Start server with profiling enabled
-        local profile_config_json="{\"profiler\": \"torch\", \"torch_profiler_dir\": \"$profile_dir\"}"
-        VLLM_SERVER_DEV_MODE=1 \
-            vllm serve --profiler-config "$profile_config_json" "${common_args_array[@]}" > "$vllm_log" 2>&1 &
+        VLLM_SERVER_DEV_MODE=1 VLLM_TORCH_PROFILER_DIR=$profile_dir \
+            vllm serve "${common_args_array[@]}" > "$vllm_log" 2>&1 &
    else
        # Start server without profiling
        VLLM_SERVER_DEV_MODE=1 \
@@ -119,7 +112,7 @@ start_server() {
        # since that we should always have permission to send signal to the server process.
        kill -0 $server_pid 2> /dev/null || break

-        RESPONSE=$(curl -s -X GET "http://${HOSTNAME}:8004/health" -w "%{http_code}" -o /dev/stdout)
+        RESPONSE=$(curl -s -X GET "http://0.0.0.0:8004/health" -w "%{http_code}" -o /dev/stdout)
        STATUS_CODE=$(echo "$RESPONSE" | tail -n 1)
        if [[ "$STATUS_CODE" -eq 200 ]]; then
            server_started=1
@@ -179,7 +172,6 @@ run_benchmark() {
        --goodput e2el:$MAX_LATENCY_ALLOWED_MS \
        --num-prompts 1000 \
        --random-prefix-len $prefix_len \
-        --host "$HOSTNAME" \
        --port 8004 &> "$bm_log"
    throughput=$(grep "Request throughput (req/s):" "$bm_log" | sed 's/[^0-9.]//g')
    e2el=$(grep "P99 E2EL (ms):" "$bm_log" | awk '{print $NF}')
@@ -195,7 +187,7 @@ run_benchmark() {
        request_rate=$((${throughput%.*} + 1))
        while ((request_rate > 0)); do
            # clear prefix cache
-            curl -X POST http://${HOSTNAME}:8004/reset_prefix_cache
+            curl -X POST http://0.0.0.0:8004/reset_prefix_cache
            sleep 5
            bm_log="$LOG_FOLDER/bm_log_${max_num_seqs}_${max_num_batched_tokens}_requestrate_${request_rate}.txt"
            vllm bench serve \
@@ -211,7 +203,6 @@ run_benchmark() {
                --goodput e2el:$MAX_LATENCY_ALLOWED_MS \
                --num-prompts 100 \
                --random-prefix-len $prefix_len \
-                --host "$HOSTNAME" \
                --port 8004 &> "$bm_log"
            throughput=$(grep "Request throughput (req/s):" "$bm_log" | sed 's/[^0-9.]//g')
            e2el=$(grep "P99 E2EL (ms):" "$bm_log" | awk '{print $NF}')
@@ -312,7 +303,6 @@ if (( $(echo "$best_throughput > 0" | bc -l) )); then
        --goodput e2el:$MAX_LATENCY_ALLOWED_MS \
        --num-prompts 100 \
        --random-prefix-len $prefix_len \
-        --host "$HOSTNAME" \
        --port 8004 \
        --profile &> "$bm_log"
 else
@@ -620,7 +620,7 @@ def get_tokenizer(
        kwargs["use_fast"] = False
    if tokenizer_mode == "mistral":
        try:
-            from vllm.tokenizers.mistral import MistralTokenizer
+            from vllm.transformers_utils.tokenizer import MistralTokenizer
        except ImportError as e:
            raise ImportError(
                "MistralTokenizer requires vllm package.\n"
@@ -1,120 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Micro benchmark comparing built-in hash(), SHA-256, and xxHash.
-
-This focuses on a single test payload shaped like the prefix-cache hash input:
-    (32-byte bytes object, 32-int tuple)
-
-Usage:
-    python benchmarks/hash_micro_benchmark.py --iterations 20000
-"""
-
-from __future__ import annotations
-
-import argparse
-import random
-import statistics
-import time
-from collections.abc import Callable, Iterable
-
-from vllm.utils.hashing import sha256, xxhash
-
-
-def _generate_test_data(seed: int) -> tuple[bytes, tuple[int, ...]]:
-    """Generate a deterministic test payload."""
-    random.seed(seed)
-    bytes_data = bytes(random.getrandbits(8) for _ in range(32))
-    int_tuple = tuple(random.randint(1, 1_000_000) for _ in range(32))
-    return (bytes_data, int_tuple)
-
-
-def _benchmark_func(func: Callable[[tuple], object], data: tuple, iterations: int):
-    """Return (avg_seconds, std_seconds) for hashing `data` `iterations` times."""
-    times: list[float] = []
-
-    # Warm-up to avoid first-run noise.
-    for _ in range(200):
-        func(data)
-
-    for _ in range(iterations):
-        start = time.perf_counter()
-        func(data)
-        end = time.perf_counter()
-        times.append(end - start)
-
-    avg = statistics.mean(times)
-    std = statistics.stdev(times) if len(times) > 1 else 0.0
-    return avg, std
-
-
-def _run_benchmarks(
-    benchmarks: Iterable[tuple[str, Callable[[tuple], object]]],
-    data: tuple,
-    iterations: int,
-):
-    """Yield (name, avg, std) for each benchmark, skipping unavailable ones."""
-    for name, func in benchmarks:
-        try:
-            avg, std = _benchmark_func(func, data, iterations)
-        except ModuleNotFoundError as exc:
-            print(f"Skipping {name}: {exc}")
-            continue
-        yield name, avg, std
-
-
-def builtin_hash(data: tuple) -> int:
-    """Wrapper for Python's built-in hash()."""
-    return hash(data)
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument(
-        "--iterations",
-        type=int,
-        default=10_000,
-        help="Number of measured iterations per hash function.",
-    )
-    parser.add_argument(
-        "--seed", type=int, default=42, help="Random seed for test payload."
-    )
-    args = parser.parse_args()
-
-    data = _generate_test_data(args.seed)
-    benchmarks = (
-        ("SHA256 (pickle)", sha256),
-        ("xxHash (pickle)", xxhash),
-        ("built-in hash()", builtin_hash),
-    )
-
-    print("=" * 60)
-    print("HASH FUNCTION MICRO BENCHMARK")
-    print("=" * 60)
-    print("Test data: (32-byte bytes object, 32-int tuple)")
-    print(f"Iterations: {args.iterations:,}")
-    print("=" * 60)
-
-    results = list(_run_benchmarks(benchmarks, data, args.iterations))
-    builtin_entry = next((r for r in results if r[0] == "built-in hash()"), None)
-
-    print("\nResults:")
-    for name, avg, std in results:
-        print(f"  {name:16s}: {avg * 1e6:8.2f} ± {std * 1e6:6.2f} μs")
-
-    if builtin_entry:
-        _, builtin_avg, _ = builtin_entry
-        print("\n" + "=" * 60)
-        print("SUMMARY (relative to built-in hash())")
-        print("=" * 60)
-        for name, avg, _ in results:
-            if name == "built-in hash()":
-                continue
-            speed_ratio = avg / builtin_avg
-            print(f"• {name} is {speed_ratio:.1f}x slower than built-in hash()")
-    else:
-        print("\nBuilt-in hash() result missing; cannot compute speed ratios.")
-
-
-if __name__ == "__main__":
-    main()
@@ -32,11 +32,12 @@ def benchmark_propose(args):

        model_config = ModelConfig(
            model="facebook/opt-125m",
+            task="generate",
            max_model_len=args.num_token + args.num_spec_token,
            tokenizer="facebook/opt-125m",
            tokenizer_mode="auto",
            dtype="auto",
-            seed=0,
+            seed=None,
            trust_remote_code=False,
        )
        proposer = NgramProposer(
@@ -107,10 +108,7 @@ def benchmark_batched_propose(args):
        device_config=DeviceConfig(device=current_platform.device_type),
        parallel_config=ParallelConfig(),
        load_config=LoadConfig(),
-        scheduler_config=SchedulerConfig(
-            max_model_len=model_config.max_model_len,
-            is_encoder_decoder=model_config.is_encoder_decoder,
-        ),
+        scheduler_config=SchedulerConfig(),
    )

    # monkey patch vllm.v1.worker.gpu_model_runner.get_pp_group
@@ -1,110 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""
-Simple benchmark to compare prefix-cache block hashing algorithms.
-
-Example:
-    python benchmark_prefix_block_hash.py --num-blocks 20000 --block-size 32
-"""
-
-from __future__ import annotations
-
-import argparse
-import random
-import statistics
-import sys
-import time
-from collections.abc import Callable, Iterable, Sequence
-
-from vllm.utils.hashing import get_hash_fn_by_name
-from vllm.v1.core.kv_cache_utils import BlockHash, hash_block_tokens, init_none_hash
-
-SUPPORTED_ALGOS = ("sha256", "sha256_cbor", "xxhash", "xxhash_cbor")
-
-
-def _generate_blocks(
-    num_blocks: int, block_size: int, vocab_size: int, seed: int
-) -> list[list[int]]:
-    rng = random.Random(seed)
-    return [
-        [rng.randrange(vocab_size) for _ in range(block_size)]
-        for _ in range(num_blocks)
-    ]
-
-
-def _hash_all_blocks(
-    hash_fn: Callable[[object], bytes],
-    blocks: Iterable[Sequence[int]],
-) -> float:
-    parent_hash: BlockHash | None = None
-    start = time.perf_counter()
-    for block in blocks:
-        parent_hash = hash_block_tokens(hash_fn, parent_hash, block, extra_keys=None)
-    end = time.perf_counter()
-    return end - start
-
-
-def _benchmark(
-    hash_algo: str,
-    blocks: list[list[int]],
-    trials: int,
-) -> tuple[float, float, float] | None:
-    try:
-        hash_fn = get_hash_fn_by_name(hash_algo)
-        init_none_hash(hash_fn)
-        timings = [_hash_all_blocks(hash_fn, blocks) for _ in range(trials)]
-    except ModuleNotFoundError as exc:
-        print(f"Skipping {hash_algo}: {exc}", file=sys.stderr)
-        return None
-
-    avg = statistics.mean(timings)
-    best = min(timings)
-    # throughput: tokens / second
-    tokens_hashed = len(blocks) * len(blocks[0])
-    throughput = tokens_hashed / best
-    return avg, best, throughput
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument("--num-blocks", type=int, default=10000, help="Block count.")
-    parser.add_argument("--block-size", type=int, default=32, help="Tokens per block.")
-    parser.add_argument(
-        "--vocab-size", type=int, default=32000, help="Token id range [0, vocab_size)."
-    )
-    parser.add_argument("--seed", type=int, default=0, help="Random seed.")
-    parser.add_argument(
-        "--trials", type=int, default=5, help="Number of timed trials per algorithm."
-    )
-    parser.add_argument(
-        "--algorithms",
-        nargs="+",
-        default=SUPPORTED_ALGOS,
-        choices=SUPPORTED_ALGOS,
-        help="Hash algorithms to benchmark.",
-    )
-    args = parser.parse_args()
-
-    blocks = _generate_blocks(
-        args.num_blocks, args.block_size, args.vocab_size, args.seed
-    )
-    print(
-        f"Benchmarking {len(args.algorithms)} algorithms on "
-        f"{args.num_blocks} blocks (block size={args.block_size})."
-    )
-
-    for algo in args.algorithms:
-        result = _benchmark(algo, blocks, args.trials)
-        if result is None:
-            continue
-
-        avg, best, throughput = result
-        print(
-            f"{algo:14s} avg: {avg:.6f}s  best: {best:.6f}s  "
-            f"throughput: {throughput / 1e6:.2f}M tokens/s"
-        )
-
-
-if __name__ == "__main__":
-    main()
@@ -40,7 +40,7 @@ from vllm.engine.arg_utils import EngineArgs
 from vllm.utils.argparse_utils import FlexibleArgumentParser

 try:
-    from vllm.tokenizers import get_tokenizer
+    from vllm.transformers_utils.tokenizer import get_tokenizer
 except ImportError:
    from backend_request_func import get_tokenizer

@@ -46,7 +46,7 @@ from tqdm.asyncio import tqdm
 from transformers import PreTrainedTokenizerBase

 try:
-    from vllm.tokenizers import get_tokenizer
+    from vllm.transformers_utils.tokenizer import get_tokenizer
 except ImportError:
    from backend_request_func import get_tokenizer

@@ -574,7 +574,7 @@ async def benchmark(
    )
    print(
        "{:<40} {:<10.2f}".format(
-            "Total token throughput (tok/s):", metrics.total_token_throughput
+            "Total Token throughput (tok/s):", metrics.total_token_throughput
        )
    )

@@ -963,7 +963,8 @@ def create_argument_parser():
    parser.add_argument(
        "--profile",
        action="store_true",
-        help="Use vLLM Profiling. --profiler-config must be provided on the server.",
+        help="Use Torch Profiler. The endpoint must be launched with "
+        "VLLM_TORCH_PROFILER_DIR to enable profiler.",
    )
    parser.add_argument(
        "--result-dir",
@@ -14,9 +14,6 @@ from tqdm import tqdm

 import vllm._custom_ops as ops
 from vllm.model_executor.layers.layernorm import RMSNorm
-from vllm.model_executor.layers.quantization.utils.fp8_utils import (
-    per_token_group_quant_fp8,
-)


@dataclass
@@ -25,7 +22,6 @@ class bench_params_t:
    hidden_size: int
    add_residual: bool
    dtype: torch.dtype
-    group_size: list[int]

    def description(self):
        return (
@@ -33,7 +29,6 @@ class bench_params_t:
            f"x D {self.hidden_size} "
            f"x R {self.add_residual} "
            f"x DT {self.dtype}"
-            f"x GS {self.group_size}"
        )


@@ -43,11 +38,10 @@ def get_bench_params() -> list[bench_params_t]:
    HIDDEN_SIZES = list(range(1024, 8129, 1024))
    ADD_RESIDUAL = [True, False]
    DTYPES = [torch.bfloat16, torch.float]
-    GROUP_SIZES = [[1, 64], [1, 128]]

-    combinations = product(NUM_TOKENS, HIDDEN_SIZES, ADD_RESIDUAL, DTYPES, GROUP_SIZES)
+    combinations = product(NUM_TOKENS, HIDDEN_SIZES, ADD_RESIDUAL, DTYPES)
    bench_params = list(
-        map(lambda x: bench_params_t(x[0], x[1], x[2], x[3], x[4]), combinations)
+        map(lambda x: bench_params_t(x[0], x[1], x[2], x[3]), combinations)
    )
    return bench_params

@@ -58,7 +52,6 @@ def unfused_int8_impl(
    x: torch.Tensor,
    residual: torch.Tensor | None,
    quant_dtype: torch.dtype,
-    group_size: list[int],
 ):
    # Norm
    torch_out = None
@@ -76,7 +69,6 @@ def unfused_fp8_impl(
    x: torch.Tensor,
    residual: torch.Tensor | None,
    quant_dtype: torch.dtype,
-    group_size: list[int],
 ):
    # Norm
    torch_out = None
@@ -89,63 +81,23 @@ def unfused_fp8_impl(
    torch_out, _ = ops.scaled_fp8_quant(torch_out)


-def unfused_groupwise_fp8_impl(
-    rms_norm_layer: RMSNorm,
-    x: torch.Tensor,
-    residual: torch.Tensor | None,
-    quant_dtype: torch.dtype,
-    group_size: list[int],
-):
-    # Norm
-    torch_out = None
-    if residual is None:
-        torch_out = rms_norm_layer.forward_cuda(x, residual)
-    else:
-        torch_out, _ = rms_norm_layer.forward_cuda(x, residual)
-
-    # Quant
-    torch_out, _ = per_token_group_quant_fp8(
-        torch_out, group_size=group_size[1], use_ue8m0=False
-    )
-
-
 def fused_impl(
    rms_norm_layer: RMSNorm,  # this stores the weights
    x: torch.Tensor,
    residual: torch.Tensor | None,
    quant_dtype: torch.dtype,
-    group_size: list[int],
 ):
    out, _ = ops.rms_norm_dynamic_per_token_quant(
        x, rms_norm_layer.weight, 1e-6, quant_dtype, residual=residual
    )


-def fused_groupwise_impl(
-    rms_norm_layer: RMSNorm,  # this stores the weights
-    x: torch.Tensor,
-    residual: torch.Tensor | None,
-    quant_dtype: torch.dtype,
-    group_size: list[int],
-):
-    out, _ = ops.rms_norm_per_block_quant(
-        x,
-        rms_norm_layer.weight,
-        1e-6,
-        quant_dtype,
-        group_size,
-        residual=residual,
-        is_scale_transposed=True,
-    )
-
-
 # Bench functions
 def bench_fn(
    rms_norm_layer: RMSNorm,
    x: torch.Tensor,
    residual: torch.Tensor,
    quant_dtype: torch.dtype,
-    group_size: list[int],
    label: str,
    sub_label: str,
    fn: Callable,
@@ -158,11 +110,10 @@ def bench_fn(
        "x": x,
        "residual": residual,
        "quant_dtype": quant_dtype,
-        "group_size": group_size,
        "fn": fn,
    }
    return TBenchmark.Timer(
-        stmt="fn(rms_norm_layer, x, residual, quant_dtype, group_size)",
+        stmt="fn(rms_norm_layer, x, residual, quant_dtype)",
        globals=globals,
        label=label,
        sub_label=sub_label,
@@ -196,7 +147,6 @@ def bench(params: bench_params_t, label: str, sub_label: str) -> Iterable[TMeasu
            x,
            residual,
            torch.int8,
-            params.group_size,
            label,
            sub_label,
            unfused_int8_impl,
@@ -211,7 +161,6 @@ def bench(params: bench_params_t, label: str, sub_label: str) -> Iterable[TMeasu
            x,
            residual,
            torch.float8_e4m3fn,
-            params.group_size,
            label,
            sub_label,
            unfused_fp8_impl,
@@ -226,7 +175,6 @@ def bench(params: bench_params_t, label: str, sub_label: str) -> Iterable[TMeasu
            x,
            residual,
            torch.int8,
-            params.group_size,
            label,
            sub_label,
            fused_impl,
@@ -241,7 +189,6 @@ def bench(params: bench_params_t, label: str, sub_label: str) -> Iterable[TMeasu
            x,
            residual,
            torch.float8_e4m3fn,
-            params.group_size,
            label,
            sub_label,
            fused_impl,
@@ -249,36 +196,6 @@ def bench(params: bench_params_t, label: str, sub_label: str) -> Iterable[TMeasu
        )
    )

-    # unfused groupwise fp8 impl.
-    timers.append(
-        bench_fn(
-            layer,
-            x,
-            residual,
-            torch.float8_e4m3fn,
-            params.group_size,
-            label,
-            sub_label,
-            unfused_groupwise_fp8_impl,
-            "unfused_groupwise_fp8_impl",
-        )
-    )
-
-    # fused groupwise fp8 impl.
-    timers.append(
-        bench_fn(
-            layer,
-            x,
-            residual,
-            torch.float8_e4m3fn,
-            params.group_size,
-            label,
-            sub_label,
-            fused_groupwise_impl,
-            "fused_groupwise_fp8_impl",
-        )
-    )
-
    print_timers(timers)

    return timers
@@ -1,244 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from dataclasses import dataclass
-from enum import Enum
-from itertools import product
-from typing import Any
-
-import torch
-import torch.utils.benchmark as TBenchmark
-from torch.utils.benchmark import Measurement as TMeasurement
-
-from vllm.model_executor.layers.quantization.utils.fp8_utils import (
-    _per_token_group_quant_fp8_colmajor,
-    silu_mul_per_token_group_quant_fp8_colmajor,
-)
-from vllm.triton_utils import triton
-from vllm.utils.deep_gemm import is_deep_gemm_e8m0_used
-
-from .utils import ArgPool, Bench, CudaGraphBenchParams
-
-GROUP_SIZE = 128
-FLOAT8_T = torch.float8_e4m3fn
-
-
-def print_timers(timers: list[TMeasurement], cuda_graph_nops: int):
-    print(
-        f"Note : The timings reported above is for {cuda_graph_nops} "
-        "consecutive invocations of the benchmarking functions. "
-        f"Please divide by {cuda_graph_nops} for single invocation "
-        "timings."
-    )
-    compare = TBenchmark.Compare(timers)
-    compare.print()
-
-
-class ImplType(Enum):
-    SILU_MUL_PER_TOKEN_GROUP_QUANT_FP8_COLMAJOR = 1
-    REFERENCE = 2
-
-    def get_impl(self):
-        if self == ImplType.SILU_MUL_PER_TOKEN_GROUP_QUANT_FP8_COLMAJOR:
-            return silu_mul_per_token_group_quant_fp8_colmajor
-        elif self == ImplType.REFERENCE:
-            return reference
-        raise ValueError(f"Unrecognized ImplType {self}")
-
-
-@dataclass
-class BenchmarkTensors:
-    input: torch.Tensor
-    output: torch.Tensor
-
-    # Reference act output tensor
-    ref_act_out: torch.Tensor
-    ref_quant_out: torch.Tensor
-
-    @staticmethod
-    def make(T: int, N: int) -> "BenchmarkTensors":
-        assert T % GROUP_SIZE == 0
-        assert N % (GROUP_SIZE * 2) == 0
-
-        input = torch.rand((T, N), dtype=torch.bfloat16, device="cuda")
-
-        # silu_mul_per_token_group_quant_fp8_colmajor output.
-        output = torch.rand((T, N // 2), dtype=torch.bfloat16, device="cuda").to(
-            FLOAT8_T
-        )
-
-        # reference output.
-        ref_act_out = torch.empty((T, N // 2), dtype=torch.bfloat16, device="cuda")
-        ref_quant_out = torch.empty(
-            (T, N // 2), dtype=torch.bfloat16, device="cuda"
-        ).to(FLOAT8_T)
-
-        return BenchmarkTensors(
-            input=input,
-            output=output,
-            ref_act_out=ref_act_out,
-            ref_quant_out=ref_quant_out,
-        )
-
-    @property
-    def T(self):
-        return self.input.size(0)
-
-    @property
-    def N(self):
-        return self.input.size(1)
-
-    def make_impl_kwargs(self, impl_type: ImplType) -> dict[str, Any]:
-        if impl_type == ImplType.SILU_MUL_PER_TOKEN_GROUP_QUANT_FP8_COLMAJOR:
-            return {
-                "input": self.input,
-                "output": self.output,
-                "use_ue8m0": is_deep_gemm_e8m0_used(),
-            }
-        elif impl_type == ImplType.REFERENCE:
-            return {
-                "input": self.input,
-                "act_out": self.ref_act_out,
-                "quant_out": self.ref_quant_out,
-                "use_ue8m0": is_deep_gemm_e8m0_used(),
-            }
-        raise ValueError(f"Unrecognized impl_type {impl_type}")
-
-
-def reference_quant(x: torch.Tensor, quant_out: torch.Tensor, use_ue8m0: bool):
-    """
-    Reference triton quant kernel from,
-    vllm.model_executor.layers.quantization.utils.fp8_utils
-    """
-    assert quant_out.size() == x.size()
-    # Allocate the scale tensor column-major format.
-    shape = (x.shape[-1] // GROUP_SIZE,) + x.shape[:-1]
-    x_q = quant_out
-    x_s = torch.empty(shape, device=x.device, dtype=torch.float32).permute(-1, -2)
-
-    M = x.numel() // GROUP_SIZE
-    N = GROUP_SIZE
-    BLOCK = triton.next_power_of_2(N)
-    # heuristics for number of warps
-    num_warps = min(max(BLOCK // 256, 1), 8)
-    num_stages = 1
-
-    finfo = torch.finfo(FLOAT8_T)
-    fp8_min = finfo.min
-    fp8_max = finfo.max
-
-    _per_token_group_quant_fp8_colmajor[(M,)](
-        x,
-        x_q,
-        x_s,
-        GROUP_SIZE,
-        x.shape[1],
-        x.stride(0),
-        x_s.stride(1),
-        eps=1e-10,
-        fp8_min=fp8_min,
-        fp8_max=fp8_max,
-        use_ue8m0=use_ue8m0,
-        BLOCK=BLOCK,
-        num_warps=num_warps,
-        num_stages=num_stages,
-    )
-    return x_q, x_s
-
-
-def reference(
-    input: torch.Tensor,
-    act_out: torch.Tensor,
-    quant_out: torch.Tensor,
-    use_ue8m0: bool,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    torch.ops._C.silu_and_mul(act_out, input)
-    return reference_quant(act_out, quant_out, use_ue8m0)
-
-
-def bench_impl(
-    bench_tensors: list[BenchmarkTensors], impl_type: ImplType
-) -> TMeasurement:
-    T = bench_tensors[0].T
-    N = bench_tensors[0].N
-
-    arg_pool_size = len(bench_tensors)
-    kwargs_list = [bt.make_impl_kwargs(impl_type) for bt in bench_tensors]
-
-    # warmup
-    for kwargs in kwargs_list:
-        impl_type.get_impl()(**kwargs)
-    torch.cuda.synchronize()
-
-    # Merge into a single kwargs and qualify arguments as ArgPool
-    kwargs = {k: ArgPool([]) for k in kwargs_list[0]}
-    for _kwargs in kwargs_list:
-        for k, v in _kwargs.items():
-            kwargs[k].values.append(v)
-
-    cuda_graph_params = None
-    cuda_graph_params = CudaGraphBenchParams(arg_pool_size)
-    timer = None
-    with Bench(
-        cuda_graph_params,
-        "silu-mul-quant",
-        f"num_tokens={T}, N={N}",
-        impl_type.name,
-        impl_type.get_impl(),
-        **kwargs,
-    ) as bench:
-        timer = bench.run()
-    return timer
-
-
-def test_correctness(T: int, N: int):
-    print(f"Testing num_tokens={T}, N={N} ...")
-
-    bench_tensor = BenchmarkTensors.make(T, N)
-
-    def output_from_impl(impl: ImplType) -> tuple[torch.Tensor, torch.Tensor]:
-        return impl.get_impl()(**bench_tensor.make_impl_kwargs(impl))
-
-    # reference output
-    ref_out_q, ref_out_s = output_from_impl(ImplType.REFERENCE)
-
-    # test ouptut
-    out_q, out_s = output_from_impl(
-        ImplType.SILU_MUL_PER_TOKEN_GROUP_QUANT_FP8_COLMAJOR
-    )
-
-    torch.testing.assert_close(ref_out_q.to(torch.float32), out_q.to(torch.float32))
-    torch.testing.assert_close(ref_out_s, out_s)
-
-
-def run(Ts: list[int], Ns: list[int], arg_pool_size: int) -> list[TMeasurement]:
-    timers = []
-    for N, T in product(Ns, Ts):
-        test_correctness(T, N)
-
-        bench_tensors: list[BenchmarkTensors] = [
-            BenchmarkTensors.make(T, N) for _ in range(arg_pool_size)
-        ]
-
-        silu_mul_quant_timer = bench_impl(
-            bench_tensors, ImplType.SILU_MUL_PER_TOKEN_GROUP_QUANT_FP8_COLMAJOR
-        )
-        timers.append(silu_mul_quant_timer)
-        reference_timer = bench_impl(bench_tensors, ImplType.REFERENCE)
-        timers.append(reference_timer)
-
-        print_timers(
-            [silu_mul_quant_timer, reference_timer], cuda_graph_nops=arg_pool_size
-        )
-
-    print_timers(timers, cuda_graph_nops=arg_pool_size)
-
-    return timers
-
-
-if __name__ == "__main__":
-    T = [128 * i for i in range(1, 16)] + [2048 * i for i in range(1, 65)]
-    N = [2048, 4096, 8192]
-
-    print(f"T = {T}, N = {N}")
-    run(T, N, arg_pool_size=8)
@@ -237,7 +237,6 @@ def marlin_create_bench_fn(bt: BenchmarkTensors) -> Callable:
            b_q_weight=w_q,
            b_bias=None,
            b_scales=w_s,
-            a_scales=None,
            global_scale=None,
            b_zeros=w_zp,
            g_idx=g_idx,
@@ -263,7 +263,7 @@ def bench_run(

    results.append(
        benchmark.Timer(
-            stmt="output = gptq_marlin_gemm(a, None, marlin_q_w, marlin_s, None, marlin_s2, marlin_zp, marlin_g_idx, marlin_sort_indices, marlin_workspace.scratch, quant_type, size_m, size_n, size_k, is_k_full, False, False, False)",  # noqa: E501
+            stmt="output = gptq_marlin_gemm(a, None, marlin_q_w, marlin_s, marlin_s2, marlin_zp, marlin_g_idx, marlin_sort_indices, marlin_workspace.scratch, quant_type, size_m, size_n, size_k, is_k_full, False, False, False)",  # noqa: E501
            globals=globals,
            label=label,
            sub_label=sub_label,
@@ -273,7 +273,7 @@ def bench_run(

    results.append(
        benchmark.Timer(
-            stmt="output = gptq_marlin_gemm(a, None, marlin_q_w, marlin_s, None, marlin_s2, marlin_zp, marlin_g_idx, marlin_sort_indices, marlin_workspace.scratch, quant_type, size_m, size_n, size_k, is_k_full, False, True, False)",  # noqa: E501
+            stmt="output = gptq_marlin_gemm(a, None, marlin_q_w, marlin_s, marlin_s2, marlin_zp, marlin_g_idx, marlin_sort_indices, marlin_workspace.scratch, quant_type, size_m, size_n, size_k, is_k_full, False, True, False)",  # noqa: E501
            globals=globals,
            label=label,
            sub_label=sub_label,
@@ -1,150 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Benchmark script comparing torch.cat vs direct copy for k_nope/k_pe concatenation
-in MLA (Multi-head Latent Attention) prefill.
-
-This validates that the optimization from commit 8d4142bd is beneficial across
-various batch sizes, not just the originally tested batch size of 32768.
-"""
-
-import time
-from collections.abc import Callable
-
-import torch
-
-# DeepSeek-V3 MLA dimensions
-NUM_HEADS = 128
-QK_NOPE_HEAD_DIM = 128
-PE_DIM = 64
-
-
-def cat_method(k_nope: torch.Tensor, k_pe: torch.Tensor) -> torch.Tensor:
-    """Original torch.cat approach with expand."""
-    return torch.cat((k_nope, k_pe.expand((*k_nope.shape[:-1], -1))), dim=-1)
-
-
-def direct_copy_method(k_nope: torch.Tensor, k_pe: torch.Tensor) -> torch.Tensor:
-    """Optimized direct copy approach (avoids expand + cat overhead)."""
-    k = torch.empty(
-        (*k_nope.shape[:-1], k_nope.shape[-1] + k_pe.shape[-1]),
-        dtype=k_nope.dtype,
-        device=k_nope.device,
-    )
-    k[..., : k_nope.shape[-1]] = k_nope
-    k[..., k_nope.shape[-1] :] = k_pe
-    return k
-
-
-def benchmark_method(
-    method: Callable,
-    k_nope: torch.Tensor,
-    k_pe: torch.Tensor,
-    num_warmup: int = 10,
-    num_iters: int = 100,
-) -> float:
-    """Benchmark a concatenation method and return mean latency in ms."""
-    # Warmup
-    for _ in range(num_warmup):
-        _ = method(k_nope, k_pe)
-    torch.cuda.synchronize()
-
-    # Benchmark
-    start = time.perf_counter()
-    for _ in range(num_iters):
-        _ = method(k_nope, k_pe)
-    torch.cuda.synchronize()
-    end = time.perf_counter()
-
-    return (end - start) / num_iters * 1000  # Convert to ms
-
-
-@torch.inference_mode()
-def run_benchmark(dtype: torch.dtype, dtype_name: str):
-    """Run benchmark for a specific dtype."""
-    torch.set_default_device("cuda")
-
-    # Batch sizes to test (powers of 2 from 32 to 65536)
-    batch_sizes = [32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768, 65536]
-
-    print("=" * 80)
-    print("Benchmark: torch.cat vs direct copy for MLA k_nope/k_pe concatenation")
-    print("=" * 80)
-    print(
-        f"Tensor shapes: k_nope=[B, {NUM_HEADS}, {QK_NOPE_HEAD_DIM}], "
-        f"k_pe=[B, 1, {PE_DIM}]"
-    )
-    print(f"dtype: {dtype_name}")
-    print()
-    print(
-        f"{'Batch Size':>12} | {'cat (ms)':>10} | {'direct (ms)':>12} | "
-        f"{'Speedup':>8} | {'Reduction':>10}"
-    )
-    print("-" * 70)
-
-    results = []
-    for batch_size in batch_sizes:
-        # Create input tensors (generate in float32 then convert for FP8 compatibility)
-        k_nope = torch.randn(
-            batch_size, NUM_HEADS, QK_NOPE_HEAD_DIM, dtype=torch.float32, device="cuda"
-        ).to(dtype)
-        k_pe = torch.randn(
-            batch_size, 1, PE_DIM, dtype=torch.float32, device="cuda"
-        ).to(dtype)
-
-        # Benchmark both methods
-        cat_time = benchmark_method(cat_method, k_nope, k_pe)
-        direct_time = benchmark_method(direct_copy_method, k_nope, k_pe)
-
-        speedup = cat_time / direct_time
-        reduction = (1 - direct_time / cat_time) * 100
-
-        results.append((batch_size, cat_time, direct_time, speedup, reduction))
-
-        print(
-            f"{batch_size:>12} | {cat_time:>10.3f} | {direct_time:>12.3f} | "
-            f"{speedup:>7.2f}x | {reduction:>9.1f}%"
-        )
-
-    print("=" * 80)
-
-    # Summary statistics
-    speedups = [r[3] for r in results]
-    print("\nSpeedup summary:")
-    print(f"  Min:  {min(speedups):.2f}x")
-    print(f"  Max:  {max(speedups):.2f}x")
-    print(f"  Mean: {sum(speedups) / len(speedups):.2f}x")
-
-    # Find crossover point
-    crossover_batch = None
-    for batch_size, _, _, speedup, _ in results:
-        if speedup >= 1.0:
-            crossover_batch = batch_size
-            break
-
-    print("\nConclusion:")
-    if crossover_batch:
-        print(f"  - Direct copy becomes beneficial at batch size >= {crossover_batch}")
-    # Filter for large batches (>= 512 which is typical for prefill)
-    large_batch_speedups = [r[3] for r in results if r[0] >= 512]
-    if large_batch_speedups:
-        avg_large = sum(large_batch_speedups) / len(large_batch_speedups)
-        print(f"  - For batch sizes >= 512: avg speedup = {avg_large:.2f}x")
-    print("  - MLA prefill typically uses large batches, so optimization is effective")
-
-    return results
-
-
-@torch.inference_mode()
-def main():
-    # Test bfloat16
-    print("\n")
-    run_benchmark(torch.bfloat16, "bfloat16")
-
-    # Test float8_e4m3fn
-    print("\n")
-    run_benchmark(torch.float8_e4m3fn, "float8_e4m3fn")
-
-
-if __name__ == "__main__":
-    main()
@@ -24,15 +24,12 @@ def get_topk_ids(num_tokens: int, num_experts: int, topk: int) -> torch.Tensor:
 num_tokens_range = [1, 16, 256, 4096]
 num_experts_range = [16, 64, 224, 256, 280, 512]
 topk_range = [1, 2, 8]
-ep_size_range = [1, 8]
-configs = list(
-    itertools.product(num_tokens_range, num_experts_range, topk_range, ep_size_range)
-)
+configs = list(itertools.product(num_tokens_range, num_experts_range, topk_range))


@triton.testing.perf_report(
    triton.testing.Benchmark(
-        x_names=["num_tokens", "num_experts", "topk", "ep_size"],
+        x_names=["num_tokens", "num_experts", "topk"],
        x_vals=configs,
        line_arg="provider",
        line_vals=["vllm"],
@@ -41,26 +38,16 @@ configs = list(
        args={},
    )
 )
-def benchmark(num_tokens, num_experts, topk, ep_size, provider):
+def benchmark(num_tokens, num_experts, topk, provider):
    """Benchmark function for Triton."""
    block_size = 256
-    torch.cuda.manual_seed_all(0)
    topk_ids = get_topk_ids(num_tokens, num_experts, topk)

-    e_map = None
-    if ep_size != 1:
-        local_e = num_experts // ep_size
-        e_ids = torch.randperm(num_experts, device="cuda", dtype=torch.int32)[:local_e]
-        e_map = torch.full((num_experts,), -1, device="cuda", dtype=torch.int32)
-        e_map[e_ids] = torch.arange(local_e, device="cuda", dtype=torch.int32)
-
    quantiles = [0.5, 0.2, 0.8]

    if provider == "vllm":
        ms, min_ms, max_ms = triton.testing.do_bench(
-            lambda: moe_align_block_size(
-                topk_ids, block_size, num_experts, e_map, ignore_invalid_experts=True
-            ),
+            lambda: moe_align_block_size(topk_ids, block_size, num_experts),
            quantiles=quantiles,
        )

@@ -0,0 +1,75 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import torch
+
+from vllm._custom_ops import moe_fused_gate
+from vllm.model_executor.layers.fused_moe.fused_moe import (
+    grouped_topk as vllm_compiled_grouped_topk,
+)
+from vllm.triton_utils import triton
+
+
+def biased_grouped_topk_org(scores, bias, num_expert_group, topk_group, topk):
+    return vllm_compiled_grouped_topk(
+        hidden_states=scores,
+        gating_output=scores,
+        topk=topk,
+        renormalize=True,
+        num_expert_group=num_expert_group,
+        topk_group=topk_group,
+        scoring_func="sigmoid",
+        e_score_correction_bias=bias,
+    )
+
+
+def biased_grouped_topk_org_kernel(scores, bias, num_expert_group, topk_group, topk):
+    return moe_fused_gate(scores, bias, num_expert_group, topk_group, topk)
+
+
+seq_length_range = [5000, 10000, 15000, 20000, 25000, 30000, 35000, 40000]
+configs = [(sq,) for sq in seq_length_range]
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["seq_length"],
+        x_vals=[list(_) for _ in configs],
+        line_arg="provider",
+        line_vals=["original", "kernel"],
+        line_names=["Original", "SGL Kernel"],
+        styles=[("blue", "-"), ("red", "-")],
+        ylabel="us",
+        plot_name="moe-fused-gate-performance",
+        args={},
+    )
+)
+def benchmark(seq_length, provider):
+    dtype = torch.bfloat16
+    device = torch.device("cuda")
+    num_experts, num_expert_group, topk_group, topk = 256, 8, 4, 8
+
+    scores = torch.randn((seq_length, num_experts), device=device, dtype=dtype)
+    bias = torch.rand(num_experts, device=device, dtype=dtype)
+
+    quantiles = [0.5, 0.2, 0.8]
+
+    if provider == "original":
+        ms, min_ms, max_ms = triton.testing.do_bench(
+            lambda: biased_grouped_topk_org(
+                scores.clone(), bias.clone(), num_expert_group, topk_group, topk
+            ),
+            quantiles=quantiles,
+        )
+    elif provider == "kernel":
+        ms, min_ms, max_ms = triton.testing.do_bench(
+            lambda: biased_grouped_topk_org_kernel(
+                scores.clone(), bias.clone(), num_expert_group, topk_group, topk
+            ),
+            quantiles=quantiles,
+        )
+
+    return 1000 * ms, 1000 * max_ms, 1000 * min_ms
+
+
+if __name__ == "__main__":
+    benchmark.run(print_data=True)
@@ -99,6 +99,7 @@ def benchmark_mrope(
    # the parameters to compute the q k v size based on tp_size
    mrope_helper_class = get_rope(
        head_size=head_dim,
+        rotary_dim=head_dim,
        max_position=max_position,
        is_neox_style=is_neox_style,
        rope_parameters=rope_parameters,
@@ -32,8 +32,8 @@ def get_benchmark(head_size, rotary_dim, is_neox_style, device):
    def benchmark(batch_size, seq_len, num_heads, provider):
        dtype = torch.bfloat16
        max_position = 8192
-        rope_parameters = {"partial_rotary_factor": rotary_dim / head_size}
-        rope = get_rope(head_size, max_position, is_neox_style, rope_parameters)
+        base = 10000
+        rope = get_rope(head_size, rotary_dim, max_position, base, is_neox_style)
        rope = rope.to(dtype=dtype, device=device)
        cos_sin_cache = rope.cos_sin_cache.to(dtype=torch.float, device=device)

@@ -2,7 +2,7 @@

 This directory includes benchmarks between DeepSeek's DeepGEMM block fp8 kernels against vLLM's existing triton and CUTLASS-based kernels.

-Currently, this just includes dense GEMMs and only works on Hopper GPUs.
+Currently this just includes dense GEMMs and only works on Hopper GPUs.

 ## Setup

@@ -251,6 +251,17 @@ if ((AVX512_FOUND AND NOT AVX512_DISABLED) OR (ASIMD_FOUND AND NOT APPLE_SILICON
        endif()

        # Build ACL with CMake
+        set(ARM_COMPUTE_BUILD_SHARED_LIB "OFF")
+        set(CMAKE_BUILD_TYPE "Release")
+        set(ARM_COMPUTE_ARCH "armv8.2-a")
+        set(ARM_COMPUTE_ENABLE_ASSERTS "OFF")
+        set(ARM_COMPUTE_ENABLE_CPPTHREADS "OFF")
+        set(ONEDNN_ENABLE_PRIMITIVE "MATMUL;REORDER")
+        set(ARM_COMPUTE_ENABLE_OPENMP "ON")
+        set(ARM_COMPUTE_ENABLE_WERROR "OFF")
+        set(ARM_COMPUTE_BUILD_EXAMPLES "OFF")
+        set(ARM_COMPUTE_BUILD_TESTING "OFF")
+
        set(_cmake_config_cmd
             ${CMAKE_COMMAND} -G Ninja -B build 
            -DARM_COMPUTE_BUILD_SHARED_LIB=OFF 
@@ -38,7 +38,7 @@ else()
  FetchContent_Declare(
          vllm-flash-attn
          GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
-          GIT_TAG 86f8f157cf82aa2342743752b97788922dd7de43
+          GIT_TAG 71bb26f6295449be880344b93b51791cc009237d
          GIT_PROGRESS TRUE
          # Don't share the vllm-flash-attn build between build types
          BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
@@ -140,21 +140,16 @@ function(vllm_prepare_torch_gomp_shim TORCH_GOMP_SHIM_DIR)
  run_python(_VLLM_TORCH_GOMP_PATH
    "
 import os, glob
-import torch
-torch_pkg = os.path.dirname(torch.__file__)
-site_root = os.path.dirname(torch_pkg)
-
-# Search both torch.libs and torch/lib
-roots = [os.path.join(site_root, 'torch.libs'), os.path.join(torch_pkg, 'lib')]
-candidates = []
-for root in roots:
-    if not os.path.isdir(root):
-        continue
-    candidates.extend(glob.glob(os.path.join(root, 'libgomp*.so*')))
-
-print(candidates[0] if candidates else '')
+try:
+  import torch
+  torch_pkg = os.path.dirname(torch.__file__)
+  site_root = os.path.dirname(torch_pkg)
+  torch_libs = os.path.join(site_root, 'torch.libs')
+  print(glob.glob(os.path.join(torch_libs, 'libgomp-*.so*'))[0])
+except:
+  print('')
 "
-    "failed to probe for libgomp")
+    "failed to probe torch.libs for libgomp")

  if(_VLLM_TORCH_GOMP_PATH STREQUAL "" OR NOT EXISTS "${_VLLM_TORCH_GOMP_PATH}")
    return()
@@ -500,13 +495,7 @@ function (define_extension_target MOD_NAME)
    set(SOABI_KEYWORD "")
  endif()

-  run_python(IS_FREETHREADED_PYTHON
-    "import sysconfig; print(1 if sysconfig.get_config_var(\"Py_GIL_DISABLED\") else 0)"
-    "Failed to determine whether interpreter is free-threaded")
-
-  # Free-threaded Python doesn't yet support the stable ABI (see PEP 803/809),
-  # so avoid using the stable ABI under free-threading only.
-  if (ARG_USE_SABI AND NOT IS_FREETHREADED_PYTHON)
+  if (ARG_USE_SABI)
    Python_add_library(${MOD_NAME} MODULE USE_SABI ${ARG_USE_SABI} ${SOABI_KEYWORD} "${ARG_SOURCES}")
  else()
    Python_add_library(${MOD_NAME} MODULE ${SOABI_KEYWORD} "${ARG_SOURCES}")
@@ -16,8 +16,7 @@ __global__ void merge_attn_states_kernel(
    scalar_t* output, float* output_lse, const scalar_t* prefix_output,
    const float* prefix_lse, const scalar_t* suffix_output,
    const float* suffix_lse, const uint num_tokens, const uint num_heads,
-    const uint head_size, const uint prefix_head_stride,
-    const uint output_head_stride) {
+    const uint head_size) {
  using pack_128b_t = uint4;
  const uint pack_size = 16 / sizeof(scalar_t);
  const uint threads_per_head = head_size / pack_size;
@@ -35,13 +34,11 @@ __global__ void merge_attn_states_kernel(
  const uint head_idx = token_head_idx % num_heads;

  const uint pack_offset = pack_idx * pack_size;  // (0~15)*8, etc.
-  const uint src_head_offset = token_idx * num_heads * prefix_head_stride +
-                               head_idx * prefix_head_stride;
-  const uint dst_head_offset = token_idx * num_heads * output_head_stride +
-                               head_idx * output_head_stride;
-  const scalar_t* prefix_head_ptr = prefix_output + src_head_offset;
-  const scalar_t* suffix_head_ptr = suffix_output + src_head_offset;
-  scalar_t* output_head_ptr = output + dst_head_offset;
+  const uint head_offset =
+      token_idx * num_heads * head_size + head_idx * head_size;
+  const scalar_t* prefix_head_ptr = prefix_output + head_offset;
+  const scalar_t* suffix_head_ptr = suffix_output + head_offset;
+  scalar_t* output_head_ptr = output + head_offset;

  float p_lse = prefix_lse[head_idx * num_tokens + token_idx];
  float s_lse = suffix_lse[head_idx * num_tokens + token_idx];
@@ -143,7 +140,7 @@ __global__ void merge_attn_states_kernel(
            reinterpret_cast<float*>(prefix_lse.data_ptr()),                \
            reinterpret_cast<scalar_t*>(suffix_output.data_ptr()),          \
            reinterpret_cast<float*>(suffix_lse.data_ptr()), num_tokens,    \
-            num_heads, head_size, prefix_head_stride, output_head_stride);  \
+            num_heads, head_size);                                          \
  }

 /*@brief Merges the attention states from prefix and suffix
@@ -169,11 +166,17 @@ void merge_attn_states_launcher(torch::Tensor& output,
  const uint num_tokens = output.size(0);
  const uint num_heads = output.size(1);
  const uint head_size = output.size(2);
-  const uint prefix_head_stride = prefix_output.stride(1);
-  const uint output_head_stride = output.stride(1);
  const uint pack_size = 16 / sizeof(scalar_t);
  TORCH_CHECK(head_size % pack_size == 0,
              "headsize must be multiple of pack_size:", pack_size);
+  TORCH_CHECK(output.stride(-2) == head_size && output.stride(-1) == 1,
+              "output heads must be contiguous in memory");
+  TORCH_CHECK(
+      prefix_output.stride(-2) == head_size && prefix_output.stride(-1) == 1,
+      "prefix_output heads must be contiguous in memory");
+  TORCH_CHECK(
+      suffix_output.stride(-2) == head_size && suffix_output.stride(-1) == 1,
+      "suffix_output heads must be contiguous in memory");
  float* output_lse_ptr = nullptr;
  if (output_lse.has_value()) {
    output_lse_ptr = output_lse.value().data_ptr<float>();
@@ -1,7 +1,6 @@
 #pragma once

 #include <torch/all.h>
-#include <c10/util/Optional.h>

 #include <map>
 #include <vector>
@@ -42,12 +41,11 @@ void convert_fp8(torch::Tensor& dst_cache, torch::Tensor& src_cache,
                 const double scale, const std::string& kv_cache_dtype);

 void gather_and_maybe_dequant_cache(
-    torch::Tensor const& src_cache,     // [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...]
-    torch::Tensor const& dst,           // [TOT_TOKENS, ENTRIES...]
-    torch::Tensor const& block_table,   // [BATCH, BLOCK_INDICES]
-    torch::Tensor const& cu_seq_lens,   // [BATCH+1]
-    torch::Tensor const& token_to_seq,  // [MAX_TOKEN_ACROSS_CHUNKS]
-    int64_t num_tokens, const std::string& kv_cache_dtype,
+    torch::Tensor const& src_cache,    // [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...]
+    torch::Tensor const& dst,          // [TOT_TOKENS, ENTRIES...]
+    torch::Tensor const& block_table,  // [BATCH, BLOCK_INDICES]
+    torch::Tensor const& cu_seq_lens,  // [BATCH+1]
+    int64_t batch_size, const std::string& kv_cache_dtype,
    torch::Tensor const& scale,
    std::optional<torch::Tensor> seq_starts = std::nullopt);

@@ -59,15 +57,6 @@ void cp_gather_cache(
    torch::Tensor const& cu_seq_lens,  // [BATCH+1]
    int64_t batch_size, std::optional<torch::Tensor> seq_starts = std::nullopt);

-// Gather and upconvert FP8 KV cache to BF16 workspace
-void cp_gather_and_upconvert_fp8_kv_cache(
-    torch::Tensor const& src_cache,         // [NUM_BLOCKS, BLOCK_SIZE, 656]
-    torch::Tensor const& dst,               // [TOT_TOKENS, 576]
-    torch::Tensor const& block_table,       // [BATCH, BLOCK_INDICES]
-    torch::Tensor const& seq_lens,          // [BATCH]
-    torch::Tensor const& workspace_starts,  // [BATCH]
-    int64_t batch_size);
-
 // Indexer K quantization and cache function
 void indexer_k_quant_and_cache(
    torch::Tensor& k,             // [num_tokens, head_dim]
@@ -82,4 +71,4 @@ void cp_gather_indexer_k_quant_cache(
    torch::Tensor& dst_k,           // [num_tokens, head_dim]
    torch::Tensor& dst_scale,  // [num_tokens, head_dim / quant_block_size * 4]
    const torch::Tensor& block_table,   // [batch_size, num_blocks]
-    const torch::Tensor& cu_seq_lens);  // [batch_size + 1]
+    const torch::Tensor& cu_seq_lens);  // [batch_size + 1]
@@ -2,7 +2,6 @@
 #include <ATen/cuda/CUDAContext.h>
 #include <c10/cuda/CUDAGuard.h>
 #include <c10/cuda/CUDAException.h>
-#include <c10/util/Optional.h>

 #include "cuda_utils.h"
 #include "cuda_compat.h"
@@ -515,8 +514,7 @@ __global__ void indexer_k_quant_and_cache_kernel(
    const int quant_block_size,                // quantization block size
    const int cache_block_size,                // cache block size
    const int cache_stride,  // stride for each token in kv_cache
-
-    const bool use_ue8m0  // use ue8m0 scale format
+    const bool use_ue8m0     // use ue8m0 scale format
 ) {
  constexpr int VEC_SIZE = 4;
  const int64_t token_idx = blockIdx.x;
@@ -907,79 +905,91 @@ void convert_fp8(torch::Tensor& dst_cache, torch::Tensor& src_cache,
 namespace vllm {

 // grid is launched with dimensions (batch, num_splits)
-template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt,
-          int ENTRY_SIZE, int CTA_SIZE>
+template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
 __global__ void gather_and_maybe_dequant_cache(
-    const cache_t* __restrict__ src_cache,     // [NUM_BLOCKS, BLOCK_SIZE,
-                                               // ENTRIES...]
-    scalar_t* __restrict__ dst,                // [TOT_TOKENS, ENTRIES...]
-    const int32_t* __restrict__ block_table,   // [BATCH, BLOCK_INDICES]
-    const int32_t* __restrict__ cu_seq_lens,   // [BATCH+1]
-    const int32_t* __restrict__ token_to_seq,  // [MAX_TOKEN_ACROSS_CHUNK]
-    const int32_t num_tokens, const int32_t block_size,
+    const cache_t* __restrict__ src_cache,    // [NUM_BLOCKS, BLOCK_SIZE,
+                                              // ENTRIES...]
+    scalar_t* __restrict__ dst,               // [TOT_TOKENS, ENTRIES...]
+    const int32_t* __restrict__ block_table,  // [BATCH, BLOCK_INDICES]
+    const int32_t* __restrict__ cu_seq_lens,  // [BATCH+1]
+    const int32_t block_size, const int32_t entry_size,
    const int64_t block_table_stride, const int64_t cache_block_stride,
    const int64_t cache_entry_stride, const int64_t dst_entry_stride,
    const float* __restrict__ scale,
    const int32_t* __restrict__ seq_starts) {  // Optional: starting offsets per
                                               // batch
-  constexpr int vec_size = sizeof(float4) / sizeof(scalar_t);
-  using ltype = vllm::vec_n_t<cache_t, vec_size>;
-  using stype = vllm::vec_n_t<scalar_t, vec_size>;
-  // We are adding this for code readability which will be optimized out when
-  // build in release.
-  assert(CTA_SIZE == blockDim.x);

-#pragma unroll
-  for (int token_id = blockIdx.x; token_id < num_tokens;
-       token_id += gridDim.x) {
-    int64_t batch_id = token_to_seq[token_id];
-    int64_t batch_start = cu_seq_lens[batch_id];
-    int64_t batch_end = cu_seq_lens[batch_id + 1];
-    int32_t batch_offset = token_id - batch_start;
+  const int64_t bid = blockIdx.x;  // Batch ID
+  const int32_t num_splits = gridDim.y;
+  const int32_t split = blockIdx.y;
+  const int32_t seq_start = cu_seq_lens[bid];
+  const int32_t seq_end = cu_seq_lens[bid + 1];
+  const int32_t seq_len = seq_end - seq_start;
+  const int32_t tot_blocks = cuda_utils::ceil_div(seq_len, block_size);
+  const int32_t split_blocks = cuda_utils::ceil_div(tot_blocks, num_splits);

-    if (token_id >= batch_end) return;
-    int32_t offset = 0;
-    if (seq_starts != nullptr) {
-      offset = seq_starts[batch_id];
-    }
-    batch_offset += offset;
-    int32_t block_table_id = batch_offset / block_size;
-    int32_t slot_id = batch_offset % block_size;
-    int32_t block_table_offset = batch_id * block_table_stride + block_table_id;
-    int32_t block_id = block_table[block_table_offset];
-    int64_t cache_offset =
-        block_id * cache_block_stride + slot_id * cache_entry_stride;
-    constexpr int32_t vec_iter_cnt = ENTRY_SIZE / vec_size;
-    scalar_t* dst_ = dst + token_id * dst_entry_stride;
-    cache_t* src_ = const_cast<cache_t*>(src_cache) + cache_offset;
+  const int32_t split_start = split * split_blocks;
+  const int32_t split_end = min((split + 1) * split_blocks, tot_blocks);

-#pragma unroll
-    for (int idx = threadIdx.x; idx < vec_iter_cnt; idx += CTA_SIZE) {
+  const bool is_active_split = (split_start < tot_blocks);
+  const bool is_last_split = (split_end == tot_blocks);
+
+  if (!is_active_split) return;
+
+  int32_t full_blocks_end = split_end;
+  int32_t partial_block_size = 0;
+
+  // Adjust the pointer for the block_table for this batch.
+  // If seq_starts is provided, compute an offset based on (seq_starts[bid] /
+  // page_size)
+  const int32_t batch_offset = bid * block_table_stride;
+  int32_t offset = 0;
+  if (seq_starts != nullptr) {
+    offset = seq_starts[bid] / block_size;
+  }
+  const int32_t* batch_block_table = block_table + batch_offset + offset;
+
+  // Adjust dst pointer based on the cumulative sequence lengths.
+  dst += seq_start * dst_entry_stride;
+
+  if (is_last_split) {
+    partial_block_size = seq_len % block_size;
+    if (partial_block_size) full_blocks_end -= 1;
+  }
+
+  auto copy_entry = [&](const cache_t* __restrict__ _src,
+                        scalar_t* __restrict__ _dst) {
+    for (int i = threadIdx.x; i < entry_size; i += blockDim.x) {
      if constexpr (kv_dt == Fp8KVCacheDataType::kAuto) {
-        reinterpret_cast<stype*>(dst_)[idx] =
-            static_cast<stype>(reinterpret_cast<ltype*>(src_)[idx]);
+        _dst[i] = static_cast<scalar_t>(_src[i]);
      } else {
-        ltype loaded_val = reinterpret_cast<ltype*>(src_)[idx];
-        stype store_val;
-#pragma unroll
-        for (int j = 0; j < vec_size; ++j) {
-          store_val.val[j] = fp8::scaled_convert<scalar_t, cache_t, kv_dt>(
-              loaded_val.val[j], *scale);
-        }
-        reinterpret_cast<stype*>(dst_)[idx] = store_val;
+        _dst[i] =
+            fp8::scaled_convert<scalar_t, cache_t, kv_dt>(_src[i], *scale);
      }
    }
-    // process tail
-    constexpr int32_t tail_cnt = ENTRY_SIZE % vec_size;
-    dst_ = dst_ + ENTRY_SIZE - tail_cnt;
-    src_ = src_ + ENTRY_SIZE - tail_cnt;
-#pragma unroll
-    for (int idx = threadIdx.x; idx < tail_cnt; idx += CTA_SIZE) {
-      if constexpr (kv_dt == Fp8KVCacheDataType::kAuto) {
-        dst_[idx] = static_cast<scalar_t>(src_[idx]);
-      } else {
-        dst_[idx] =
-            fp8::scaled_convert<scalar_t, cache_t, kv_dt>(src_[idx], *scale);
+  };
+
+  const auto loop_end =
+      std::min((int64_t)full_blocks_end, block_table_stride - offset);
+  for (int pid = split_start; pid < loop_end; ++pid) {
+    auto block_id = batch_block_table[pid];
+    auto block_start_ptr = src_cache + block_id * cache_block_stride;
+    auto block_dst_ptr = dst + pid * block_size * dst_entry_stride;
+    for (int eid = 0; eid < block_size; ++eid) {
+      copy_entry(block_start_ptr + eid * cache_entry_stride,
+                 block_dst_ptr + eid * dst_entry_stride);
+    }
+  }
+
+  if (partial_block_size) {
+    if (offset + full_blocks_end < block_table_stride) {
+      auto block_id = batch_block_table[full_blocks_end];
+      auto block_start_ptr = src_cache + block_id * cache_block_stride;
+      auto block_dst_ptr =
+          dst + full_blocks_end * block_size * dst_entry_stride;
+      for (int eid = 0; eid < partial_block_size; ++eid) {
+        copy_entry(block_start_ptr + eid * cache_entry_stride,
+                   block_dst_ptr + eid * dst_entry_stride);
      }
    }
  }
@@ -991,38 +1001,34 @@ __global__ void gather_and_maybe_dequant_cache(
 // SCALAR_T is the data type of the destination tensor.
 // CACHE_T is the stored data type of kv-cache.
 // KV_DTYPE is the real data type of kv-cache.
-#define CALL_GATHER_CACHE(SCALAR_T, CACHE_T, KV_DTYPE)                        \
-  vllm::gather_and_maybe_dequant_cache<SCALAR_T, CACHE_T, KV_DTYPE, 576,      \
-                                       thread_block_size>                     \
-      <<<grid, block, 0, stream>>>(                                           \
-          reinterpret_cast<CACHE_T*>(src_cache.data_ptr()),                   \
-          reinterpret_cast<SCALAR_T*>(dst.data_ptr()),                        \
-          block_table.data_ptr<int32_t>(), cu_seq_lens.data_ptr<int32_t>(),   \
-          token_to_seq.data_ptr<int32_t>(), num_tokens, block_size,           \
-          block_table_stride, cache_block_stride, cache_entry_stride,         \
-          dst_entry_stride, reinterpret_cast<const float*>(scale.data_ptr()), \
-          seq_starts_ptr);
+#define CALL_GATHER_CACHE(SCALAR_T, CACHE_T, KV_DTYPE)                      \
+  vllm::gather_and_maybe_dequant_cache<SCALAR_T, CACHE_T, KV_DTYPE>         \
+      <<<grid, block, 0, stream>>>(                                         \
+          reinterpret_cast<CACHE_T*>(src_cache.data_ptr()),                 \
+          reinterpret_cast<SCALAR_T*>(dst.data_ptr()),                      \
+          block_table.data_ptr<int32_t>(), cu_seq_lens.data_ptr<int32_t>(), \
+          block_size, entry_size, block_table_stride, cache_block_stride,   \
+          cache_entry_stride, dst_entry_stride,                             \
+          reinterpret_cast<const float*>(scale.data_ptr()), seq_starts_ptr);

 // Gather sequences from the cache into the destination tensor.
 //  - cu_seq_lens contains the cumulative sequence lengths for each batch
 //  - block_table contains the cache block indices for each sequence
-//  - token_to_seq contains the back mapping from token_id to batch_id
 //  - Optionally, seq_starts (if provided) offsets the starting block index by
 //  (seq_starts[bid] / page_size)
 void gather_and_maybe_dequant_cache(
-    torch::Tensor const& src_cache,     // [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...]
-    torch::Tensor const& dst,           // [TOT_TOKENS, ENTRIES...]
-    torch::Tensor const& block_table,   // [BATCH, BLOCK_INDICES]
-    torch::Tensor const& cu_seq_lens,   // [BATCH+1]
-    torch::Tensor const& token_to_seq,  // [MAX_TOKEN_ACROSS_CHUNKS]
-    int64_t num_tokens, const std::string& kv_cache_dtype,
+    torch::Tensor const& src_cache,    // [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...]
+    torch::Tensor const& dst,          // [TOT_TOKENS, ENTRIES...]
+    torch::Tensor const& block_table,  // [BATCH, BLOCK_INDICES]
+    torch::Tensor const& cu_seq_lens,  // [BATCH+1]
+    int64_t batch_size, const std::string& kv_cache_dtype,
    torch::Tensor const& scale,
    std::optional<torch::Tensor> seq_starts = std::nullopt) {
  at::cuda::OptionalCUDAGuard device_guard(src_cache.device());
  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();

  int32_t block_size = src_cache.size(1);
-  int32_t head_dim = dst.size(-1);
+  int32_t entry_size = src_cache.flatten(2, -1).size(2);

  TORCH_CHECK(block_table.dtype() == torch::kInt32,
              "block_table must be int32");
@@ -1032,9 +1038,6 @@ void gather_and_maybe_dequant_cache(
    TORCH_CHECK(seq_starts.value().dtype() == torch::kInt32,
                "seq_starts must be int32");
  }
-  TORCH_CHECK(head_dim == 576,
-              "gather_and_maybe_dequant_cache only support the head_dim to 576 "
-              "for better performance")

  TORCH_CHECK(src_cache.device() == dst.device(),
              "src_cache and dst must be on the same device");
@@ -1052,9 +1055,10 @@ void gather_and_maybe_dequant_cache(
  int64_t cache_entry_stride = src_cache.stride(1);
  int64_t dst_entry_stride = dst.stride(0);

-  constexpr int32_t thread_block_size = 64;
-  dim3 grid(num_tokens);
-  dim3 block(thread_block_size);
+  // 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(1024);

  const int32_t* seq_starts_ptr =
      seq_starts.has_value() ? seq_starts.value().data_ptr<int32_t>() : nullptr;
@@ -1063,82 +1067,6 @@ void gather_and_maybe_dequant_cache(
 }

 namespace vllm {
-
-// Gather and upconvert FP8 KV cache tokens to BF16 workspace
-// Similar to cp_gather_cache but specifically for FP8->BF16 conversion
-__global__ void cp_gather_and_upconvert_fp8_kv_cache(
-    const uint8_t* __restrict__ src_cache,    // [NUM_BLOCKS, BLOCK_SIZE, 656]
-    __nv_bfloat16* __restrict__ dst,          // [TOT_TOKENS, 576]
-    const int32_t* __restrict__ block_table,  // [BATCH, BLOCK_INDICES]
-    const int32_t* __restrict__ seq_lens,     // [BATCH]
-    const int32_t* __restrict__ workspace_starts,  // [BATCH]
-    const int32_t block_size, const int32_t head_dim,
-    const int64_t block_table_stride, const int64_t cache_block_stride,
-    const int64_t cache_entry_stride, const int64_t dst_entry_stride) {
-  const int64_t bid = blockIdx.x;  // Batch ID
-  const int32_t num_splits = gridDim.y;
-  const int32_t split = blockIdx.y;
-  const int32_t seq_start = workspace_starts[bid];
-  const int32_t seq_len = seq_lens[bid];
-  const int32_t tot_slots = seq_len;
-  const int32_t split_slots = cuda_utils::ceil_div(tot_slots, num_splits);
-
-  const int32_t split_start = split * split_slots;
-  const int32_t split_end = min((split + 1) * split_slots, tot_slots);
-
-  const bool is_active_split = (split_start < tot_slots);
-
-  if (!is_active_split) return;
-
-  // Adjust the pointer for the block_table for this batch
-  const int32_t batch_offset = bid * block_table_stride;
-  int32_t offset = split_start;
-  int32_t offset_div = offset / block_size;
-  offset = offset % block_size;
-  const int32_t* batch_block_table = block_table + batch_offset;
-
-  // Adjust dst pointer based on the cumulative sequence lengths
-  dst += seq_start * dst_entry_stride;
-
-  const int tid = threadIdx.x;
-
-  // Process each token in this split
-  for (int pid = split_start; pid < split_end; ++pid) {
-    auto block_id = batch_block_table[offset_div];
-    const uint8_t* token_ptr =
-        src_cache + block_id * cache_block_stride + offset * cache_entry_stride;
-    __nv_bfloat16* dst_ptr = dst + pid * dst_entry_stride;
-
-    // FP8 format: 512 bytes fp8 + 16 bytes scales + 128 bytes rope (64 bf16)
-    const uint8_t* no_pe_ptr = token_ptr;
-    const float* scales_ptr = reinterpret_cast<const float*>(token_ptr + 512);
-    const __nv_bfloat16* rope_ptr =
-        reinterpret_cast<const __nv_bfloat16*>(token_ptr + 512 + 16);
-
-    // Parallelize fp8 dequant (512 elements) and rope copy (64 elements)
-    if (tid < 512) {
-      // FP8 dequantization
-      const int tile = tid >> 7;  // each tile is 128 elements
-      const float scale = scales_ptr[tile];
-      const uint8_t val = no_pe_ptr[tid];
-      dst_ptr[tid] =
-          fp8::scaled_convert<__nv_bfloat16, uint8_t,
-                              vllm::Fp8KVCacheDataType::kFp8E4M3>(val, scale);
-    } else if (tid < 576) {
-      // Rope copy (64 bf16 elements)
-      const int rope_idx = tid - 512;
-      dst_ptr[512 + rope_idx] = rope_ptr[rope_idx];
-    }
-
-    // Move to next token
-    offset += 1;
-    if (offset == block_size) {
-      offset_div += 1;
-      offset = 0;
-    }
-  }
-}
-
 template <typename scalar_t>
 // Note(hc): The cp_gather_cache allows seq_starts to no longer be divisible by
 // block_size.
@@ -1280,57 +1208,6 @@ void cp_gather_cache(
  }
 }

-void cp_gather_and_upconvert_fp8_kv_cache(
-    torch::Tensor const& src_cache,         // [NUM_BLOCKS, BLOCK_SIZE, 656]
-    torch::Tensor const& dst,               // [TOT_TOKENS, 576]
-    torch::Tensor const& block_table,       // [BATCH, BLOCK_INDICES]
-    torch::Tensor const& seq_lens,          // [BATCH]
-    torch::Tensor const& workspace_starts,  // [BATCH]
-    int64_t batch_size) {
-  at::cuda::OptionalCUDAGuard device_guard(src_cache.device());
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-
-  int32_t block_size = src_cache.size(1);
-  int32_t head_dim = dst.size(1);
-
-  TORCH_CHECK(block_table.dtype() == torch::kInt32,
-              "block_table must be int32");
-  TORCH_CHECK(seq_lens.dtype() == torch::kInt32, "seq_lens must be int32");
-  TORCH_CHECK(workspace_starts.dtype() == torch::kInt32,
-              "workspace_starts must be int32");
-
-  TORCH_CHECK(src_cache.device() == dst.device(),
-              "src_cache and dst must be on the same device");
-  TORCH_CHECK(src_cache.device() == block_table.device(),
-              "src_cache and block_table must be on the same device");
-  TORCH_CHECK(src_cache.device() == seq_lens.device(),
-              "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");
-
-  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");
-
-  int64_t block_table_stride = block_table.stride(0);
-  int64_t cache_block_stride = src_cache.stride(0);
-  int64_t cache_entry_stride = src_cache.stride(1);
-  int64_t dst_entry_stride = dst.stride(0);
-
-  // 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_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,
-      dst_entry_stride);
-}
-
 // Macro to dispatch the kernel based on the data type.
 #define CALL_INDEXER_K_QUANT_AND_CACHE(KV_T, CACHE_T, KV_DTYPE)         \
  vllm::indexer_k_quant_and_cache_kernel<KV_T, CACHE_T, KV_DTYPE>       \
@@ -13,18 +13,6 @@
  #define AMX_DISPATCH(...) case cpu_attention::ISA::AMX:
 #endif

-#ifdef __aarch64__
-  #include "cpu_attn_neon.hpp"
-  #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;     \
@@ -53,7 +41,6 @@
  [&] {                                                                       \
    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,   \
@@ -86,8 +73,6 @@ torch::Tensor get_scheduler_metadata(
    isa = cpu_attention::ISA::VEC;
  } else if (isa_hint == "vec16") {
    isa = cpu_attention::ISA::VEC16;
-  } else if (isa_hint == "neon") {
-    isa = cpu_attention::ISA::NEON;
  } else {
    TORCH_CHECK(false, "Unsupported CPU attention ISA hint: " + isa_hint);
  }
@@ -117,6 +102,7 @@ torch::Tensor get_scheduler_metadata(
  input.casual = casual;
  input.isa = isa;
  input.enable_kv_split = enable_kv_split;
+  TORCH_CHECK(casual, "Only supports casual mask for now.");

  VLLM_DISPATCH_FLOATING_TYPES(dtype, "get_scheduler_metadata", [&]() {
    CPU_ATTN_DISPATCH_CASE_HEADDIM(head_dim, [&] {
@@ -172,8 +158,6 @@ void cpu_attn_reshape_and_cache(
      return cpu_attention::ISA::VEC;
    } else if (isa == "vec16") {
      return cpu_attention::ISA::VEC16;
-    } else if (isa == "neon") {
-      return cpu_attention::ISA::NEON;
    } else {
      TORCH_CHECK(false, "Invalid ISA type: " + isa);
    }
@@ -14,7 +14,7 @@
 #include "utils.hpp"

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

 template <ISA isa, typename scalar_t, int64_t head_dim>
 class AttentionImpl {};
@@ -143,12 +143,6 @@ struct AttentionMetadata {
      case ISA::VEC:
        ss << "VEC, ";
        break;
-      case ISA::VEC16:
-        ss << "VEC16, ";
-        break;
-      case ISA::NEON:
-        ss << "NEON, ";
-        break;
    }
    ss << "workitem_group_num: " << workitem_group_num
       << ", reduction_item_num: " << reduction_item_num
@@ -186,7 +180,7 @@ struct AttentionMetadata {
 //  - Intermediate outputs: q_tile_size * head_dim * output_buffer_elem_size + 2
 //  * q_tile_size * 4, partial output, max + sum (float)
 // Reduction scratchpad contains:
-//  - flags: bool array to indicate whether the split is finished
+//  - flags: bool array to indicate wether the split is finished
 //  - outputs: split_num * q_tile_size * head_dim * output_buffer_elem_size
 //  - max, sum: 2 * split_num * q_tile_size * 4
 class AttentionScratchPad {
@@ -847,7 +841,7 @@ struct VecTypeTrait<c10::BFloat16> {
 };
 #endif

-#if !defined(__powerpc__) && !defined(__s390x__)
+#if !defined(__powerpc__)
 template <>
 struct VecTypeTrait<c10::Half> {
  using vec_t = vec_op::FP16Vec16;
@@ -1246,8 +1240,14 @@ class AttentionMainLoop {
        // rescale sum and partial outputs
        if (need_rescale) {
          // compute rescale factor
+#ifdef DEFINE_FAST_EXP
+          vec_op::FP32Vec16 rescale_factor_vec(rescale_factor);
+          rescale_factor_vec = fast_exp(rescale_factor_vec);
+          rescale_factor = rescale_factor_vec.get_last_elem();
+#else
          rescale_factor = std::exp(rescale_factor);
          vec_op::FP32Vec16 rescale_factor_vec(rescale_factor);
+#endif

          // rescale sum
          new_sum_val += rescale_factor * init_sum_val;
@@ -1883,8 +1883,15 @@ class AttentionMainLoop {
                                   : curr_output_buffer;
          float rescale_factor = final_max > curr_max ? curr_max - final_max
                                                      : final_max - curr_max;
+
+#ifdef DEFINE_FAST_EXP
+          vec_op::FP32Vec16 rescale_factor_vec(rescale_factor);
+          rescale_factor_vec = fast_exp(rescale_factor_vec);
+          rescale_factor = rescale_factor_vec.get_last_elem();
+#else
          rescale_factor = std::exp(rescale_factor);
          vec_op::FP32Vec16 rescale_factor_vec(rescale_factor);
+#endif

          local_sum[head_idx] = final_max > curr_max
                                    ? final_sum + rescale_factor * curr_sum
@@ -60,54 +60,4 @@

 #endif

-#ifdef __aarch64__
-  // Implementation copied from Arm Optimized Routines (expf AdvSIMD)
-  // https://github.com/ARM-software/optimized-routines/blob/master/math/aarch64/advsimd/expf.c
-  #include <limits>
-  #define DEFINE_FAST_EXP                                                      \
-    const float32x4_t inv_ln2 = vdupq_n_f32(0x1.715476p+0f);                   \
-    const float ln2_hi = 0x1.62e4p-1f;                                         \
-    const float ln2_lo = 0x1.7f7d1cp-20f;                                      \
-    const float c0 = 0x1.0e4020p-7f;                                           \
-    const float c2 = 0x1.555e66p-3f;                                           \
-    const float32x4_t ln2_c02 = {ln2_hi, ln2_lo, c0, c2};                      \
-    const uint32x4_t exponent_bias = vdupq_n_u32(0x3f800000);                  \
-    const float32x4_t c1 = vdupq_n_f32(0x1.573e2ep-5f);                        \
-    const float32x4_t c3 = vdupq_n_f32(0x1.fffdb6p-2f);                        \
-    const float32x4_t c4 = vdupq_n_f32(0x1.ffffecp-1f);                        \
-    const float32x4_t pos_special_bound = vdupq_n_f32(0x1.5d5e2ap+6f);         \
-    const float32x4_t neg_special_bound = vnegq_f32(pos_special_bound);        \
-    const float32x4_t inf =                                                    \
-        vdupq_n_f32(std::numeric_limits<float>::infinity());                   \
-    const float32x4_t zero = vdupq_n_f32(0.0f);                                \
-    auto neon_expf = [&](float32x4_t values) __attribute__((always_inline)) {  \
-      float32x4_t n = vrndaq_f32(vmulq_f32(values, inv_ln2));                  \
-      float32x4_t r = vfmsq_laneq_f32(values, n, ln2_c02, 0);                  \
-      r = vfmsq_laneq_f32(r, n, ln2_c02, 1);                                   \
-      uint32x4_t e = vshlq_n_u32(vreinterpretq_u32_s32(vcvtq_s32_f32(n)), 23); \
-      float32x4_t scale = vreinterpretq_f32_u32(vaddq_u32(e, exponent_bias));  \
-      float32x4_t r2 = vmulq_f32(r, r);                                        \
-      float32x4_t p = vfmaq_laneq_f32(c1, r, ln2_c02, 2);                      \
-      float32x4_t q = vfmaq_laneq_f32(c3, r, ln2_c02, 3);                      \
-      q = vfmaq_f32(q, p, r2);                                                 \
-      p = vmulq_f32(c4, r);                                                    \
-      float32x4_t poly = vfmaq_f32(p, q, r2);                                  \
-      poly = vfmaq_f32(scale, poly, scale);                                    \
-      const uint32x4_t hi_mask = vcgeq_f32(values, pos_special_bound);         \
-      const uint32x4_t lo_mask = vcleq_f32(values, neg_special_bound);         \
-      poly = vbslq_f32(hi_mask, inf, poly);                                    \
-      return vbslq_f32(lo_mask, zero, poly);                                   \
-    };                                                                         \
-    auto fast_exp = [&](vec_op::FP32Vec16& vec)                                \
-                        __attribute__((always_inline)) {                       \
-                          float32x4x4_t result;                                \
-                          result.val[0] = neon_expf(vec.reg.val[0]);           \
-                          result.val[1] = neon_expf(vec.reg.val[1]);           \
-                          result.val[2] = neon_expf(vec.reg.val[2]);           \
-                          result.val[3] = neon_expf(vec.reg.val[3]);           \
-                          return vec_op::FP32Vec16(result);                    \
-                        };
-
-#endif  // __aarch64__
-
 #endif
@@ -1,386 +0,0 @@
-#ifndef CPU_ATTN_NEON_HPP
-#define CPU_ATTN_NEON_HPP
-
-#include "cpu_attn_impl.hpp"
-#include <arm_neon.h>
-#include <type_traits>
-namespace cpu_attention {
-
-namespace {
-
-#define BLOCK_SIZE_ALIGNMENT 32
-#define HEAD_SIZE_ALIGNMENT 32
-#define MAX_Q_HEAD_NUM_PER_ITER 16
-
-// These do not use vectorized class for loading / converting
-// because csrc/cpu/cpu_types_arm.hpp does not have fallback options
-// for vec_op::BF16Vec* / vec_op::BF16Vec* on Arm HW that
-// doesn't support BF16.
-// We don't use vec_op::FP32Vec* or vec_op::FP16Vec* for consistency.
-template <typename kv_cache_t>
-FORCE_INLINE void load_row8_B_as_f32(const kv_cache_t* p, float32x4_t& b0,
-                                     float32x4_t& b1);
-
-template <>
-FORCE_INLINE void load_row8_B_as_f32<float>(const float* p, float32x4_t& b0,
-                                            float32x4_t& b1) {
-  b0 = vld1q_f32(p + 0);
-  b1 = vld1q_f32(p + 4);
-}
-
-template <>
-FORCE_INLINE void load_row8_B_as_f32<c10::Half>(const c10::Half* p,
-                                                float32x4_t& b0,
-                                                float32x4_t& b1) {
-  const float16_t* h = reinterpret_cast<const float16_t*>(p);
-  float16x8_t v = vld1q_f16(h);
-  b0 = vcvt_f32_f16(vget_low_f16(v));
-  b1 = vcvt_f32_f16(vget_high_f16(v));
-}
-
-template <>
-FORCE_INLINE void load_row8_B_as_f32<c10::BFloat16>(const c10::BFloat16* p,
-                                                    float32x4_t& b0,
-                                                    float32x4_t& b1) {
-  const uint16_t* u = reinterpret_cast<const uint16_t*>(p);
-#ifdef ARM_BF16_SUPPORT
-  uint16x8_t u0 = vld1q_u16(u);
-  bfloat16x8_t bf0 = vreinterpretq_bf16_u16(u0);
-  b0 = vcvtq_low_f32_bf16(bf0);
-  b1 = vcvtq_high_f32_bf16(bf0);
-#else
-  uint16x8_t x0 = vld1q_u16(u);
-  uint32x4_t lo = vshlq_n_u32(vmovl_u16(vget_low_u16(x0)), 16);
-  uint32x4_t hi = vshlq_n_u32(vmovl_u16(vget_high_u16(x0)), 16);
-  b0 = vreinterpretq_f32_u32(lo);
-  b1 = vreinterpretq_f32_u32(hi);
-#endif
-}
-
-// 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
-template <int32_t M, typename kv_cache_t>
-FORCE_INLINE void gemm_micro_neon_fmla_Mx8_Ku4(
-    const float* __restrict A,       // [M x K],
-    const kv_cache_t* __restrict B,  // [K x 8],
-    float* __restrict C,             // [M x 8],
-    int64_t lda, int64_t ldb, int64_t ldc, int32_t K, bool accumulate) {
-  // kernel supports max M of 8, as it'd spill for larger M
-  static_assert(1 <= M && M <= 8, "M must be in [1,8]");
-
-// helpers for per-M codegen
-#define ROWS_APPLY(OP) OP(0) OP(1) OP(2) OP(3) OP(4) OP(5) OP(6) OP(7)
-#define IF_M(i) if constexpr (M > (i))
-
-  // A row base pointers
-#define DECL_A(i) const float* a##i = A + (i) * lda;
-  ROWS_APPLY(DECL_A)
-#undef DECL_A
-
-  // declare 2 accumulators per row of M
-#define DECL_ACC(i) float32x4_t acc##i##_0, acc##i##_1;
-  ROWS_APPLY(DECL_ACC)
-#undef DECL_ACC
-
-  // initialize accumulators
-#define INIT_ACC(i)                              \
-  IF_M(i) {                                      \
-    if (accumulate) {                            \
-      acc##i##_0 = vld1q_f32(C + (i) * ldc + 0); \
-      acc##i##_1 = vld1q_f32(C + (i) * ldc + 4); \
-    } else {                                     \
-      acc##i##_0 = vdupq_n_f32(0.f);             \
-      acc##i##_1 = vdupq_n_f32(0.f);             \
-    }                                            \
-  }
-  ROWS_APPLY(INIT_ACC)
-#undef INIT_ACC
-
-  int32_t k = 0;
-
-  // K unrolled by 4
-  for (; k + 3 < K; k += 4) {
-    // load A[k..k+3] for each active row (M)
-#define LOAD_A4(i)     \
-  float32x4_t a##i##v; \
-  IF_M(i) a##i##v = vld1q_f32(a##i + k);
-    ROWS_APPLY(LOAD_A4)
-#undef LOAD_A4
-
-    // helper: FMA lane L from aiv
-#define FMAS_LANE(i, aiv, L)                              \
-  IF_M(i) {                                               \
-    acc##i##_0 = vfmaq_laneq_f32(acc##i##_0, b0, aiv, L); \
-    acc##i##_1 = vfmaq_laneq_f32(acc##i##_1, b1, aiv, L); \
-  }
-
-    // k + 0
-    {
-      float32x4_t b0, b1;
-      load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 0) * ldb, b0, b1);
-#define STEP_K0(i) FMAS_LANE(i, a##i##v, 0)
-      ROWS_APPLY(STEP_K0)
-#undef STEP_K0
-    }
-    // k + 1
-    {
-      float32x4_t b0, b1;
-      load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 1) * ldb, b0, b1);
-#define STEP_K1(i) FMAS_LANE(i, a##i##v, 1)
-      ROWS_APPLY(STEP_K1)
-#undef STEP_K1
-    }
-    // k + 2
-    {
-      float32x4_t b0, b1;
-      load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 2) * ldb, b0, b1);
-#define STEP_K2(i) FMAS_LANE(i, a##i##v, 2)
-      ROWS_APPLY(STEP_K2)
-#undef STEP_K2
-    }
-    // k + 3
-    {
-      float32x4_t b0, b1;
-      load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 3) * ldb, b0, b1);
-#define STEP_K3(i) FMAS_LANE(i, a##i##v, 3)
-      ROWS_APPLY(STEP_K3)
-#undef STEP_K3
-    }
-#undef FMAS_LANE
-  }
-
-  // K tail
-  for (; k < K; ++k) {
-    float32x4_t b0, b1;
-    load_row8_B_as_f32<kv_cache_t>(B + (int64_t)k * ldb, b0, b1);
-#define TAIL_ROW(i)                             \
-  IF_M(i) {                                     \
-    float32x4_t ai = vdupq_n_f32(*(a##i + k));  \
-    acc##i##_0 = vfmaq_f32(acc##i##_0, b0, ai); \
-    acc##i##_1 = vfmaq_f32(acc##i##_1, b1, ai); \
-  }
-    ROWS_APPLY(TAIL_ROW)
-#undef TAIL_ROW
-  }
-
-  // store accumulators to C
-#define STORE_ROW(i)                          \
-  IF_M(i) {                                   \
-    vst1q_f32(C + (i) * ldc + 0, acc##i##_0); \
-    vst1q_f32(C + (i) * ldc + 4, acc##i##_1); \
-  }
-  ROWS_APPLY(STORE_ROW)
-#undef STORE_ROW
-
-#undef ROWS_APPLY
-#undef IF_M
-}
-
-template <int32_t N, typename kv_cache_t>
-FORCE_INLINE void gemm_macro_neon_fmla_Mx8_Ku4(const float* __restrict A,
-                                               const kv_cache_t* __restrict B,
-                                               float* __restrict C, int32_t M,
-                                               int32_t K, int64_t lda,
-                                               int64_t ldb, int64_t ldc,
-                                               bool accumulate) {
-  // micro kernel is Mx8
-  static_assert(N % 8 == 0, "N must be a multiple of 8");
-  for (int32_t m = 0; m < M;) {
-    int32_t mb = (M - m >= 8) ? 8 : (M - m >= 4) ? 4 : (M - m >= 2) ? 2 : 1;
-    const float* Ab = A + m * lda;
-    float* Cb = C + m * ldc;
-
-    for (int32_t n = 0; n < N; n += 8) {
-      const kv_cache_t* Bn = B + n;
-      float* Cn = Cb + n;
-      switch (mb) {
-        case 8:
-          gemm_micro_neon_fmla_Mx8_Ku4<8, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
-                                                      K, accumulate);
-          break;
-        case 4:
-          gemm_micro_neon_fmla_Mx8_Ku4<4, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
-                                                      K, accumulate);
-          break;
-        case 2:
-          gemm_micro_neon_fmla_Mx8_Ku4<2, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
-                                                      K, accumulate);
-          break;
-        default:
-          gemm_micro_neon_fmla_Mx8_Ku4<1, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
-                                                      K, accumulate);
-          break;
-      }
-    }
-    // no tail loop for N as it's guaranteed to be a multiple of 8
-    m += mb;
-  }
-}
-
-template <typename kv_cache_t>
-class TileGemmNeonFMLA {
- public:
-  template <AttentionGemmPhase phase, int32_t k_size>
-  FORCE_INLINE static void gemm(const int32_t m_size,
-                                float* __restrict__ a_tile,
-                                kv_cache_t* __restrict__ b_tile,
-                                float* __restrict__ c_tile, const int64_t lda,
-                                const int64_t ldb, const int64_t ldc,
-                                const int32_t block_size,
-                                const int32_t dynamic_k_size,
-                                const bool accum_c) {
-    if constexpr (phase == AttentionGemmPhase::QK) {
-      gemm_macro_neon_fmla_Mx8_Ku4<BLOCK_SIZE_ALIGNMENT, kv_cache_t>(
-          a_tile, b_tile, c_tile, m_size, k_size, lda, ldb, ldc, accum_c);
-    } else {
-      gemm_macro_neon_fmla_Mx8_Ku4<HEAD_SIZE_ALIGNMENT, kv_cache_t>(
-          a_tile, b_tile, c_tile, m_size, dynamic_k_size, lda, ldb, ldc,
-          accum_c);
-    }
-  }
-};
-
-}  // namespace
-
-// this is similar to "ISA::VEC" at the moment
-template <typename scalar_t, int64_t head_dim>
-class AttentionImpl<ISA::NEON, scalar_t, head_dim> {
- public:
-  using query_t = scalar_t;
-  using q_buffer_t = float;
-  using kv_cache_t = scalar_t;
-  using logits_buffer_t = float;
-  using partial_output_buffer_t = float;
-  using prob_buffer_t = float;
-
-  constexpr static int64_t BlockSizeAlignment =
-      BLOCK_SIZE_ALIGNMENT;  // KV token num unit of QK and PV phases
-  constexpr static int64_t HeadDimAlignment =
-      HEAD_SIZE_ALIGNMENT;  // headdim num unit of PV phase
-  constexpr static int64_t MaxQHeadNumPerIteration = MAX_Q_HEAD_NUM_PER_ITER;
-  constexpr static int64_t HeadDim = 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);
-  // the gemm micro kernel is Mx8
-  static_assert(HeadDimAlignment % 8 == 0);
-  static_assert(BlockSizeAlignment % 8 == 0);
-
- public:
-  template <template <typename tile_gemm_t> typename attention>
-  FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) {
-    attention<TileGemmNeonFMLA<kv_cache_t>> attention_iteration;
-    attention_iteration(CPU_ATTENTION_PARAMS);
-  }
-
-  // k_cache_token_group_stride: stride of K cache when move to next
-  // BlockSizeAlignment tokens in a block
-  constexpr static int64_t k_cache_token_group_stride(
-      const int32_t block_size) {
-    return BlockSizeAlignment;  // layout of k_cache block is [head_dim,
-                                // block_size], row-major
-  }
-
-  // v_cache_token_group_stride: stride of V cache when move to next
-  // BlockSizeAlignment tokens in a block
-  constexpr static int64_t v_cache_token_group_stride(
-      const int32_t block_size) {
-    return head_dim * BlockSizeAlignment;  // layout of v_cache is [block_size,
-                                           // head_dim], row-major
-  }
-
-  // v_cache_head_group_stride: stride of V cache when move to next
-  // HeadDimAlignment head dims in a block
-  constexpr static int64_t v_cache_head_group_stride(const int32_t block_size) {
-    return HeadDimAlignment;  // layout of v_cache is [block_size, head_dim],
-                              // row-major
-  }
-
-  // Copy q to q_buffer and cast it to fp32
-  static void copy_q_heads_tile(
-      scalar_t* __restrict__ src,  // [q_num, q_heads_per_kv, head_size]
-      float* __restrict__ q_buffer, const int32_t q_num,
-      const int32_t q_heads_per_kv, const int64_t q_num_stride,
-      const int64_t q_head_stride, float scale) {
-    static_assert(head_dim % 16 == 0);
-    constexpr int32_t unroll_size = head_dim / 16;
-    using load_vec_t = typename VecTypeTrait<scalar_t>::vec_t;
-
-    vec_op::FP32Vec16 scale_vec(scale);
-    for (int32_t q_num_idx = 0; q_num_idx < q_num; ++q_num_idx) {
-      for (int32_t q_head_idx = 0; q_head_idx < q_heads_per_kv; ++q_head_idx) {
-        scalar_t* __restrict__ curr_q =
-            src + q_num_idx * q_num_stride + q_head_idx * q_head_stride;
-        float* __restrict__ curr_q_buffer =
-            q_buffer + q_num_idx * q_heads_per_kv * head_dim +
-            q_head_idx * head_dim;
-
-        vec_op::unroll_loop<int32_t, unroll_size>([&](int32_t i) {
-          load_vec_t vec(curr_q);
-          vec_op::FP32Vec16 fp32_vec(vec);
-          fp32_vec = fp32_vec * scale_vec;
-          fp32_vec.save(curr_q_buffer);
-
-          curr_q += 16;
-          curr_q_buffer += 16;
-        });
-      }
-    }
-  }
-
-  // reshape K as column-major and V as row-major
-  static void reshape_and_cache(
-      const scalar_t* __restrict__ key, const scalar_t* __restrict__ value,
-      scalar_t* __restrict__ key_cache, scalar_t* __restrict__ value_cache,
-      const int64_t* __restrict__ slot_mapping, const int64_t token_num,
-      const int64_t key_token_num_stride, const int64_t value_token_num_stride,
-      const int64_t head_num, const int64_t key_head_num_stride,
-      const int64_t value_head_num_stride, const int64_t num_blocks,
-      const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
-      const int64_t block_size, const int64_t block_size_stride) {
-#pragma omp parallel for collapse(2)
-    for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
-      for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
-        const int64_t pos = slot_mapping[token_idx];
-        if (pos < 0) {
-          // skip
-          continue;
-        }
-
-        const int64_t block_idx = pos / block_size;
-        const int64_t block_offset = pos % block_size;
-        {
-          // Write Key
-          const scalar_t* key_start_ptr = key +
-                                          token_idx * key_token_num_stride +
-                                          head_idx * key_head_num_stride;
-          scalar_t* key_cache_start_ptr =
-              key_cache + block_idx * num_blocks_stride +
-              head_idx * cache_head_num_stride + block_offset;
-
-#pragma GCC unroll 8
-          for (int64_t i = 0, j = 0; i < head_dim; ++i, j += block_size) {
-            key_cache_start_ptr[j] = key_start_ptr[i];
-          }
-        }
-        {
-          // Write Value
-          const scalar_t* value_start_ptr = value +
-                                            token_idx * value_token_num_stride +
-                                            head_idx * value_head_num_stride;
-          scalar_t* value_cache_start_ptr =
-              value_cache + block_idx * num_blocks_stride +
-              head_idx * cache_head_num_stride + block_offset * head_dim;
-          std::memcpy(value_cache_start_ptr, value_start_ptr,
-                      sizeof(scalar_t) * head_dim);
-        }
-      }
-    }
-  }
-};
-}  // namespace cpu_attention
-
-#endif  // #ifndef CPU_ATTN_NEON_HPP
@@ -4,7 +4,6 @@

 #include <vecintrin.h>
 #include <cmath>
-#include <limits>
 #include <torch/all.h>
 namespace vec_op {

@@ -175,9 +174,8 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
  }

  explicit FP32Vec8(const BF16Vec8& v) {
-    // On big-endian s390x, place BF16 first to get correct byte order
-    reg.val[0] = (__vector float)vec_mergeh(v.reg, zero);
-    reg.val[1] = (__vector float)vec_mergel(v.reg, zero);
+    reg.val[0] = (__vector float)vec_mergeh(zero, v.reg);
+    reg.val[1] = (__vector float)vec_mergel(zero, v.reg);
  }

  float reduce_sum() const {
@@ -191,257 +189,51 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
  }

  FP32Vec8 exp() const {
-    f32x4x2_t out;
-
-    const __vector float log2e = vec_splats(1.44269504088896341f);
-    const __vector float one = vec_splats(1.0f);
-    const __vector float min_x = vec_splats(-87.3f);
-    const __vector float max_x = vec_splats(88.7f);
-
-    // 5th-degree minimax polynomial for 2^r (r in [0,1))
-    const __vector float c1 = vec_splats(0.6931471805599453f);
-    const __vector float c2 = vec_splats(0.240226506959101f);
-    const __vector float c3 = vec_splats(0.05550410866482158f);
-    const __vector float c4 = vec_splats(0.009618129107628477f);
-    const __vector float c5 = vec_splats(0.0013333558146428443f);
-
-    for (int i = 0; i < 2; i++) {
-      __vector float x = reg.val[i];
-
-      x = vec_max(x, min_x);
-      x = vec_min(x, max_x);
-
-      __vector float y = vec_mul(x, log2e);
-
-      __vector float kf = vec_floor(y);
-      __vector float r = vec_sub(y, kf);
-
-      __vector signed int k = vec_signed(kf);
-      const __vector signed int min_k = vec_splats((signed int)-126);
-      const __vector signed int max_k = vec_splats((signed int)127);
-      k = vec_min(vec_max(k, min_k), max_k);
-
-      // Build 2^k from exponent bits
-      __vector signed int exp_int = vec_add(k, vec_splats((signed int)127));
-      __vector unsigned int bits = (__vector unsigned int)exp_int;
-      bits = vec_sl(bits, vec_splats((unsigned int)23));
-      __vector float pow2k = (__vector float)bits;
-
-      // Improved minimax polynomial
-      __vector float poly = vec_madd(c5, r, c4);
-      poly = vec_madd(poly, r, c3);
-      poly = vec_madd(poly, r, c2);
-      poly = vec_madd(poly, r, c1);
-      poly = vec_madd(poly, r, one);
-
-      out.val[i] = vec_mul(pow2k, poly);
-    }
-
-    return FP32Vec8(out);
+    // TODO: Vectorize this
+    AliasReg ar;
+    ar.reg = reg;
+    f32x4x4_t ret;
+    ret.val[0][0] = std::exp(ar.values[0]);
+    ret.val[0][1] = std::exp(ar.values[1]);
+    ret.val[0][2] = std::exp(ar.values[2]);
+    ret.val[0][3] = std::exp(ar.values[3]);
+    ret.val[1][0] = std::exp(ar.values[4]);
+    ret.val[1][1] = std::exp(ar.values[5]);
+    ret.val[1][2] = std::exp(ar.values[6]);
+    ret.val[1][3] = std::exp(ar.values[7]);
+    return FP32Vec8(f32x4x2_t({ret.val[0], ret.val[1]}));
  }

  FP32Vec8 tanh() const {
-    // tanh(x) = (exp(2x) - 1) / (exp(2x) + 1)
-    const __vector float one = vec_splats(1.0f);
-    const __vector float two = vec_splats(2.0f);
-    const __vector float zero = vec_splats(0.0f);
-    const __vector float sat =
-        vec_splats(9.0f);  // beyond this, tanh(x) ~ sign(x)
-
-    f32x4x2_t out;
-
-    for (int i = 0; i < 2; i++) {
-      __vector float x = reg.val[i];
-      __vector float ax = vec_abs(x);
-
-      // sign(x): +1 or -1
-      __vector float sign = vec_sel(vec_splats(-1.0f), one, vec_cmpgt(x, zero));
-
-      // saturation mask: |x| > sat
-      __vector __bool int saturated = vec_cmpgt(ax, sat);
-
-      // 2x
-      __vector float two_x = vec_mul(x, two);
-
-      // Build a temporary FP32Vec8 with both lanes = 2x, reuse exp()
-      f32x4x2_t tmp;
-      tmp.val[0] = two_x;
-      tmp.val[1] = two_x;
-      FP32Vec8 exp_2x_vec(tmp);
-
-      FP32Vec8 e2x = exp_2x_vec.exp();
-      __vector float e = e2x.reg.val[i];
-
-      // tanh(x) = (e - 1) / (e + 1)
-      __vector float num = vec_sub(e, one);
-      __vector float den = vec_add(e, one);
-
-      __vector float t = vec_div(num, den);
-
-      // For large |x|, clamp to sign(x)
-      out.val[i] = vec_sel(t, sign, saturated);
-    }
-
-    return FP32Vec8(out);
+    // TODO: Vectorize this
+    AliasReg ar;
+    ar.reg = reg;
+    f32x4x4_t ret;
+    ret.val[0][0] = std::tanh(ar.values[0]);
+    ret.val[0][1] = std::tanh(ar.values[1]);
+    ret.val[0][2] = std::tanh(ar.values[2]);
+    ret.val[0][3] = std::tanh(ar.values[3]);
+    ret.val[1][0] = std::tanh(ar.values[4]);
+    ret.val[1][1] = std::tanh(ar.values[5]);
+    ret.val[1][2] = std::tanh(ar.values[6]);
+    ret.val[1][3] = std::tanh(ar.values[7]);
+    return FP32Vec8(f32x4x2_t({ret.val[0], ret.val[1]}));
  }

  FP32Vec8 er() const {
-    // A&S 7.1.26 approximation:
-    // erf(x) = sign(x) * (1 - ((((a5*t + a4)*t + a3)*t + a2)*t + a1) * t *
-    // exp(-x^2)) t = 1 / (1 + p*|x|),  p = 0.3275911
-
-    const __vector float one = vec_splats(1.0f);
-    const __vector float zero = vec_splats(0.0f);
-    const __vector float p = vec_splats(0.3275911f);
-
-    // Polynomial coeffs
-    const __vector float a1 = vec_splats(0.254829592f);
-    const __vector float a2 = vec_splats(-0.284496736f);
-    const __vector float a3 = vec_splats(1.421413741f);
-    const __vector float a4 = vec_splats(-1.453152027f);
-    const __vector float a5 = vec_splats(1.061405429f);
-
-    // Threshold where erf(x) ~ sign(x)
-    const __vector float sat = vec_splats(6.0f);
-
-    f32x4x2_t out;
-
-    for (int lane = 0; lane < 2; lane++) {
-      __vector float x = reg.val[lane];
-      __vector float ax = vec_abs(x);
-
-      // sign(x)
-      __vector float sign = vec_sel(vec_splats(-1.0f), one, vec_cmpgt(x, zero));
-
-      // |x| > 6 → erf(x) = ±1
-      __vector __bool int saturated = vec_cmpgt(ax, sat);
-
-      // t = 1 / (1 + p * |x|)
-      __vector float t = vec_madd(p, ax, one);
-      t = vec_div(one, t);
-
-      // poly = a5
-      __vector float poly = a5;
-      poly = vec_madd(poly, t, a4);
-      poly = vec_madd(poly, t, a3);
-      poly = vec_madd(poly, t, a2);
-      poly = vec_madd(poly, t, a1);
-
-      // full polynomial: poly = poly * t
-      poly = vec_mul(poly, t);
-
-      // Compute exp(-x^2)
-      __vector float x2 = vec_mul(x, x);
-      __vector float neg_x2 = vec_neg(x2);
-
-      f32x4x2_t tmp;
-      tmp.val[0] = neg_x2;
-      tmp.val[1] = neg_x2;
-      FP32Vec8 exp_neg_x2(tmp);
-
-      FP32Vec8 e = exp_neg_x2.exp();
-      __vector float ex = e.reg.val[lane];
-
-      // erf(x) = sign * (1 - poly * exp(-x^2))
-      __vector float term = vec_mul(poly, ex);
-      __vector float y = vec_sub(one, term);
-      y = vec_mul(y, sign);
-
-      // saturated → ±1
-      __vector float sat_val = vec_mul(sign, one);
-      out.val[lane] = vec_sel(y, sat_val, saturated);
-    }
-
-    return FP32Vec8(out);
-  }
-  // Elementwise sigmoid(x) = 1 / (1 + exp(-x))
-  FP32Vec8 sigmoid() const {
-    const __vector float one = vec_splats(1.0f);
-
-    f32x4x2_t neg;
-    for (int i = 0; i < 2; ++i) {
-      neg.val[i] = vec_neg(reg.val[i]);
-    }
-
-    FP32Vec8 neg_x(neg);
-    FP32Vec8 e = neg_x.exp();  // exp(-x)
-
-    f32x4x2_t denom;
-    for (int i = 0; i < 2; ++i) {
-      denom.val[i] = vec_add(one, e.reg.val[i]);
-    }
-
-    FP32Vec8 denom_vec(denom);
-    FP32Vec8 one_vec(1.0f);
-
-    return one_vec / denom_vec;
-  }
-
-  // Tanh-based GELU:
-  // gelu(x) = 0.5 * x * (1 + tanh(√(2/π) * (x + 0.044715 * x^3)))
-  FP32Vec8 gelu_tanh() const {
-    const __vector float k_s2pi = vec_splats(0.7978845608028654f);  // √(2/π)
-    const __vector float k_0_0447 = vec_splats(0.044715f);
-
-    f32x4x2_t x2, x3, inner;
-    for (int i = 0; i < 2; ++i) {
-      __vector float x = reg.val[i];
-      x2.val[i] = vec_mul(x, x);                            // x^2
-      x3.val[i] = vec_mul(x2.val[i], x);                    // x^3
-      __vector float t = vec_madd(k_0_0447, x3.val[i], x);  // x + 0.044715*x^3
-      inner.val[i] = vec_mul(k_s2pi, t);                    // √(2/π)*(...)
-    }
-
-    FP32Vec8 inner_vec(inner);
-    FP32Vec8 t = inner_vec.tanh();  // tanh part
-
-    FP32Vec8 one_vec(1.0f);
-    FP32Vec8 half_vec(0.5f);
-
-    FP32Vec8 x_vec(*this);
-    return x_vec * half_vec * (one_vec + t);
-  }
-
-  // Erf-based GELU:
-  // gelu(x) = 0.5 * x * (1 + erf(x / √2))
-  FP32Vec8 gelu_erf() const {
-    const __vector float inv_sqrt2 = vec_splats(0.7071067811865476f);  // 1/√2
-    FP32Vec8 x_vec(*this);
-
-    f32x4x2_t scaled;
-    for (int i = 0; i < 2; ++i) {
-      scaled.val[i] = vec_mul(reg.val[i], inv_sqrt2);
-    }
-    FP32Vec8 x_scaled(scaled);
-
-    FP32Vec8 erf_x = x_scaled.er();
-
-    FP32Vec8 one_vec(1.0f);
-    FP32Vec8 half_vec(0.5f);
-
-    return x_vec * half_vec * (one_vec + erf_x);
-  }
-
-  // Elementwise reciprocal: 1/x (scalar per lane, for correctness)
-  FP32Vec8 rcp() const {
-    AliasReg in, out;
-    in.reg = reg;
-
-    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
-      out.values[i] = 1.0f / in.values[i];
-    }
-    return FP32Vec8(out.reg);
-  }
-
-  // Elementwise rsqrt(x) = 1 / sqrt(x) (scalar per lane, for correctness)
-  FP32Vec8 rsqrt() const {
-    AliasReg in, out;
-    in.reg = reg;
-
-    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
-      out.values[i] = 1.0f / std::sqrt(in.values[i]);
-    }
-    return FP32Vec8(out.reg);
+    // TODO: Vectorize this
+    AliasReg ar;
+    ar.reg = reg;
+    f32x4x4_t ret;
+    ret.val[0][0] = std::erf(ar.values[0]);
+    ret.val[0][1] = std::erf(ar.values[1]);
+    ret.val[0][2] = std::erf(ar.values[2]);
+    ret.val[0][3] = std::erf(ar.values[3]);
+    ret.val[1][0] = std::erf(ar.values[4]);
+    ret.val[1][1] = std::erf(ar.values[5]);
+    ret.val[1][2] = std::erf(ar.values[6]);
+    ret.val[1][3] = std::erf(ar.values[7]);
+    return FP32Vec8(f32x4x2_t({ret.val[0], ret.val[1]}));
  }

  FP32Vec8 operator*(const FP32Vec8& b) const {
@@ -524,11 +316,10 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
  }

  explicit FP32Vec16(const BF16Vec16& v) {
-    // On big-endian s390x, place BF16 first to get correct byte order
-    reg.val[0] = (__vector float)vec_mergeh(v.reg.val[0], zero);
-    reg.val[1] = (__vector float)vec_mergel(v.reg.val[0], zero);
-    reg.val[2] = (__vector float)vec_mergeh(v.reg.val[1], zero);
-    reg.val[3] = (__vector float)vec_mergel(v.reg.val[1], zero);
+    reg.val[0] = (__vector float)vec_mergeh(zero, v.reg.val[0]);
+    reg.val[1] = (__vector float)vec_mergel(zero, v.reg.val[0]);
+    reg.val[2] = (__vector float)vec_mergeh(zero, v.reg.val[1]);
+    reg.val[3] = (__vector float)vec_mergel(zero, v.reg.val[1]);
  }

  explicit FP32Vec16(const BF16Vec8& v) : FP32Vec16(FP32Vec8(v)) {}
@@ -585,23 +376,6 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
    return result;
  }

-  FP32Vec16 max(const FP32Vec16& b) const {
-    return FP32Vec16(f32x4x4_t({vec_max(reg.val[0], b.reg.val[0]),
-                                vec_max(reg.val[1], b.reg.val[1]),
-                                vec_max(reg.val[2], b.reg.val[2]),
-                                vec_max(reg.val[3], b.reg.val[3])}));
-  }
-
-  float reduce_max() const {
-    AliasReg ar;
-    ar.reg = reg;
-    float result = ar.values[0];
-    unroll_loop<int, VEC_ELEM_NUM>([&result, &ar](int i) {
-      if (ar.values[i] > result) result = ar.values[i];
-    });
-    return result;
-  }
-
  void save(float* ptr) const {
    vec_xst(reg.val[0], 0, ptr);
    vec_xst(reg.val[1], 16, ptr);
@@ -628,14 +402,15 @@ struct VecType<c10::BFloat16> {
  using vec_type = BF16Vec8;
 };

-// On s390x, FP16 (Half) is not natively supported, use FP32 vectors instead
-using FP16Vec16 = FP32Vec16;
-
 template <typename T>
 void storeFP32(float v, T* ptr) {
  *ptr = v;
 }

+inline void fma(FP32Vec16& acc, FP32Vec16& a, FP32Vec16& b) {
+  acc = acc + a * b;
+}
+
 namespace c10 {
 struct BFloat16 {
  uint16_t value;  // Assume BFloat16 is defined as a struct containing a 16-bit
@@ -654,79 +429,6 @@ inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
  #define __VEC_CLASS_FP_NAN (1 << 6)
 #endif

-// Optimized FMA (Fused Multiply-Add) implementations using IBM Z vector
-// intrinsics
-
-// FP32Vec4 FMA: acc = acc + (a * b) or equivalently acc = fma(a, b, acc)
-FORCE_INLINE void fma(FP32Vec4& acc, const FP32Vec4& a, const FP32Vec4& b) {
-  acc.reg = vec_madd(a.reg, b.reg, acc.reg);
-}
-
-// FP32Vec8 FMA: acc = acc + (a * b)
-FORCE_INLINE void fma(FP32Vec8& acc, const FP32Vec8& a, const FP32Vec8& b) {
-  acc.reg.val[0] = vec_madd(a.reg.val[0], b.reg.val[0], acc.reg.val[0]);
-  acc.reg.val[1] = vec_madd(a.reg.val[1], b.reg.val[1], acc.reg.val[1]);
-}
-
-// FP32Vec16 FMA: acc = acc + (a * b)
-FORCE_INLINE void fma(FP32Vec16& acc, const FP32Vec16& a, const FP32Vec16& b) {
-  acc.reg.val[0] = vec_madd(a.reg.val[0], b.reg.val[0], acc.reg.val[0]);
-  acc.reg.val[1] = vec_madd(a.reg.val[1], b.reg.val[1], acc.reg.val[1]);
-  acc.reg.val[2] = vec_madd(a.reg.val[2], b.reg.val[2], acc.reg.val[2]);
-  acc.reg.val[3] = vec_madd(a.reg.val[3], b.reg.val[3], acc.reg.val[3]);
-}
-
-// Multiply-Subtract: acc = acc - (a * b)
-FORCE_INLINE void fms(FP32Vec4& acc, const FP32Vec4& a, const FP32Vec4& b) {
-  acc.reg = vec_msub(a.reg, b.reg, acc.reg);
-}
-
-FORCE_INLINE void fms(FP32Vec8& acc, const FP32Vec8& a, const FP32Vec8& b) {
-  acc.reg.val[0] = vec_msub(a.reg.val[0], b.reg.val[0], acc.reg.val[0]);
-  acc.reg.val[1] = vec_msub(a.reg.val[1], b.reg.val[1], acc.reg.val[1]);
-}
-
-FORCE_INLINE void fms(FP32Vec16& acc, const FP32Vec16& a, const FP32Vec16& b) {
-  acc.reg.val[0] = vec_msub(a.reg.val[0], b.reg.val[0], acc.reg.val[0]);
-  acc.reg.val[1] = vec_msub(a.reg.val[1], b.reg.val[1], acc.reg.val[1]);
-  acc.reg.val[2] = vec_msub(a.reg.val[2], b.reg.val[2], acc.reg.val[2]);
-  acc.reg.val[3] = vec_msub(a.reg.val[3], b.reg.val[3], acc.reg.val[3]);
-}
-
-// Negative Multiply-Add: acc = -(a * b) + acc
-FORCE_INLINE void nfma(FP32Vec4& acc, const FP32Vec4& a, const FP32Vec4& b) {
-  acc.reg = vec_nmadd(a.reg, b.reg, acc.reg);
-}
-
-FORCE_INLINE void nfma(FP32Vec8& acc, const FP32Vec8& a, const FP32Vec8& b) {
-  acc.reg.val[0] = vec_nmadd(a.reg.val[0], b.reg.val[0], acc.reg.val[0]);
-  acc.reg.val[1] = vec_nmadd(a.reg.val[1], b.reg.val[1], acc.reg.val[1]);
-}
-
-FORCE_INLINE void nfma(FP32Vec16& acc, const FP32Vec16& a, const FP32Vec16& b) {
-  acc.reg.val[0] = vec_nmadd(a.reg.val[0], b.reg.val[0], acc.reg.val[0]);
-  acc.reg.val[1] = vec_nmadd(a.reg.val[1], b.reg.val[1], acc.reg.val[1]);
-  acc.reg.val[2] = vec_nmadd(a.reg.val[2], b.reg.val[2], acc.reg.val[2]);
-  acc.reg.val[3] = vec_nmadd(a.reg.val[3], b.reg.val[3], acc.reg.val[3]);
-}
-
-// Negative Multiply-Subtract: acc = -(a * b) - acc
-FORCE_INLINE void nfms(FP32Vec4& acc, const FP32Vec4& a, const FP32Vec4& b) {
-  acc.reg = vec_nmsub(a.reg, b.reg, acc.reg);
-}
-
-FORCE_INLINE void nfms(FP32Vec8& acc, const FP32Vec8& a, const FP32Vec8& b) {
-  acc.reg.val[0] = vec_nmsub(a.reg.val[0], b.reg.val[0], acc.reg.val[0]);
-  acc.reg.val[1] = vec_nmsub(a.reg.val[1], b.reg.val[1], acc.reg.val[1]);
-}
-
-FORCE_INLINE void nfms(FP32Vec16& acc, const FP32Vec16& a, const FP32Vec16& b) {
-  acc.reg.val[0] = vec_nmsub(a.reg.val[0], b.reg.val[0], acc.reg.val[0]);
-  acc.reg.val[1] = vec_nmsub(a.reg.val[1], b.reg.val[1], acc.reg.val[1]);
-  acc.reg.val[2] = vec_nmsub(a.reg.val[2], b.reg.val[2], acc.reg.val[2]);
-  acc.reg.val[3] = vec_nmsub(a.reg.val[3], b.reg.val[3], acc.reg.val[3]);
-}
-
 const static __vector unsigned char omask = {2,  3,  6,  7,  10, 11, 14, 15,
                                             18, 19, 22, 23, 26, 27, 30, 31};
 const static __vector unsigned int bias = {0x00007fff, 0x00007fff, 0x00007fff,
@@ -739,24 +441,13 @@ const static __vector unsigned int one = {1, 1, 1, 1};
 inline BF16Vec8::BF16Vec8(const FP32Vec8& v) {
  __vector unsigned int inp0 = (__vector unsigned int)(v.reg.val[0]);
  __vector unsigned int inp1 = (__vector unsigned int)(v.reg.val[1]);
-  __vector unsigned int lsb0 = inp0 >> sh16;
-  __vector unsigned int lsb1 = inp1 >> sh16;
-  lsb0 = lsb0 & one;
-  lsb1 = lsb1 & one;
-  __vector unsigned int rnd0 = lsb0 + bias;
-  __vector unsigned int rnd1 = lsb1 + bias;
-  inp0 = inp0 + rnd0;
-  inp1 = inp1 + rnd1;
  int cc;
  __vector __bool int sel0 =
      vec_fp_test_data_class(v.reg.val[0], __VEC_CLASS_FP_NAN, &cc);
  __vector __bool int sel1 =
      vec_fp_test_data_class(v.reg.val[1], __VEC_CLASS_FP_NAN, &cc);
-  inp0 = vec_sel(inp0, nan, sel0);
-  inp1 = vec_sel(inp1, nan, sel1);
-  inp0 = inp0 >> sh16;
-  inp1 = inp1 >> sh16;
-
+  inp0 = vec_sel(inp0, nan, sel0) >> sh16;
+  inp1 = vec_sel(inp1, nan, sel1) >> sh16;
  reg = (__vector signed short)vec_perm(inp0, inp1, omask);
 }

@@ -765,22 +456,6 @@ inline BF16Vec16::BF16Vec16(const FP32Vec16& v) {
  __vector unsigned int inp1 = (__vector unsigned int)(v.reg.val[1]);
  __vector unsigned int inp2 = (__vector unsigned int)(v.reg.val[2]);
  __vector unsigned int inp3 = (__vector unsigned int)(v.reg.val[3]);
-  __vector unsigned int lsb0 = inp0 >> sh16;
-  __vector unsigned int lsb1 = inp1 >> sh16;
-  __vector unsigned int lsb2 = inp2 >> sh16;
-  __vector unsigned int lsb3 = inp3 >> sh16;
-  lsb0 = lsb0 & one;
-  lsb1 = lsb1 & one;
-  lsb2 = lsb2 & one;
-  lsb3 = lsb3 & one;
-  __vector unsigned int rnd0 = lsb0 + bias;
-  __vector unsigned int rnd1 = lsb1 + bias;
-  __vector unsigned int rnd2 = lsb2 + bias;
-  __vector unsigned int rnd3 = lsb3 + bias;
-  inp0 = inp0 + rnd0;
-  inp1 = inp1 + rnd1;
-  inp2 = inp2 + rnd2;
-  inp3 = inp3 + rnd3;
  int cc;
  __vector __bool int sel0 =
      vec_fp_test_data_class(v.reg.val[0], __VEC_CLASS_FP_NAN, &cc);
@@ -790,164 +465,15 @@ inline BF16Vec16::BF16Vec16(const FP32Vec16& v) {
      vec_fp_test_data_class(v.reg.val[2], __VEC_CLASS_FP_NAN, &cc);
  __vector __bool int sel3 =
      vec_fp_test_data_class(v.reg.val[3], __VEC_CLASS_FP_NAN, &cc);
-  inp0 = vec_sel(inp0, nan, sel0);
-  inp1 = vec_sel(inp1, nan, sel1);
-  inp2 = vec_sel(inp2, nan, sel2);
-  inp3 = vec_sel(inp3, nan, sel3);
-  inp0 = inp0 >> sh16;
-  inp1 = inp1 >> sh16;
-  inp2 = inp2 >> sh16;
-  inp3 = inp3 >> sh16;
-
+  inp0 = vec_sel(inp0, nan, sel0) >> sh16;
+  inp1 = vec_sel(inp1, nan, sel1) >> sh16;
+  inp2 = vec_sel(inp2, nan, sel2) >> sh16;
+  inp3 = vec_sel(inp3, nan, sel3) >> sh16;
  reg.val[0] = (__vector signed short)vec_perm(inp0, inp1, omask);
  reg.val[1] = (__vector signed short)vec_perm(inp2, inp3, omask);
 }

-// 1D softmax over `n` elements in `input`, writes result to `output`.
-// Uses FP32Vec8 for main body, scalar tail handling.
-// Requirement: n > 0
-FORCE_INLINE void softmax_fp32vec8(float* output, const float* input, int n) {
-  if (n <= 0) return;
-
-  // ---------- Pass 1: find max ----------
-  float max_val = -std::numeric_limits<float>::infinity();
-  int i = 0;
-
-  for (; i + FP32Vec8::VEC_ELEM_NUM <= n; i += FP32Vec8::VEC_ELEM_NUM) {
-    FP32Vec8 v(input + i);
-    FP32Vec8::AliasReg ar;
-    ar.reg = v.reg;
-    for (int j = 0; j < FP32Vec8::VEC_ELEM_NUM; ++j) {
-      if (ar.values[j] > max_val) max_val = ar.values[j];
-    }
-  }
-  for (; i < n; ++i) {
-    if (input[i] > max_val) max_val = input[i];
-  }
-
-  // ---------- Pass 2: compute exp(x - max) and sum ----------
-  float sum = 0.0f;
-  i = 0;
-
-  for (; i + FP32Vec8::VEC_ELEM_NUM <= n; i += FP32Vec8::VEC_ELEM_NUM) {
-    float tmp[FP32Vec8::VEC_ELEM_NUM];
-    for (int j = 0; j < FP32Vec8::VEC_ELEM_NUM; ++j) {
-      tmp[j] = input[i + j] - max_val;
-    }
-
-    FP32Vec8 v(tmp);
-    FP32Vec8 e = v.exp();
-
-    FP32Vec8::AliasReg ar;
-    ar.reg = e.reg;
-    for (int j = 0; j < FP32Vec8::VEC_ELEM_NUM; ++j) {
-      output[i + j] = ar.values[j];
-      sum += ar.values[j];
-    }
-  }
-
-  // Tail
-  for (; i < n; ++i) {
-    float x = input[i] - max_val;
-    float ex = std::exp(x);  // scalar tail
-    output[i] = ex;
-    sum += ex;
-  }
-
-  // ---------- Pass 3: normalize ----------
-  float inv_sum = 1.0f / sum;
-  i = 0;
-
-  for (; i + FP32Vec8::VEC_ELEM_NUM <= n; i += FP32Vec8::VEC_ELEM_NUM) {
-    float tmp[FP32Vec8::VEC_ELEM_NUM];
-    for (int j = 0; j < FP32Vec8::VEC_ELEM_NUM; ++j) {
-      tmp[j] = output[i + j] * inv_sum;
-    }
-    FP32Vec8 v(tmp);
-    v.save(output + i);
-  }
-
-  for (; i < n; ++i) {
-    output[i] *= inv_sum;
-  }
-}
-
-// 1D RMSNorm kernel:
-//   input:  x[0..n-1]
-//   weight: w[0..n-1] (gamma), may be nullptr
-//   output: y[i] = x[i] * inv_rms * (weight[i] if weight != nullptr else 1)
-//   eps: small epsilon for numerical stability
-FORCE_INLINE void rmsnorm_fp32vec8(float* output, const float* input,
-                                   const float* weight, int n, float eps) {
-  if (n <= 0) return;
-
-  // ---------- Pass 1: compute sum of squares ----------
-  float sum_sq = 0.0f;
-  int i = 0;
-
-  for (; i + FP32Vec8::VEC_ELEM_NUM <= n; i += FP32Vec8::VEC_ELEM_NUM) {
-    FP32Vec8 x_vec(input + i);
-
-    FP32Vec8 sq = x_vec * x_vec;
-
-    FP32Vec8::AliasReg ar;
-    ar.reg = sq.reg;
-    for (int j = 0; j < FP32Vec8::VEC_ELEM_NUM; ++j) {
-      sum_sq += ar.values[j];
-    }
-  }
-
-  // Tail
-  for (; i < n; ++i) {
-    float v = input[i];
-    sum_sq += v * v;
-  }
-
-  float mean_sq = sum_sq / static_cast<float>(n);
-  float inv_rms = 1.0f / std::sqrt(mean_sq + eps);
-
-  // ---------- Pass 2: scale (and apply weight if given) ----------
-  const float inv_rms_f = inv_rms;
-  i = 0;
-
-  if (weight) {
-    // with gamma
-    for (; i + FP32Vec8::VEC_ELEM_NUM <= n; i += FP32Vec8::VEC_ELEM_NUM) {
-      FP32Vec8 x_vec(input + i);
-
-      float wtmp[FP32Vec8::VEC_ELEM_NUM];
-      for (int j = 0; j < FP32Vec8::VEC_ELEM_NUM; ++j) {
-        wtmp[j] = weight[i + j];
-      }
-      FP32Vec8 w_vec(wtmp);
-
-      FP32Vec8 scale_vec(inv_rms_f);
-      FP32Vec8 y = x_vec * scale_vec * w_vec;
-      y.save(output + i);
-    }
-
-    for (; i < n; ++i) {
-      output[i] = input[i] * inv_rms_f * weight[i];
-    }
-  } else {
-    // without gamma
-    for (; i + FP32Vec8::VEC_ELEM_NUM <= n; i += FP32Vec8::VEC_ELEM_NUM) {
-      FP32Vec8 x_vec(input + i);
-      FP32Vec8 scale_vec(inv_rms_f);
-      FP32Vec8 y = x_vec * scale_vec;
-      y.save(output + i);
-    }
-
-    for (; i < n; ++i) {
-      output[i] = input[i] * inv_rms_f;
-    }
-  }
-}
-
-// Prefetch data to cache for better memory access performance
-FORCE_INLINE void prefetch(const void* addr) {
-  __builtin_prefetch(addr, 0, 3);  // 0=read, 3=high temporal locality
-}
+inline void prefetch(const void* addr) { void __dcbt(const void* addr); }

 };  // namespace vec_op

@@ -45,55 +45,31 @@ std::string init_cpu_threads_env(const std::string& cpu_ids) {
  // Memory node binding
  if (numa_available() != -1) {
    int mem_node_id = numa_node_of_cpu(omp_cpu_ids.front());
-    std::set<int> node_ids;
-    for (const auto& cpu_id : omp_cpu_ids) {
-      int node_id = numa_node_of_cpu(cpu_id);
-      if (node_id != -1) {
-        node_ids.insert(node_id);
-      }
-      if (node_id != mem_node_id) {
-        TORCH_WARN("CPU ", cpu_id, " is on NUMA node ", node_id, ", but CPU ",
-                   omp_cpu_ids.front(), " is on NUMA node ", mem_node_id,
-                   ". All CPUs should be on the same NUMA node for optimal "
-                   "performance. Memory will be bound to NUMA node ",
-                   mem_node_id, ".");
-      }
+    // Verify all CPUs are on the same NUMA node
+    for (size_t i = 1; i < omp_cpu_ids.size(); ++i) {
+      int node_id = numa_node_of_cpu(omp_cpu_ids[i]);
+      TORCH_CHECK(node_id == mem_node_id, "CPU ", omp_cpu_ids[i],
+                  " is on NUMA node ", node_id, ", but CPU ",
+                  omp_cpu_ids.front(), " is on NUMA node ", mem_node_id,
+                  ". All CPUs should be on the same NUMA node for optimal "
+                  "performance. Memory will be bound to NUMA node ",
+                  mem_node_id, ".");
    }
-    // Concatenate all node_ids into a single comma-separated string
-    if (!node_ids.empty()) {
-      std::string node_ids_str;
-      for (const int node_id : node_ids) {
-        if (!node_ids_str.empty()) {
-          node_ids_str += ",";
-        }
-        node_ids_str += std::to_string(node_id);
-      }
+    bitmask* mask = numa_parse_nodestring(std::to_string(mem_node_id).c_str());
+    bitmask* src_mask = numa_get_membind();

-      bitmask* mask = numa_parse_nodestring(node_ids_str.c_str());
-      bitmask* src_mask = numa_get_membind();
+    int pid = getpid();

-      int pid = getpid();
-
-      if (mask && src_mask) {
-        // move all existing pages to the specified numa node.
-        *(src_mask->maskp) = *(src_mask->maskp) ^ *(mask->maskp);
-        int page_num = numa_migrate_pages(pid, src_mask, mask);
-        if (page_num == -1) {
-          TORCH_WARN("numa_migrate_pages failed. errno: " +
-                     std::to_string(errno));
-        }
-
-        // restrict memory allocation node.
-        numa_set_membind(mask);
-        numa_set_strict(1);
-
-        numa_free_nodemask(mask);
-        numa_free_nodemask(src_mask);
-      } else {
-        TORCH_WARN("numa_parse_nodestring or numa_get_membind failed. errno: " +
-                   std::to_string(errno));
-      }
+    // move all existing pages to the specified numa node.
+    *(src_mask->maskp) = *(src_mask->maskp) ^ *(mask->maskp);
+    int page_num = numa_migrate_pages(pid, src_mask, mask);
+    if (page_num == -1) {
+      TORCH_WARN("numa_migrate_pages failed. errno: " + std::to_string(errno));
    }
+
+    // restrict memory allocation node.
+    numa_set_membind(mask);
+    numa_set_strict(1);
  }

  // OMP threads binding
@@ -22,10 +22,15 @@ torch::Tensor get_cuda_view_from_cpu_tensor(torch::Tensor& cpu_tensor) {
  auto strides = cpu_tensor.strides();
  auto options = cpu_tensor.options().device(torch::kCUDA);

-  // use default no-op deleter, since the memory is owned by the original CPU
-  // tensor
+  // from_blob signature: from_blob(void *data, IntArrayRef sizes, ..., Deleter,
+  // const TensorOptions &) Provide a no-op deleter. The CPU tensor holds the
+  // memory, so we don't free it here.
+  auto deleter = [](void*) {
+    // no-op, since the memory is owned by the original CPU tensor
+  };
+
  torch::Tensor cuda_tensor =
-      torch::from_blob(device_ptr, sizes, strides, options);
+      torch::from_blob(device_ptr, sizes, strides, deleter, options);

  TORCH_CHECK(cuda_tensor.device().is_cuda(),
              "Resulting tensor is not on CUDA device");
@@ -117,42 +117,3 @@
      break;                                  \
    }                                         \
  }
-
-#define VLLM_DISPATCH_BOOL(expr, const_expr, ...) \
-  if (expr) {                                     \
-    constexpr bool const_expr = true;             \
-    __VA_ARGS__();                                \
-  } else {                                        \
-    constexpr bool const_expr = false;            \
-    __VA_ARGS__();                                \
-  }
-
-#define VLLM_DISPATCH_GROUP_SIZE(group_size, const_group_size, ...) \
-  if (group_size == 128) {                                          \
-    constexpr int const_group_size = 128;                           \
-    __VA_ARGS__();                                                  \
-  } else if (group_size == 64) {                                    \
-    constexpr int const_group_size = 64;                            \
-    __VA_ARGS__();                                                  \
-  }
-
-#define VLLM_DISPATCH_RANK234(NUM_DIMS, ...)                                   \
-  switch (NUM_DIMS) {                                                          \
-    case 2: {                                                                  \
-      constexpr int tensor_rank = 2;                                           \
-      __VA_ARGS__();                                                           \
-      break;                                                                   \
-    }                                                                          \
-    case 3: {                                                                  \
-      constexpr int tensor_rank = 3;                                           \
-      __VA_ARGS__();                                                           \
-      break;                                                                   \
-    }                                                                          \
-    case 4: {                                                                  \
-      constexpr int tensor_rank = 4;                                           \
-      __VA_ARGS__();                                                           \
-      break;                                                                   \
-    }                                                                          \
-    default:                                                                   \
-      TORCH_CHECK(false, "Expects rank 2, 3 or 4 tensors but got ", NUM_DIMS); \
-  }
@@ -10,38 +10,16 @@
 namespace vllm {

 // TODO(woosuk): Further optimize this kernel.
-template <typename scalar_t, int VEC_SIZE, int NUM_DIMS>
+template <typename scalar_t, int VEC_SIZE>
 __global__ void rms_norm_kernel(
-    scalar_t* __restrict__ out,           // [..., hidden_size]
-    const scalar_t* __restrict__ input,   // [..., hidden_size]
-    const int64_t input_stride_d2,        // input.stride(-2)
-    const int64_t input_stride_d3,        // input.stride(-3)
-    const int64_t input_stride_d4,        // input.stride(-4)
-    const int64_t input_shape_d2,         // input.size(-2)
-    const int64_t input_shape_d3,         // input.size(-3)
+    scalar_t* __restrict__ out,          // [..., hidden_size]
+    const scalar_t* __restrict__ input,  // [..., hidden_size]
+    const int64_t input_stride,
    const scalar_t* __restrict__ weight,  // [hidden_size]
    const float epsilon, const int num_tokens, const int hidden_size) {
  __shared__ float s_variance;
  float variance = 0.0f;
-  const scalar_t* input_row;
-  if constexpr (NUM_DIMS == 2) {
-    // 2D for layernorm normal case [batch_size, hidden]
-    input_row = input + blockIdx.x * input_stride_d2;
-  } else if constexpr (NUM_DIMS == 3) {
-    // 3D for q/k norm [batch_size, num_heads, head_size]
-    int batch_idx = blockIdx.x / input_shape_d2;
-    int head_idx = blockIdx.x % input_shape_d2;
-    input_row =
-        input + batch_idx * input_stride_d3 + head_idx * input_stride_d2;
-  } else if constexpr (NUM_DIMS == 4) {
-    // 4D for transformers model_impl qk norm [batch, seq, head, head_dim]
-    int batch_idx = blockIdx.x / (input_shape_d3 * input_shape_d2);
-    int remaining = blockIdx.x % (input_shape_d3 * input_shape_d2);
-    int seq_idx = remaining / input_shape_d2;
-    int head_idx = remaining % input_shape_d2;
-    input_row = input + batch_idx * input_stride_d4 +
-                seq_idx * input_stride_d3 + head_idx * input_stride_d2;
-  }
+  const scalar_t* input_row = input + blockIdx.x * input_stride;

  auto vec_op = [&variance](const vec_n_t<scalar_t, VEC_SIZE>& vec) {
 #pragma unroll
@@ -186,44 +164,38 @@ void rms_norm(torch::Tensor& out,     // [..., hidden_size]
              torch::Tensor& weight,  // [hidden_size]
              double epsilon) {
  TORCH_CHECK(out.is_contiguous());
-  if (input.stride(-1) != 1) {
-    input = input.contiguous();
-  }
  TORCH_CHECK(input.stride(-1) == 1);
  TORCH_CHECK(weight.is_contiguous());

  int hidden_size = input.size(-1);

-  int num_tokens = input.numel() / hidden_size;
-  int num_dims = input.dim();
-  int64_t input_stride_d2 = input.stride(-2);
-  int64_t input_stride_d3 = (num_dims >= 3) ? input.stride(-3) : 0;
-  int64_t input_stride_d4 = (num_dims >= 4) ? input.stride(-4) : 0;
-  int64_t input_shape_d2 = (num_dims >= 3) ? input.size(-2) : 0;
-  int64_t input_shape_d3 = (num_dims >= 4) ? input.size(-3) : 0;
+  // We cannot just use `input.stride(-2)` if the tensor is not row-major.
+  // Instead, we use a 2d view to get the second-innermost stride.
+  // That way the dimensions (except the last one) can be arbitrarily permuted.
+  torch::Tensor input_view = input.view({-1, hidden_size});
+
+  int num_tokens = input_view.numel() / hidden_size;
+  int64_t input_stride = input_view.stride(-2);

  // For large num_tokens, use smaller blocks to increase SM concurrency.
  const int max_block_size = (num_tokens < 256) ? 1024 : 256;
  dim3 grid(num_tokens);
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(input_view));
  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-  VLLM_DISPATCH_RANK234(num_dims, [&] {
-    VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "rms_norm_kernel", [&] {
-      const int calculated_vec_size =
-          std::gcd(16 / sizeof(scalar_t), hidden_size);
-      const int block_size =
-          std::min(hidden_size / calculated_vec_size, max_block_size);
-      dim3 block(block_size);
-      VLLM_DISPATCH_VEC_SIZE(calculated_vec_size, [&] {
-        vllm::rms_norm_kernel<scalar_t, vec_size, tensor_rank>
-            <<<grid, block, 0, stream>>>(
-                out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(),
-                input_stride_d2, input_stride_d3, input_stride_d4,
-                input_shape_d2, input_shape_d3, weight.data_ptr<scalar_t>(),
-                epsilon, num_tokens, hidden_size);
+  VLLM_DISPATCH_FLOATING_TYPES(
+      input_view.scalar_type(), "rms_norm_kernel", [&] {
+        const int calculated_vec_size =
+            std::gcd(16 / sizeof(scalar_t), hidden_size);
+        const int block_size =
+            std::min(hidden_size / calculated_vec_size, max_block_size);
+        dim3 block(block_size);
+        VLLM_DISPATCH_VEC_SIZE(calculated_vec_size, [&] {
+          vllm::rms_norm_kernel<scalar_t, vec_size><<<grid, block, 0, stream>>>(
+              out.data_ptr<scalar_t>(), input_view.data_ptr<scalar_t>(),
+              input_stride, weight.data_ptr<scalar_t>(), epsilon, num_tokens,
+              hidden_size);
+        });
      });
-    });
-  });
 }

 #define LAUNCH_FUSED_ADD_RMS_NORM(width)                                    \
@@ -93,16 +93,16 @@ torch::Tensor dynamic_4bit_int_moe_cpu(
  }
  auto Y_all = at::empty({offsets[E], H}, x_c.options());

-  at::parallel_for(0, offsets[E], 0, [&](int64_t idx_begin, int64_t idx_end) {
+  at::parallel_for(0, E, 1, [&](int64_t e_begin, int64_t e_end) {
    c10::InferenceMode guard;
-    for (int64_t e = 0; e < E; ++e) {
-      int64_t start = std::max(offsets[e], idx_begin);
-      int64_t end = std::min(offsets[e + 1], idx_end);
-      int64_t te = end - start;
-      if (te <= 0) {
+    for (int64_t e = e_begin; e < e_end; ++e) {
+      const int64_t te = counts[e];
+      if (te == 0) {
        continue;
      }

+      const int64_t start = offsets[e];
+
      auto x_e = X_all.narrow(/*dim=*/0, /*start=*/start, /*length=*/te);

      auto w13_e = w13_packed.select(/*dim=*/0, e);
@@ -444,27 +444,23 @@ __device__ inline T apply_sigmoid(T val) {
  return cuda_cast<T, float>(sigmoid_accurate(f));
 }

-template <ScoringFunc SF, typename T>
-__device__ inline T apply_scoring(T val) {
-  if constexpr (SF == SCORING_SIGMOID) {
-    return apply_sigmoid(val);
-  } else {
-    return val;
-  }
-}
-
-template <typename T, ScoringFunc SF>
+template <typename T>
 __device__ void topk_with_k2(T* output, T const* input, T const* bias,
                             cg::thread_block_tile<32> const& tile,
                             int32_t const lane_id,
-                             int const num_experts_per_group) {
+                             int const num_experts_per_group,
+                             int const scoring_func) {
  // Get the top2 per thread
  T largest = neg_inf<T>();
  T second_largest = neg_inf<T>();

  if (num_experts_per_group > WARP_SIZE) {
    for (int i = lane_id; i < num_experts_per_group; i += WARP_SIZE) {
-      T value = apply_scoring<SF>(input[i]);
+      T value = input[i];
+      // Apply scoring function if needed
+      if (scoring_func == SCORING_SIGMOID) {
+        value = apply_sigmoid(value);
+      }
      value = value + bias[i];

      if (value > largest) {
@@ -476,11 +472,17 @@ __device__ void topk_with_k2(T* output, T const* input, T const* bias,
    }
  } else {
    for (int i = lane_id; i < num_experts_per_group; i += WARP_SIZE) {
-      T value = apply_scoring<SF>(input[i]);
+      T value = input[i];
+      // Apply scoring function if needed
+      if (scoring_func == SCORING_SIGMOID) {
+        value = apply_sigmoid(value);
+      }
      value = value + bias[i];
      largest = value;
    }
  }
+
+  __syncwarp();  // Ensure all threads have valid data before reduction
  // Get the top2 warpwise
  T max1 = cg::reduce(tile, largest, cg::greater<T>());

@@ -499,12 +501,13 @@ __device__ void topk_with_k2(T* output, T const* input, T const* bias,
  }
 }

-template <typename T, ScoringFunc SF>
+template <typename T>
 __global__ void topk_with_k2_kernel(T* output, T* input, T const* bias,
                                    int64_t const num_tokens,
                                    int64_t const num_cases,
                                    int64_t const n_group,
-                                    int64_t const num_experts_per_group) {
+                                    int64_t const num_experts_per_group,
+                                    int const scoring_func) {
  int32_t warp_id = threadIdx.x / WARP_SIZE;
  int32_t lane_id = threadIdx.x % WARP_SIZE;

@@ -522,21 +525,21 @@ __global__ void topk_with_k2_kernel(T* output, T* input, T const* bias,
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
    asm volatile("griddepcontrol.wait;");
 #endif
-    topk_with_k2<T, SF>(output, input, group_bias, tile, lane_id,
-                        num_experts_per_group);
+    topk_with_k2(output, input, group_bias, tile, lane_id,
+                 num_experts_per_group, scoring_func);
  }
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
  asm volatile("griddepcontrol.launch_dependents;");
 #endif
 }

-template <typename T, typename IdxT, ScoringFunc SF, int NGroup = -1>
+template <typename T, typename IdxT>
 __global__ void group_idx_and_topk_idx_kernel(
    T* scores, T const* group_scores, float* topk_values, IdxT* topk_indices,
    T const* bias, int64_t const num_tokens, int64_t const n_group,
    int64_t const topk_group, int64_t const topk, int64_t const num_experts,
    int64_t const num_experts_per_group, bool renormalize,
-    double routed_scaling_factor) {
+    double routed_scaling_factor, int scoring_func) {
  int32_t warp_id = threadIdx.x / WARP_SIZE;
  int32_t lane_id = threadIdx.x % WARP_SIZE;
  int32_t case_id =
@@ -546,11 +549,6 @@ __global__ void group_idx_and_topk_idx_kernel(
  topk_values += case_id * topk;
  topk_indices += case_id * topk;

-  constexpr bool kUseStaticNGroup = (NGroup > 0);
-  // use int32 to avoid implicit conversion
-  int32_t const n_group_i32 =
-      kUseStaticNGroup ? NGroup : static_cast<int32_t>(n_group);
-
  int32_t align_num_experts_per_group =
      warp_topk::round_up_to_multiple_of<WARP_SIZE>(num_experts_per_group);

@@ -576,17 +574,17 @@ __global__ void group_idx_and_topk_idx_kernel(

  if (case_id < num_tokens) {
    // calculate group_idx
-    int32_t target_num_min =
-        WARP_SIZE - n_group_i32 + static_cast<int32_t>(topk_group);
+    int32_t target_num_min = WARP_SIZE - n_group + topk_group;
    // The check is necessary to avoid abnormal input
-    if (lane_id < n_group_i32 && is_finite(group_scores[lane_id])) {
+    if (lane_id < n_group && is_finite(group_scores[lane_id])) {
      value = group_scores[lane_id];
    }

-    int count_equal_to_top_value = WARP_SIZE - n_group_i32;
+    int count_equal_to_top_value = WARP_SIZE - n_group;
    int pre_count_equal_to_top_value = 0;
    // Use loop to find the largset top_group
    while (count_equal_to_top_value < target_num_min) {
+      __syncwarp();  // Ensure all threads have valid data before reduction
      topk_group_value = cg::reduce(tile, value, cg::greater<T>());
      if (value == topk_group_value) {
        value = neg_inf<T>();
@@ -606,7 +604,7 @@ __global__ void group_idx_and_topk_idx_kernel(
  int count_equalto_topkth_group = 0;
  bool if_proceed_next_topk = topk_group_value != neg_inf<T>();
  if (case_id < num_tokens && if_proceed_next_topk) {
-    auto process_group = [&](int i_group) {
+    for (int i_group = 0; i_group < n_group; i_group++) {
      if ((group_scores[i_group] > topk_group_value) ||
          ((group_scores[i_group] == topk_group_value) &&
           (count_equalto_topkth_group < num_equalto_topkth_group))) {
@@ -615,10 +613,11 @@ __global__ void group_idx_and_topk_idx_kernel(
             i += WARP_SIZE) {
          T candidates = neg_inf<T>();
          if (i < num_experts_per_group) {
-            // apply scoring function (if any) and add bias
+            // Apply scoring function (if any) and add bias
            T input = scores[offset + i];
            if (is_finite(input)) {
-              T score = apply_scoring<SF>(input);
+              T score = (scoring_func == SCORING_SIGMOID) ? apply_sigmoid(input)
+                                                          : input;
              candidates = score + bias[offset + i];
            }
          }
@@ -628,21 +627,12 @@ __global__ void group_idx_and_topk_idx_kernel(
          count_equalto_topkth_group++;
        }
      }
-    };
-
-    if constexpr (kUseStaticNGroup) {
-#pragma unroll
-      for (int i_group = 0; i_group < NGroup; ++i_group) {
-        process_group(i_group);
-      }
-    } else {
-      for (int i_group = 0; i_group < n_group_i32; ++i_group) {
-        process_group(i_group);
-      }
    }
    queue.done();
+    __syncwarp();
    // Get the topk_idx
    queue.dumpIdx(s_topk_idx);
+    __syncwarp();
  }

  // Load the valid score value
@@ -656,13 +646,12 @@ __global__ void group_idx_and_topk_idx_kernel(
      if (i < topk) {
        // Load the score value (without bias) for normalization
        T input = scores[s_topk_idx[i]];
-        value = apply_scoring<SF>(input);
+        value =
+            (scoring_func == SCORING_SIGMOID) ? apply_sigmoid(input) : input;
        s_topk_value[i] = value;
      }
-      if (renormalize) {
-        topk_sum +=
-            cg::reduce(tile, cuda_cast<float, T>(value), cg::plus<float>());
-      }
+      topk_sum +=
+          cg::reduce(tile, cuda_cast<float, T>(value), cg::plus<float>());
    }
  }

@@ -671,9 +660,13 @@ __global__ void group_idx_and_topk_idx_kernel(
  if (case_id < num_tokens) {
    if (if_proceed_next_topk) {
      for (int i = lane_id; i < topk; i += WARP_SIZE) {
-        float base = cuda_cast<float, T>(s_topk_value[i]);
-        float value = renormalize ? (base / topk_sum * routed_scaling_factor)
-                                  : (base * routed_scaling_factor);
+        float value;
+        if (renormalize) {
+          value = cuda_cast<float, T>(s_topk_value[i]) / topk_sum *
+                  routed_scaling_factor;
+        } else {
+          value = cuda_cast<float, T>(s_topk_value[i]) * routed_scaling_factor;
+        }
        topk_indices[i] = s_topk_idx[i];
        topk_values[i] = value;
      }
@@ -691,45 +684,6 @@ __global__ void group_idx_and_topk_idx_kernel(
 #endif
 }

-template <typename T, typename IdxT, ScoringFunc SF>
-inline void launch_group_idx_and_topk_kernel(
-    cudaLaunchConfig_t const& config, T* scores, T* group_scores,
-    float* topk_values, IdxT* topk_indices, T const* bias,
-    int64_t const num_tokens, int64_t const n_group, int64_t const topk_group,
-    int64_t const topk, int64_t const num_experts,
-    int64_t const num_experts_per_group, bool const renormalize,
-    double const routed_scaling_factor) {
-  auto launch = [&](auto* kernel_instance2) {
-    cudaLaunchKernelEx(&config, kernel_instance2, scores, group_scores,
-                       topk_values, topk_indices, bias, num_tokens, n_group,
-                       topk_group, topk, num_experts, num_experts_per_group,
-                       renormalize, routed_scaling_factor);
-  };
-
-  switch (n_group) {
-    case 4: {
-      launch(&group_idx_and_topk_idx_kernel<T, IdxT, SF, 4>);
-      break;
-    }
-    case 8: {
-      launch(&group_idx_and_topk_idx_kernel<T, IdxT, SF, 8>);
-      break;
-    }
-    case 16: {
-      launch(&group_idx_and_topk_idx_kernel<T, IdxT, SF, 16>);
-      break;
-    }
-    case 32: {
-      launch(&group_idx_and_topk_idx_kernel<T, IdxT, SF, 32>);
-      break;
-    }
-    default: {
-      launch(&group_idx_and_topk_idx_kernel<T, IdxT, SF>);
-      break;
-    }
-  }
-}
-
 template <typename T, typename IdxT>
 void invokeNoAuxTc(T* scores, T* group_scores, float* topk_values,
                   IdxT* topk_indices, T const* bias, int64_t const num_tokens,
@@ -740,6 +694,7 @@ void invokeNoAuxTc(T* scores, T* group_scores, float* topk_values,
                   cudaStream_t const stream = 0) {
  int64_t num_cases = num_tokens * n_group;
  int64_t topk_with_k2_num_blocks = (num_cases - 1) / NUM_WARPS_PER_BLOCK + 1;
+  auto* kernel_instance1 = &topk_with_k2_kernel<T>;
  cudaLaunchConfig_t config;
  config.gridDim = topk_with_k2_num_blocks;
  config.blockDim = BLOCK_SIZE;
@@ -750,33 +705,16 @@ void invokeNoAuxTc(T* scores, T* group_scores, float* topk_values,
  attrs[0].val.programmaticStreamSerializationAllowed = enable_pdl;
  config.numAttrs = 1;
  config.attrs = attrs;
-  auto const sf = static_cast<ScoringFunc>(scoring_func);
-  int64_t const num_experts_per_group = num_experts / n_group;
-  auto launch_topk_with_k2 = [&](auto* kernel_instance1) {
-    cudaLaunchKernelEx(&config, kernel_instance1, group_scores, scores, bias,
-                       num_tokens, num_cases, n_group, num_experts_per_group);
-  };
-  switch (sf) {
-    case SCORING_NONE: {
-      auto* kernel_instance1 = &topk_with_k2_kernel<T, SCORING_NONE>;
-      launch_topk_with_k2(kernel_instance1);
-      break;
-    }
-    case SCORING_SIGMOID: {
-      auto* kernel_instance1 = &topk_with_k2_kernel<T, SCORING_SIGMOID>;
-      launch_topk_with_k2(kernel_instance1);
-      break;
-    }
-    default:
-      // should be guarded by higher level checks.
-      TORCH_CHECK(false, "Unsupported scoring_func in invokeNoAuxTc");
-  }
+  cudaLaunchKernelEx(&config, kernel_instance1, group_scores, scores, bias,
+                     num_tokens, num_cases, n_group, num_experts / n_group,
+                     scoring_func);

  int64_t topk_with_k_group_num_blocks =
      (num_tokens - 1) / NUM_WARPS_PER_BLOCK + 1;
  size_t dynamic_smem_in_bytes =
      warp_topk::calc_smem_size_for_block_wide<T, int32_t>(NUM_WARPS_PER_BLOCK,
                                                           topk);
+  auto* kernel_instance2 = &group_idx_and_topk_idx_kernel<T, IdxT>;
  config.gridDim = topk_with_k_group_num_blocks;
  config.blockDim = BLOCK_SIZE;
  config.dynamicSmemBytes = dynamic_smem_in_bytes;
@@ -785,24 +723,10 @@ void invokeNoAuxTc(T* scores, T* group_scores, float* topk_values,
  attrs[0].val.programmaticStreamSerializationAllowed = enable_pdl;
  config.numAttrs = 1;
  config.attrs = attrs;
-  switch (sf) {
-    case SCORING_NONE: {
-      launch_group_idx_and_topk_kernel<T, IdxT, SCORING_NONE>(
-          config, scores, group_scores, topk_values, topk_indices, bias,
-          num_tokens, n_group, topk_group, topk, num_experts,
-          num_experts_per_group, renormalize, routed_scaling_factor);
-      break;
-    }
-    case SCORING_SIGMOID: {
-      launch_group_idx_and_topk_kernel<T, IdxT, SCORING_SIGMOID>(
-          config, scores, group_scores, topk_values, topk_indices, bias,
-          num_tokens, n_group, topk_group, topk, num_experts,
-          num_experts_per_group, renormalize, routed_scaling_factor);
-      break;
-    }
-    default:
-      TORCH_CHECK(false, "Unsupported scoring_func in invokeNoAuxTc");
-  }
+  cudaLaunchKernelEx(&config, kernel_instance2, scores, group_scores,
+                     topk_values, topk_indices, bias, num_tokens, n_group,
+                     topk_group, topk, num_experts, num_experts / n_group,
+                     renormalize, routed_scaling_factor, scoring_func);
 }

 #define INSTANTIATE_NOAUX_TC(T, IdxT)                                       \
@@ -1,2 +1 @@
-sm*_kernel_*.cu
-kernel_selector.h
+kernel_*.cu
@@ -4,282 +4,134 @@ import glob
 import itertools
 import os
 import subprocess
-import sys

 import jinja2

-ARCHS = []
-SUPPORT_FP8 = False
-for arch in sys.argv[1].split(","):
-    arch = arch[: arch.index(".") + 2].replace(".", "")
-    arch = int(arch)
-    # only SM89 and SM120 fully support
-    # mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32.
-    # SM90 and SM100 can use this PTX, but it’s simulated
-    # with FP16 MMA, so it cannot achieve any acceleration.
-    if arch in [89, 120]:
-        SUPPORT_FP8 = True
-
-FILE_HEAD_COMMENT = """
-// auto generated by generate_kernels.py
+FILE_HEAD = """
+// auto generated by generate.py
 // clang-format off
-""".lstrip()

-FILE_HEAD = (
-    FILE_HEAD_COMMENT
-    + """
 #include "kernel.h"
 #include "marlin_template.h"

 namespace MARLIN_NAMESPACE_NAME {
-"""
-)
+""".strip()

 TEMPLATE = (
    "template __global__ void Marlin<"
-    "{{a_type_id}}, "
-    "{{b_type_id}}, "
-    "{{c_type_id}}, "
+    "{{scalar_t}}, "
+    "{{w_type_id}}, "
    "{{s_type_id}}, "
    "{{threads}}, "
    "{{thread_m_blocks}}, "
    "{{thread_n_blocks}}, "
    "{{thread_k_blocks}}, "
-    "{{m_block_size_8}}, "
+    "{{'true' if m_block_size_8 else 'false'}}, "
    "{{stages}}, "
    "{{group_blocks}}, "
-    "{{is_zp_float}}>"
+    "{{'true' if is_zp_float else 'false'}}>"
    "( MARLIN_KERNEL_PARAMS );"
 )

+# int8 with zero point case (vllm::kU8) is also supported,
+# we don't add it to reduce wheel size.
+SCALAR_TYPES = [
+    "vllm::kU4",
+    "vllm::kU4B8",
+    "vllm::kU8B128",
+    "vllm::kFE4M3fn",
+    "vllm::kFE2M1f",
+]
 THREAD_CONFIGS = [(128, 128, 256), (64, 256, 256), (64, 128, 128)]

 THREAD_M_BLOCKS = [0.5, 1, 2, 3, 4]
-
-QUANT_CONFIGS = [
-    # AWQ-INT4
-    {
-        "b_type": "kU4",
-        "thread_configs": THREAD_CONFIGS,
-        "thread_m_blocks": THREAD_M_BLOCKS,
-        "group_blocks": [-1, 2, 4, 8],
-    },
-    # GPTQ-INT4
-    {
-        "b_type": "kU4B8",
-        "thread_configs": THREAD_CONFIGS,
-        "thread_m_blocks": THREAD_M_BLOCKS,
-        "group_blocks": [-1, 0, 2, 4, 8],
-    },
-    # AWQ-INT8
-    {
-        "b_type": "kU8B128",
-        "thread_configs": THREAD_CONFIGS,
-        "thread_m_blocks": THREAD_M_BLOCKS,
-        "group_blocks": [-1, 0, 2, 4, 8],
-    },
-    # FP8
-    {
-        "b_type": "kFE4M3fn",
-        "thread_configs": THREAD_CONFIGS,
-        "thread_m_blocks": THREAD_M_BLOCKS,
-        "group_blocks": [-1, 8],
-    },
-    # NVFP4
-    {
-        "b_type": "kFE2M1f",
-        "s_type": "kFE4M3fn",
-        "thread_configs": THREAD_CONFIGS,
-        "thread_m_blocks": THREAD_M_BLOCKS,
-        "group_blocks": [1],
-    },
-    # MXFP4
-    {
-        "a_type": ["kBFloat16"],
-        "b_type": "kFE2M1f",
-        "s_type": "kFE8M0fnu",
-        "thread_configs": THREAD_CONFIGS,
-        "thread_m_blocks": THREAD_M_BLOCKS,
-        "group_blocks": [2],
-    },
-    # AWQ-INT4 with INT8 activation
-    {
-        "a_type": ["kS8"],
-        "b_type": "kU4",
-        "thread_configs": THREAD_CONFIGS,
-        "thread_m_blocks": [1, 2, 3, 4],
-        "group_blocks": [-1, 2, 4, 8],
-    },
-    # GPTQ-INT4 with INT8 activation
-    {
-        "a_type": ["kS8"],
-        "b_type": "kU4B8",
-        "thread_configs": THREAD_CONFIGS,
-        "thread_m_blocks": [1, 2, 3, 4],
-        "group_blocks": [-1, 2, 4, 8],
-    },
-    # GPTQ-INT4 with FP8 activation
-    {
-        "a_type": ["kFE4M3fn"],
-        "b_type": "kU4B8",
-        "thread_configs": THREAD_CONFIGS,
-        "thread_m_blocks": [1, 2, 3, 4],
-        "group_blocks": [-1, 2, 4, 8],
-    },
-    # AWQ-INT4 with FP8 activation
-    {
-        "a_type": ["kFE4M3fn"],
-        "b_type": "kU4",
-        "thread_configs": THREAD_CONFIGS,
-        "thread_m_blocks": [1, 2, 3, 4],
-        "group_blocks": [-1, 2, 4, 8],
-    },
-    # MXFP4 with FP8 activation
-    {
-        "a_type": ["kFE4M3fn"],
-        "b_type": "kFE2M1f",
-        "c_type": ["kBFloat16"],
-        "s_type": "kFE8M0fnu",
-        "thread_configs": THREAD_CONFIGS,
-        "thread_m_blocks": [1, 2, 3, 4],
-        "group_blocks": [2],
-    },
-]
+# group_blocks:
+#   = 0 : act order case
+#   = -1 : channelwise quantization
+#   > 0 : group_size=16*group_blocks
+GROUP_BLOCKS = [0, -1, 1, 2, 4, 8]
+DTYPES = ["fp16", "bf16"]


 def remove_old_kernels():
-    for filename in glob.glob(os.path.dirname(__file__) + "/*kernel_*.cu"):
+    for filename in glob.glob(os.path.dirname(__file__) + "/kernel_*.cu"):
        subprocess.call(["rm", "-f", filename])

-    filename = os.path.dirname(__file__) + "/kernel_selector.h"
-    subprocess.call(["rm", "-f", filename])
-

 def generate_new_kernels():
-    result_dict = {}
-
-    for quant_config in QUANT_CONFIGS:
-        c_types = quant_config.get("c_type", ["kFloat16", "kBFloat16"])
-        a_types = quant_config.get("a_type", ["kFloat16", "kBFloat16"])
-        b_type = quant_config["b_type"]
-        all_group_blocks = quant_config["group_blocks"]
-        all_m_blocks = quant_config["thread_m_blocks"]
-        all_thread_configs = quant_config["thread_configs"]
-
-        for a_type, c_type in itertools.product(a_types, c_types):
-            if not SUPPORT_FP8 and a_type == "kFE4M3fn":
-                continue
-            if "16" in a_type and "16" in c_type and a_type != c_type:
-                continue
-            s_type = quant_config.get("s_type", c_type)
-            if (a_type, b_type, c_type) not in result_dict:
-                result_dict[(a_type, b_type, c_type)] = []
-
-            for group_blocks, m_blocks, thread_configs in itertools.product(
-                all_group_blocks, all_m_blocks, all_thread_configs
-            ):
-                thread_k, thread_n, threads = thread_configs
-
-                if threads == 256:
-                    # for small batch (m_blocks == 1),
-                    #     we only need (128, 128, 256)
-                    # for large batch (m_blocks > 1),
-                    #     we only need (64, 256, 256)
-                    if m_blocks <= 1 and (thread_k, thread_n) != (128, 128):
-                        continue
-                    if m_blocks > 1 and (thread_k, thread_n) != (64, 256):
-                        continue
-
-                config = {
-                    "threads": threads,
-                    "s_type": s_type,
-                    "thread_m_blocks": max(m_blocks, 1),
-                    "thread_k_blocks": thread_k // 16,
-                    "thread_n_blocks": thread_n // 16,
-                    "m_block_size_8": "true" if m_blocks == 0.5 else "false",
-                    "stages": "pipe_stages",
-                    "group_blocks": group_blocks,
-                    "is_zp_float": "false",
-                }
-
-                result_dict[(a_type, b_type, c_type)].append(config)
-
-    kernel_selector_str = FILE_HEAD_COMMENT
-
-    for (a_type, b_type, c_type), config_list in result_dict.items():
+    for scalar_type, dtype in itertools.product(SCALAR_TYPES, DTYPES):
        all_template_str_list = []
-        for config in config_list:
-            s_type = config["s_type"]
+
+        for group_blocks, m_blocks, thread_configs in itertools.product(
+            GROUP_BLOCKS, THREAD_M_BLOCKS, THREAD_CONFIGS
+        ):
+            # act order case only support gptq-int4 and gptq-int8
+            if group_blocks == 0 and scalar_type not in [
+                "vllm::kU4B8",
+                "vllm::kU8B128",
+            ]:
+                continue
+            if thread_configs[2] == 256:
+                # for small batch (m_blocks == 1), we only need (128, 128, 256)
+                # for large batch (m_blocks > 1), we only need (64, 256, 256)
+                if m_blocks <= 1 and thread_configs[0] != 128:
+                    continue
+                if m_blocks > 1 and thread_configs[0] != 64:
+                    continue
+
+            # we only support channelwise quantization and group_size == 128
+            # for fp8
+            if scalar_type == "vllm::kFE4M3fn" and group_blocks not in [-1, 8]:
+                continue
+            # nvfp4 only supports group_size == 16
+            # mxfp4 only supports group_size == 32
+            if scalar_type == "vllm::kFE2M1f" and group_blocks not in [1, 2]:
+                continue
+            # other quantization methods don't support group_size = 16
+            if scalar_type != "vllm::kFE2M1f" and group_blocks == 1:
+                continue
+
+            k_blocks = thread_configs[0] // 16
+            n_blocks = thread_configs[1] // 16
+            threads = thread_configs[2]
+
+            c_dtype = "half" if dtype == "fp16" else "nv_bfloat16"
+
+            if scalar_type == "vllm::kFE2M1f" and group_blocks == 1:
+                s_type = "vllm::kFE4M3fn"
+            elif scalar_type == "vllm::kFE2M1f" and group_blocks == 2:
+                s_type = "vllm::kFE8M0fnu"
+                if dtype == "fp16":
+                    # we cannot safely dequantize e8m0 to fp16, so skip this
+                    continue
+            elif dtype == "fp16":
+                s_type = "vllm::kFloat16"
+            elif dtype == "bf16":
+                s_type = "vllm::kBFloat16"
+
            template_str = jinja2.Template(TEMPLATE).render(
-                a_type_id=f"vllm::{a_type}.id()",
-                b_type_id=f"vllm::{b_type}.id()",
-                c_type_id=f"vllm::{c_type}.id()",
-                s_type_id=f"vllm::{s_type}.id()",
-                **config,
+                scalar_t=c_dtype,
+                w_type_id=scalar_type + ".id()",
+                s_type_id=s_type + ".id()",
+                threads=threads,
+                thread_m_blocks=max(m_blocks, 1),
+                thread_n_blocks=n_blocks,
+                thread_k_blocks=k_blocks,
+                m_block_size_8=m_blocks == 0.5,
+                stages="pipe_stages",
+                group_blocks=group_blocks,
+                is_zp_float=False,
            )
+
            all_template_str_list.append(template_str)

-            conditions = [
-                f"a_type == vllm::{a_type}",
-                f"b_type == vllm::{b_type}",
-                f"c_type == vllm::{c_type}",
-                f"s_type == vllm::{s_type}",
-                f"threads == {config['threads']}",
-                f"thread_m_blocks == {config['thread_m_blocks']}",
-                f"thread_n_blocks == {config['thread_n_blocks']}",
-                f"thread_k_blocks == {config['thread_k_blocks']}",
-                f"m_block_size_8 == {config['m_block_size_8']}",
-                f"group_blocks == {config['group_blocks']}",
-                f"is_zp_float == {config['is_zp_float']}",
-            ]
-            conditions = " && ".join(conditions)
-
-            if kernel_selector_str == FILE_HEAD_COMMENT:
-                kernel_selector_str += f"if ({conditions})\n  kernel = "
-            else:
-                kernel_selector_str += f"else if ({conditions})\n  kernel = "
-
-            kernel_template2 = (
-                "Marlin<{{a_type_id}}, {{b_type_id}}, {{c_type_id}}, "
-                "{{s_type_id}}, {{threads}}, {{thread_m_blocks}}, "
-                "{{thread_n_blocks}}, {{thread_k_blocks}}, "
-                "{{m_block_size_8}}, {{stages}}, {{group_blocks}}, "
-                "{{is_zp_float}}>;"
-            )
-
-            kernel_selector_str += (
-                jinja2.Template(kernel_template2).render(
-                    a_type_id=f"vllm::{a_type}.id()",
-                    b_type_id=f"vllm::{b_type}.id()",
-                    c_type_id=f"vllm::{c_type}.id()",
-                    s_type_id=f"vllm::{s_type}.id()",
-                    **config,
-                )
-                + "\n"
-            )
-
        file_content = FILE_HEAD + "\n\n"
        file_content += "\n\n".join(all_template_str_list) + "\n\n}\n"
-        if a_type == "kFE4M3fn":
-            filename = f"sm89_kernel_{a_type[1:]}_{b_type[1:]}_{c_type[1:]}.cu"
-        else:
-            filename = f"sm80_kernel_{a_type[1:]}_{b_type[1:]}_{c_type[1:]}.cu"
-
-        filename = filename.lower()
+        filename = f"kernel_{dtype}_{scalar_type[6:].lower()}.cu"

        with open(os.path.join(os.path.dirname(__file__), filename), "w") as f:
            f.write(file_content)

-    if not SUPPORT_FP8 and kernel_selector_str != FILE_HEAD_COMMENT:
-        kernel_selector_str += (
-            "else if (a_type == vllm::kFE4M3fn)\n"
-            "  TORCH_CHECK(false, "
-            '"marlin kernel with fp8 activation is not built.");'
-        )
-
-    with open(os.path.join(os.path.dirname(__file__), "kernel_selector.h"), "w") as f:
-        f.write(kernel_selector_str)
-

 if __name__ == "__main__":
    remove_old_kernels()
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
Alexander Matveev	ba25edb264	add file Signed-off-by: Alexander Matveev <amatveev@redhat.com>	2025-11-26 13:56:58 +00:00
Barbara SuslovaandAlexander Matveev	9f3812d001	aligning Signed-off-by: Barbara Suslova <barbara.suslova@axel-t.com>	2025-11-25 15:31:48 +00:00
Barbara SuslovaandAlexander Matveev	d8dbb8c0c4	change the logic of passing variables Signed-off-by: Barbara Suslova <barbara.suslova@axel-t.com>	2025-11-25 15:31:46 +00:00
Barbara SuslovaandAlexander Matveev	4d9f1ad714	lint Signed-off-by: Barbara Suslova <barbara.suslova@axel-t.com>	2025-11-25 15:28:46 +00:00
Barbara SuslovaandAlexander Matveev	c13434bbb8	add fusion of shared expert and fused_moe_gate Signed-off-by: Barbara Suslova <barbara.suslova@axel-t.com>	2025-11-25 15:28:17 +00:00
				`@@ -1 +0,0 @@`
				`Meta-Llama-4-Maverick-17B-128E-Instruct-FP8.yaml`