Merge branch 'main' into fix-mtp-dummy-run-assertion

[Perf] Improve default triton fused moe configs (#34846 )
Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
2026-02-23 16:54:20 -08:00 · 2026-02-23 16:01:28 -08:00 · 2026-02-23 16:30:56 -05:00 · 2026-02-23 17:14:51 +00:00 · 2026-02-23 08:38:05 -08:00 · 2026-02-23 08:38:00 -08:00
363 changed files with 40677 additions and 7220 deletions
@@ -10,7 +10,7 @@ steps:
      docker build
      --build-arg max_jobs=16
      --build-arg REMOTE_VLLM=1
-      --build-arg ARG_PYTORCH_ROCM_ARCH='gfx90a;gfx942'
+      --build-arg ARG_PYTORCH_ROCM_ARCH='gfx942;gfx950'
      --build-arg VLLM_BRANCH=$BUILDKITE_COMMIT
      --tag "rocm/vllm-ci:${BUILDKITE_COMMIT}"
      -f docker/Dockerfile.rocm
@@ -14,7 +14,7 @@ BUILDKITE_COMMIT=$3
 aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY"

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

 # push
-docker push "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-cpu
+docker push "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-arm64-cpu
@@ -1,64 +0,0 @@
-#!/bin/bash
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-# Setup script for Prime-RL integration tests
-# This script prepares the environment for running Prime-RL tests with nightly vLLM
-
-set -euo pipefail
-
-SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
-REPO_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)"
-PRIME_RL_REPO="https://github.com/PrimeIntellect-ai/prime-rl.git"
-PRIME_RL_DIR="${REPO_ROOT}/prime-rl"
-
-if command -v rocm-smi &> /dev/null || command -v rocminfo &> /dev/null; then
-    echo "AMD GPU detected. Prime-RL currently only supports NVIDIA. Skipping..."
-    exit 0
-fi
-
-echo "Setting up Prime-RL integration test environment..."
-
-# Clean up any existing Prime-RL directory
-if [ -d "${PRIME_RL_DIR}" ]; then
-    echo "Removing existing Prime-RL directory..."
-    rm -rf "${PRIME_RL_DIR}"
-fi
-
-# Install UV if not available
-if ! command -v uv &> /dev/null; then
-    echo "Installing UV package manager..."
-    curl -LsSf https://astral.sh/uv/install.sh | sh
-    source "$HOME"/.local/bin/env
-fi
-
-# Clone Prime-RL repository at specific branch for reproducible tests
-PRIME_RL_BRANCH="integ-vllm-main"
-echo "Cloning Prime-RL repository at branch: ${PRIME_RL_BRANCH}..."
-git clone --branch "${PRIME_RL_BRANCH}" --single-branch "${PRIME_RL_REPO}" "${PRIME_RL_DIR}"
-cd "${PRIME_RL_DIR}"
-
-echo "Setting up UV project environment..."
-export UV_PROJECT_ENVIRONMENT=/usr/local
-ln -s /usr/bin/python3 /usr/local/bin/python
-
-# Remove vllm pin from pyproject.toml
-echo "Removing vllm pin from pyproject.toml..."
-sed -i '/vllm==/d' pyproject.toml
-
-# Sync Prime-RL dependencies
-echo "Installing Prime-RL dependencies..."
-uv sync --inexact && uv sync --inexact --all-extras
-
-# Verify installation
-echo "Verifying installations..."
-uv run python -c "import vllm; print(f'vLLM version: {vllm.__version__}')"
-uv run python -c "import prime_rl; print('Prime-RL imported successfully')"
-
-echo "Prime-RL integration test environment setup complete!"
-
-echo "Running Prime-RL integration tests..."
-export WANDB_MODE=offline # this makes this test not require a WANDB_API_KEY
-uv run pytest -vs tests/integration/test_rl.py -m gpu
-
-echo "Prime-RL integration tests completed!"
@@ -67,7 +67,7 @@ steps:
  timeout_in_minutes: 30
  mirror_hardwares: [amdexperimental, amdproduction, amdtentative]
  agent_pool: mi325_1
-  grade: Blocking
+  # grade: Blocking
  source_file_dependencies:
  - vllm/
  - tests/test_inputs.py
@@ -549,7 +549,7 @@ steps:
  - tests/samplers
  - tests/conftest.py
  commands:
-    - pytest -v -s -m samplers
+    - pytest -v -s samplers

 - label: LoRA Test %N # 20min each
  timeout_in_minutes: 30
@@ -1107,18 +1107,6 @@ steps:
  commands:
    - pytest -v -s models/quantization

-# This test is used only in PR development phase to test individual models and should never run on main
- label: Custom Models Test
-  mirror_hardwares: [amdexperimental, amdproduction]
-  agent_pool: mi325_1
-  # grade: Blocking
-  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*
-
 - label: Transformers Nightly Models Test
  mirror_hardwares: [amdexperimental]
  agent_pool: mi325_1
@@ -1634,21 +1622,6 @@ steps:
    - 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]
-  agent_pool: mi325_2
-  # grade: Blocking
-  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
-
 ##### EPLB Accuracy Tests #####
 - label: DeepSeek V2-Lite Accuracy
  mirror_hardwares: [amdexperimental, amdproduction]
@@ -1709,7 +1682,6 @@ steps:
  # in /vllm/tools/pre_commit/generate_nightly_torch_test.py
  mirror_hardwares: [amdexperimental, amdproduction, amdtentative]
  agent_pool: mi355_1
-  grade: Blocking
  soft_fail: true
  source_file_dependencies:
  - requirements/nightly_torch_test.txt
@@ -1720,7 +1692,6 @@ steps:
  timeout_in_minutes: 15
  mirror_hardwares: [amdexperimental, amdproduction, amdtentative]
  agent_pool: mi355_1
-  grade: Blocking
  source_file_dependencies:
  - vllm/
  - tests/multimodal
@@ -1733,7 +1704,6 @@ steps:
  timeout_in_minutes: 30
  mirror_hardwares: [amdexperimental, amdproduction, amdtentative]
  agent_pool: mi355_1
-  grade: Blocking
  source_file_dependencies:
  - vllm/
  - tests/test_inputs.py
@@ -1763,7 +1733,6 @@ steps:
  timeout_in_minutes: 20
  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_1
-  # grade: Blocking
  source_file_dependencies:
  - tests/standalone_tests/python_only_compile.sh
  - setup.py
@@ -1774,7 +1743,6 @@ steps:
  timeout_in_minutes: 30
  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi355_1
-  # grade: Blocking
  fast_check: true
  torch_nightly: true
  source_file_dependencies:
@@ -1791,7 +1759,6 @@ steps:
 - label: Entrypoints Unit Tests # 5min
  mirror_hardwares: [amdexperimental, amdproduction, amdtentative]
  agent_pool: mi355_1
-  grade: Blocking
  timeout_in_minutes: 10
  working_dir: "/vllm-workspace/tests"
  fast_check: true
@@ -1806,7 +1773,6 @@ steps:
  timeout_in_minutes: 40
  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi355_1
-  # grade: Blocking
  working_dir: "/vllm-workspace/tests"
  fast_check: true
  torch_nightly: true
@@ -1824,7 +1790,6 @@ steps:
  timeout_in_minutes: 130
  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_1
-  # grade: Blocking
  working_dir: "/vllm-workspace/tests"
  fast_check: true
  torch_nightly: true
@@ -1841,7 +1806,6 @@ steps:
  timeout_in_minutes: 50
  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_1
-  # grade: Blocking
  working_dir: "/vllm-workspace/tests"
  fast_check: true
  torch_nightly: true
@@ -1860,7 +1824,6 @@ steps:
  timeout_in_minutes: 50
  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_1
-  # grade: Blocking
  working_dir: "/vllm-workspace/tests"
  fast_check: true
  torch_nightly: true
@@ -1875,7 +1838,6 @@ steps:
  timeout_in_minutes: 50
  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_1
-  # grade: Blocking
  working_dir: "/vllm-workspace/tests"
  fast_check: true
  torch_nightly: true
@@ -1890,7 +1852,6 @@ steps:
  timeout_in_minutes: 50
  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_4
-  # grade: Blocking
  working_dir: "/vllm-workspace/tests"
  num_gpus: 4
  source_file_dependencies:
@@ -1952,7 +1913,6 @@ steps:
  timeout_in_minutes: 10
  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_8
-  # grade: Blocking
  gpu: h100
  num_gpus: 8
  working_dir: "/vllm-workspace/tests"
@@ -1973,7 +1933,6 @@ steps:
 - label: EPLB Algorithm Test # 5min
  mirror_hardwares: [amdexperimental, amdproduction, amdtentative]
  agent_pool: mi355_1
-  grade: Blocking
  timeout_in_minutes: 15
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -1985,7 +1944,6 @@ steps:
 - label: EPLB Execution Test # 10min
  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi355_4
-  # grade: Blocking
  timeout_in_minutes: 20
  working_dir: "/vllm-workspace/tests"
  num_gpus: 4
@@ -2000,7 +1958,6 @@ steps:
  timeout_in_minutes: 20
  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi355_2
-  # grade: Blocking
  num_gpus: 2
  source_file_dependencies:
  - vllm/
@@ -2020,7 +1977,6 @@ steps:
  timeout_in_minutes: 20
  mirror_hardwares: [amdexperimental, amdproduction, amdtentative]
  agent_pool: mi355_1
-  grade: Blocking
  source_file_dependencies:
  - vllm/
  - tests/test_regression
@@ -2033,7 +1989,6 @@ steps:
  timeout_in_minutes: 15
  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi355_1
-  # grade: Blocking
  source_file_dependencies:
  - vllm/
  - tests/engine
@@ -2050,7 +2005,6 @@ steps:
  # The test uses 4 GPUs, but we schedule it on 8-GPU machines for stability.
  # See discussion here: https://github.com/vllm-project/vllm/pull/31040
  agent_pool: mi355_8
-  # grade: Blocking
  source_file_dependencies:
    - vllm/
    - tests/v1
@@ -2064,7 +2018,6 @@ steps:
  timeout_in_minutes: 50
  mirror_hardwares: [amdexperimental, amdproduction, amdtentative]
  agent_pool: mi355_1
-  grade: Blocking
  source_file_dependencies:
    - vllm/
    - tests/v1
@@ -2075,7 +2028,6 @@ steps:
  timeout_in_minutes: 60
  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_1
-  # grade: Blocking
  source_file_dependencies:
    - vllm/
    - tests/v1
@@ -2103,7 +2055,6 @@ steps:
 - label: V1 Test attention (H100) # 10min
  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_1
-  # grade: Blocking
  timeout_in_minutes: 30
  gpu: h100
  source_file_dependencies:
@@ -2143,7 +2094,6 @@ steps:
 - label: V1 Test others (CPU) # 5 mins
  mirror_hardwares: [amdexperimental, amdproduction, amdtentative]
  agent_pool: mi355_1
-  grade: Blocking
  source_file_dependencies:
    - vllm/
    - tests/v1
@@ -2161,7 +2111,6 @@ steps:
  timeout_in_minutes: 45
  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi355_1
-  # grade: Blocking
  working_dir: "/vllm-workspace/examples"
  source_file_dependencies:
  - vllm/entrypoints
@@ -2196,7 +2145,6 @@ steps:
  timeout_in_minutes: 15
  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi355_1
-  # grade: Blocking
  source_file_dependencies:
  - vllm/
  - tests/cuda
@@ -2208,20 +2156,18 @@ steps:
  timeout_in_minutes: 75
  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_1
-  # grade: Blocking
  source_file_dependencies:
  - vllm/model_executor/layers
  - vllm/sampling_metadata.py
  - tests/samplers
  - tests/conftest.py
  commands:
-    - pytest -v -s -m samplers
+    - pytest -v -s samplers

 - label: LoRA Test %N # 20min each
  timeout_in_minutes: 30
  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_1
-  # grade: Blocking
  source_file_dependencies:
  - vllm/lora
  - tests/lora
@@ -2242,7 +2188,6 @@ steps:
  timeout_in_minutes: 30
  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi355_1
-  # grade: Blocking
  torch_nightly: true
  source_file_dependencies:
    - vllm/
@@ -2259,7 +2204,6 @@ steps:
  timeout_in_minutes: 30
  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi355_1
-  # grade: Blocking
  torch_nightly: true
  source_file_dependencies:
  - vllm/
@@ -2305,7 +2249,6 @@ steps:
  timeout_in_minutes: 75
  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi355_1
-  # grade: Blocking
  source_file_dependencies:
  - csrc/
  - tests/kernels/core
@@ -2317,7 +2260,6 @@ steps:
  timeout_in_minutes: 35
  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi355_1
-  # grade: Blocking
  source_file_dependencies:
  - csrc/attention/
  - vllm/v1/attention
@@ -2332,7 +2274,6 @@ steps:
  timeout_in_minutes: 90
  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_1
-  # grade: Blocking
  source_file_dependencies:
  - csrc/quantization/
  - vllm/model_executor/layers/quantization
@@ -2345,7 +2286,6 @@ steps:
  timeout_in_minutes: 60
  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi355_1
-  # grade: Blocking
  source_file_dependencies:
  - csrc/quantization/cutlass_w8a8/moe/
  - csrc/moe/
@@ -2362,7 +2302,6 @@ steps:
  timeout_in_minutes: 45
  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi355_1
-  # grade: Blocking
  source_file_dependencies:
  - csrc/mamba/
  - tests/kernels/mamba
@@ -2406,7 +2345,6 @@ steps:
  torch_nightly: true
  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi355_1
-  # grade: Blocking
  source_file_dependencies:
  - vllm/engine/arg_utils.py
  - vllm/config/model.py
@@ -2423,7 +2361,6 @@ steps:
  timeout_in_minutes: 20
  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi355_1
-  # grade: Blocking
  working_dir: "/vllm-workspace/.buildkite"
  source_file_dependencies:
  - benchmarks/
@@ -2434,7 +2371,6 @@ steps:
  timeout_in_minutes: 20
  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi355_1
-  # grade: Blocking
  source_file_dependencies:
  - vllm/
  - tests/benchmarks/
@@ -2445,7 +2381,6 @@ steps:
  timeout_in_minutes: 90
  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi355_1
-  # grade: Blocking
  source_file_dependencies:
  - csrc/
  - vllm/model_executor/layers/quantization
@@ -2466,7 +2401,6 @@ steps:
  timeout_in_minutes: 75
  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_1
-  # grade: Blocking
  source_file_dependencies:
  - csrc/
  - vllm/model_executor/layers/quantization
@@ -2478,7 +2412,6 @@ steps:
  timeout_in_minutes: 15
  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi355_1
-  # grade: Blocking
  source_file_dependencies:
  - csrc/
  - vllm/entrypoints/openai/
@@ -2495,7 +2428,6 @@ steps:
  timeout_in_minutes: 45
  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi355_1
-  # grade: Blocking
  torch_nightly: true
  source_file_dependencies:
  - vllm/
@@ -2508,7 +2440,6 @@ steps:
  timeout_in_minutes: 45
  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi355_1
-  # grade: Blocking
  torch_nightly: true
  source_file_dependencies:
  - vllm/model_executor/models/
@@ -2528,7 +2459,6 @@ steps:
  timeout_in_minutes: 45
  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_1
-  # grade: Blocking
  torch_nightly: true
  source_file_dependencies:
  - vllm/
@@ -2541,7 +2471,6 @@ steps:
 - label: Basic Models Test (Other CPU) # 5min
  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi355_1
-  # grade: Blocking
  timeout_in_minutes: 10
  torch_nightly: true
  source_file_dependencies:
@@ -2556,7 +2485,6 @@ steps:
  timeout_in_minutes: 25
  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi355_1
-  # grade: Blocking
  torch_nightly: true
  source_file_dependencies:
  - vllm/
@@ -2570,7 +2498,6 @@ steps:
  timeout_in_minutes: 45
  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_1
-  # grade: Blocking
  torch_nightly: true
  source_file_dependencies:
  - vllm/model_executor/models/
@@ -2591,7 +2518,6 @@ steps:
  timeout_in_minutes: 75
  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_1
-  # grade: Blocking
  torch_nightly: true
  source_file_dependencies:
  - vllm/
@@ -2612,7 +2538,6 @@ steps:
  timeout_in_minutes: 110
  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_1
-  # grade: Blocking
  optional: true
  source_file_dependencies:
  - vllm/
@@ -2628,7 +2553,6 @@ steps:
  timeout_in_minutes: 110
  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_1
-  # grade: Blocking
  optional: true
  source_file_dependencies:
  - vllm/
@@ -2640,7 +2564,6 @@ steps:
  timeout_in_minutes: 50
  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_1
-  # grade: Blocking
  optional: true
  source_file_dependencies:
  - vllm/
@@ -2676,7 +2599,6 @@ steps:
  timeout_in_minutes: 60
  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_1
-  # grade: Blocking
  source_file_dependencies:
  - vllm/
  - tests/models/multimodal
@@ -2688,7 +2610,6 @@ steps:
  timeout_in_minutes: 100
  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_1
-  # grade: Blocking
  torch_nightly: true
  source_file_dependencies:
  - vllm/
@@ -2706,7 +2627,6 @@ steps:
  timeout_in_minutes: 10
  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi355_1
-  # grade: Blocking
  working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
  source_file_dependencies:
  - vllm/multimodal/
@@ -2721,7 +2641,6 @@ steps:
  timeout_in_minutes: 120
  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_1
-  # grade: Blocking
  optional: true
  source_file_dependencies:
  - vllm/
@@ -2736,7 +2655,6 @@ steps:
  timeout_in_minutes: 120
  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_1
-  # grade: Blocking
  optional: true
  source_file_dependencies:
  - vllm/
@@ -2751,7 +2669,6 @@ steps:
  timeout_in_minutes: 150
  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_1
-  # grade: Blocking
  optional: true
  source_file_dependencies:
  - vllm/
@@ -2766,29 +2683,15 @@ steps:
  timeout_in_minutes: 60
  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi355_1
-  # grade: Blocking
  source_file_dependencies:
  - vllm/model_executor/layers/quantization
  - tests/models/quantization
  commands:
    - pytest -v -s models/quantization

-# This test is used only in PR development phase to test individual models and should never run on main
- label: Custom Models Test
-  mirror_hardwares: [amdexperimental, amdproduction]
-  agent_pool: mi355_1
-  # grade: Blocking
-  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*
-
 - label: Transformers Nightly Models Test
  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_1
-  # grade: Blocking
  working_dir: "/vllm-workspace/"
  optional: true
  commands:
@@ -2927,7 +2830,6 @@ steps:
  timeout_in_minutes: 20
  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi355_2
-  # grade: Blocking
  working_dir: "/vllm-workspace/tests"
  num_gpus: 2
  source_file_dependencies:
@@ -2943,7 +2845,6 @@ steps:
  timeout_in_minutes: 30
  mirror_hardwares: [amdexperimental, amdmultinode]
  agent_pool: mi355_4
-  # grade: Blocking
  working_dir: "/vllm-workspace/tests"
  num_gpus: 2
  num_nodes: 2
@@ -2970,7 +2871,6 @@ steps:
  timeout_in_minutes: 90
  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_2
-  # grade: Blocking
  working_dir: "/vllm-workspace/tests"
  num_gpus: 2
  source_file_dependencies:
@@ -3010,7 +2910,6 @@ steps:
  timeout_in_minutes: 50
  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_2
-  # grade: Blocking
  working_dir: "/vllm-workspace/tests"
  num_gpus: 2
  source_file_dependencies:
@@ -3032,7 +2931,6 @@ steps:
  timeout_in_minutes: 60
  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi355_2
-  # grade: Blocking
  working_dir: "/vllm-workspace/tests"
  num_gpus: 2
  source_file_dependencies:
@@ -3066,7 +2964,6 @@ steps:
  timeout_in_minutes: 60
  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi355_4
-  # grade: Blocking
  working_dir: "/vllm-workspace/tests"
  num_gpus: 4
  source_file_dependencies:
@@ -3083,7 +2980,6 @@ steps:
  timeout_in_minutes: 30
  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi355_4
-  # grade: Blocking
  num_gpus: 4
  source_file_dependencies:
  - vllm/lora
@@ -3108,7 +3004,6 @@ steps:
  timeout_in_minutes: 45
  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi355_2
-  # grade: Blocking
  working_dir: "/vllm-workspace/tests"
  num_gpus: 2
  optional: true
@@ -3121,7 +3016,6 @@ steps:
 - label: Weight Loading Multiple GPU Test - Large Models # optional
  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_2
-  # grade: Blocking
  working_dir: "/vllm-workspace/tests"
  num_gpus: 2
  optional: true
@@ -3134,7 +3028,6 @@ steps:
 - label: NixlConnector PD accuracy tests (Distributed) # 30min
  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi355_4
-  # grade: Blocking
  timeout_in_minutes: 30
  working_dir: "/vllm-workspace/tests"
  num_gpus: 4
@@ -3148,7 +3041,6 @@ steps:
 - label: DP EP NixlConnector PD accuracy tests (Distributed) # 15min
  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi355_4
-  # grade: Blocking
  timeout_in_minutes: 15
  working_dir: "/vllm-workspace/tests"
  num_gpus: 4
@@ -3165,7 +3057,6 @@ steps:
 - label: Distributed Tests (A100) # optional
  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_4
-  # grade: Blocking
  gpu: a100
  optional: true
  num_gpus: 4
@@ -3188,7 +3079,6 @@ steps:
  optional: true
  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_4
-  # grade: Blocking
  num_gpus: 4
  working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
  source_file_dependencies:
@@ -3204,7 +3094,6 @@ steps:
  optional: true
  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_4
-  # grade: Blocking
  num_gpus: 4
  working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
  source_file_dependencies:
@@ -3219,7 +3108,6 @@ steps:
 - label: Distributed Tests (H200) # optional
  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_2
-  # grade: Blocking
  gpu: h200
  optional: true
  working_dir: "/vllm-workspace/"
@@ -3254,7 +3142,6 @@ steps:
  timeout_in_minutes: 20
  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi355_1
-  # grade: Blocking
  source_file_dependencies:
  - csrc/
  - vllm/model_executor/layers/quantization
@@ -3264,7 +3151,6 @@ steps:
 - label: LM Eval Large Models (4 Card)
  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi355_4
-  # grade: Blocking
  gpu: a100
  optional: true
  num_gpus: 4
@@ -3300,26 +3186,10 @@ steps:
    - 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]
-  agent_pool: mi355_2
-  # grade: Blocking
-  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
-
 ##### EPLB Accuracy Tests #####
 - label: DeepSeek V2-Lite Accuracy
  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi355_4
-  # grade: Blocking
  timeout_in_minutes: 60
  gpu: h100
  optional: true
@@ -3331,7 +3201,6 @@ steps:
 - label: Qwen3-30B-A3B-FP8-block Accuracy (H100)
  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi355_4
-  # grade: Blocking
  timeout_in_minutes: 60
  gpu: h100
  optional: true
@@ -3354,7 +3223,6 @@ steps:
  timeout_in_minutes: 60
  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_4
-  # grade: Blocking
  optional: true
  num_gpus: 4
  working_dir: "/vllm-workspace"
@@ -104,7 +104,6 @@ steps:
  # NEW rlhf examples
  - cd new_weight_syncing
  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf.py
-  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_async_new_apis.py

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

@@ -209,7 +209,7 @@ steps:
    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
    - CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh

- label: Pipeline + Context Parallelism (4 GPUs))
+- label: Pipeline + Context Parallelism (4 GPUs)
  timeout_in_minutes: 60
  working_dir: "/vllm-workspace/tests"
  num_devices: 4
@@ -28,16 +28,3 @@ steps:
  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
-  soft_fail: true
-  num_devices: 2
-  working_dir: "/vllm-workspace"
-  source_file_dependencies:
-  - vllm/
-  - .buildkite/scripts/run-prime-rl-test.sh
-  commands:
-    - nvidia-smi
-    - bash .buildkite/scripts/run-prime-rl-test.sh
@@ -147,6 +147,19 @@ steps:
  - pytest -v -s transformers_utils
  - pytest -v -s config

+- label: GPT-OSS Eval (H100)
+  timeout_in_minutes: 60
+  working_dir: "/vllm-workspace/"
+  device: h100
+  optional: true
+  source_file_dependencies:
+  - tests/evals/gpt_oss
+  - vllm/model_executor/models/gpt_oss.py
+  - vllm/model_executor/layers/quantization/mxfp4.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: GPT-OSS Eval (B200)
  timeout_in_minutes: 60
  working_dir: "/vllm-workspace/"
@@ -18,4 +18,4 @@ steps:
      depends_on:
      - image-build-amd
      commands:
-      - pytest -v -s -m samplers
+      - pytest -v -s samplers
@@ -55,7 +55,7 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
 /vllm/v1/kv_cache_interface.py @heheda12345
 /vllm/v1/kv_offload @ApostaC @orozery
 /vllm/v1/worker/gpu/kv_connector.py @orozery
-/vllm/v1/worker/kv_connector_model_runner_mixin.py @orozery
+/vllm/v1/worker/kv_connector_model_runner_mixin.py @orozery @NickLucche

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

 message(STATUS "Enabling moe extension.")
@@ -13,6 +13,7 @@ from torch.utils.benchmark import Measurement as TMeasurement
 from tqdm import tqdm

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


+@default_vllm_config()
 def main():
    torch.set_default_device("cuda")
    bench_params = get_bench_params()
@@ -7,6 +7,7 @@ import itertools
 import torch

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


+@default_vllm_config()
 def benchmark_activation(
    batch_size: int,
    seq_len: int,
@@ -8,6 +8,7 @@ os.environ["VLLM_USE_DEEP_GEMM"] = "0"

 import torch

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


+@default_vllm_config()
 def build_w8a8_block_fp8_runner(M, N, K, block_size, device, use_cutlass):
    """Build runner function for w8a8 block fp8 matmul."""
    factor_for_scale = 1e-2
@@ -408,18 +408,18 @@ def run_benchmarks(

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

-    # Create RMSNorm and QuantFP8 layers once for native benchmarks
-
    if "none" in quant_modes:
        # Standard AllReduce + RMSNorm
+        # Re-create VllmFusedAllreduce per config so CustomOp binds the
+        # correct forward method (native vs custom kernel).
        for custom_op in ["-rms_norm", "+rms_norm"]:
            with set_current_vllm_config(
                VllmConfig(compilation_config=CompilationConfig(custom_ops=[custom_op]))
            ):
                try:
+                    vllm_fused_allreduce = VllmFusedAllreduce(hidden_dim, dtype)
                    suffix = (
                        "_custom_rms_norm" if "+" in custom_op else "_native_rms_norm"
                    )
@@ -438,6 +438,7 @@ def run_benchmarks(
            VllmConfig(compilation_config=CompilationConfig(custom_ops=["-rms_norm"]))
        ):
            try:
+                vllm_fused_allreduce = VllmFusedAllreduce(hidden_dim, dtype)
                standard_allreduce_rmsnorm_native_compiled = torch.compile(
                    vllm_fused_allreduce.allreduce_rmsnorm,
                    fullgraph=True,
@@ -482,7 +483,7 @@ def run_benchmarks(
                "_custom_rms_norm" if "+" in rms_norm_custom_op else "_native_rms_norm"
            )
            for quant_fp8_custom_op in ["-quant_fp8", "+quant_fp8"]:
-                suffix += (
+                op_suffix = suffix + (
                    "_custom_quant_fp8"
                    if "+" in quant_fp8_custom_op
                    else "_native_quant_fp8"
@@ -495,16 +496,17 @@ def run_benchmarks(
                    )
                ):
                    try:
+                        vllm_fused_allreduce = VllmFusedAllreduce(hidden_dim, dtype)
                        time_ms = benchmark_operation(
                            vllm_fused_allreduce.allreduce_rmsnorm_fp8_quant,
                            input_tensor,
                            residual=residual,
                            scale_factor=scale_fp8,
                        )
-                        results[f"standard_allreduce{suffix}"] = time_ms
+                        results[f"standard_allreduce{op_suffix}"] = time_ms
                    except Exception as e:
                        logger.error("Standard AllReduce+RMSNorm+FP8 failed: %s", e)
-                        results[f"standard_allreduce{suffix}"] = float("inf")
+                        results[f"standard_allreduce{op_suffix}"] = float("inf")

        # Standard AllReduce + RMSNorm + FP8 Quant Native Compiled
        with set_current_vllm_config(
@@ -515,6 +517,7 @@ def run_benchmarks(
            )
        ):
            try:
+                vllm_fused_allreduce = VllmFusedAllreduce(hidden_dim, dtype)
                standard_allreduce_rmsnorm_fp8_quant_native_compiled = torch.compile(
                    vllm_fused_allreduce.allreduce_rmsnorm_fp8_quant,
                    fullgraph=True,
@@ -580,6 +583,7 @@ def run_benchmarks(
                )
            ):
                try:
+                    vllm_fused_allreduce = VllmFusedAllreduce(hidden_dim, dtype)
                    time_ms = benchmark_operation(
                        vllm_fused_allreduce.allreduce_rmsnorm_fp4_quant,
                        input_tensor,
@@ -598,6 +602,7 @@ def run_benchmarks(
            VllmConfig(compilation_config=CompilationConfig(custom_ops=["-rms_norm"]))
        ):
            try:
+                vllm_fused_allreduce = VllmFusedAllreduce(hidden_dim, dtype)
                standard_allreduce_rmsnorm_fp4_quant_native_compiled = torch.compile(
                    vllm_fused_allreduce.allreduce_rmsnorm_fp4_quant,
                    fullgraph=True,
@@ -5,12 +5,14 @@ import time

 import torch

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


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

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


+@default_vllm_config()
 def benchmark_mrope(
    model_name: str,
    num_tokens: int,
@@ -7,6 +7,7 @@ from unittest.mock import patch
 import pandas as pd
 import torch

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


+@default_vllm_config()
 def benchmark_quantization(
    batch_size,
    hidden_size,
@@ -5,6 +5,7 @@ import itertools

 import torch

+from vllm.benchmarks.lib.utils import default_vllm_config
 from vllm.model_executor.layers.rotary_embedding import get_rope
 from vllm.triton_utils import triton
 from vllm.utils.argparse_utils import FlexibleArgumentParser
@@ -29,6 +30,7 @@ def get_benchmark(head_size, rotary_dim, is_neox_style, device):
            args={},
        )
    )
+    @default_vllm_config()
    def benchmark(batch_size, seq_len, num_heads, provider):
        dtype = torch.bfloat16
        max_position = 8192
@@ -14,7 +14,8 @@ struct alignas(32) u32x8_t {
 };

 __device__ __forceinline__ void ld256(u32x8_t& val, const u32x8_t* ptr) {
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && \
+    defined(CUDA_VERSION) && CUDA_VERSION >= 12090
  asm volatile("ld.global.nc.v8.u32 {%0,%1,%2,%3,%4,%5,%6,%7}, [%8];\n"
               : "=r"(val.u0), "=r"(val.u1), "=r"(val.u2), "=r"(val.u3),
                 "=r"(val.u4), "=r"(val.u5), "=r"(val.u6), "=r"(val.u7)
@@ -35,7 +36,8 @@ __device__ __forceinline__ void ld256(u32x8_t& val, const u32x8_t* ptr) {
 }

 __device__ __forceinline__ void st256(u32x8_t& val, u32x8_t* ptr) {
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && \
+    defined(CUDA_VERSION) && CUDA_VERSION >= 12090
  asm volatile("st.global.v8.u32 [%0], {%1,%2,%3,%4,%5,%6,%7,%8};\n"
               :
               : "l"(ptr), "r"(val.u0), "r"(val.u1), "r"(val.u2), "r"(val.u3),
@@ -0,0 +1,291 @@
+/*
+ * Adapted from SGLang's sgl-kernel implementation, which was adapted from
+ * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu
+ * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp
+ *
+ * Copyright (c) 2019-2023, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+
+#include "dsv3_router_gemm_utils.h"
+
+// Custom FMA implementation using PTX assembly instructions
+__device__ __forceinline__ void fma(float2& d, float2 const& a, float2 const& b,
+                                    float2 const& c) {
+  asm volatile("fma.rn.f32x2 %0, %1, %2, %3;\n"
+               : "=l"(reinterpret_cast<uint64_t&>(d))
+               : "l"(reinterpret_cast<uint64_t const&>(a)),
+                 "l"(reinterpret_cast<uint64_t const&>(b)),
+                 "l"(reinterpret_cast<uint64_t const&>(c)));
+}
+
+// Convert 8 bfloat16 values from a uint4 to float array - optimized conversion
+template <int VPT>
+__device__ __forceinline__ void bf16_uint4_to_float8(uint4 const& vec,
+                                                     float* dst) {
+  __nv_bfloat16* bf16_ptr =
+      reinterpret_cast<__nv_bfloat16*>(const_cast<uint4*>(&vec));
+
+#pragma unroll
+  for (int i = 0; i < VPT; i++) {
+    dst[i] = __bfloat162float(bf16_ptr[i]);
+  }
+}
+
+template <typename T, int kBlockSize, int VPT, int kNumTokens, int kNumExperts,
+          int kHiddenDim>
+__global__ __launch_bounds__(128, 1) void router_gemm_kernel_bf16_output(
+    __nv_bfloat16* out, T const* mat_a, T const* mat_b) {
+  // Each block handles one expert column
+  int const n_idx = blockIdx.x;
+  int const tid = threadIdx.x;
+  constexpr int kWarpSize = 32;
+  constexpr int kNumWarps = kBlockSize / kWarpSize;
+  // Constants for this kernel
+  constexpr int k_elems_per_k_iteration = VPT * kBlockSize;
+  constexpr int k_iterations =
+      kHiddenDim / k_elems_per_k_iteration;  // Total K iterations
+
+  // Initialize accumulators for all M rows
+  float acc[kNumTokens] = {};
+
+  // Shared memory for warp-level reduction
+  __shared__ float sm_reduction[kNumTokens][kNumWarps];  // kNumWarps
+
+  // B matrix is in column-major order, so we can directly load a column for the
+  // n_idx expert
+  T const* b_col = mat_b + n_idx * kHiddenDim;
+
+  // Pre-compute k_base values for each iteration to help compiler optimize
+  int k_bases[k_iterations];
+#pragma unroll
+  for (int ki = 0; ki < k_iterations; ki++) {
+    k_bases[ki] = ki * k_elems_per_k_iteration + tid * VPT;
+  }
+
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
+  asm volatile("griddepcontrol.wait;");
+#endif
+
+  // Process the GEMM in chunks
+  for (int ki = 0; ki < k_iterations; ki++) {
+    int const k_base = k_bases[ki];
+
+    // Load B matrix values using vector load (8 bf16 values)
+    uint4 b_vec = *reinterpret_cast<uint4 const*>(b_col + k_base);
+
+    // Convert B values to float
+    float b_float[VPT];
+    bf16_uint4_to_float8<VPT>(b_vec, b_float);
+
+// Process each token
+#pragma unroll
+    for (int m_idx = 0; m_idx < kNumTokens; m_idx++) {
+      // Load both rows of A matrix using vector loads
+      uint4 a_vec = *reinterpret_cast<uint4 const*>(
+          mat_a + (m_idx * kHiddenDim) + k_base);
+
+      // Convert A values to float
+      float a_float[VPT];
+      bf16_uint4_to_float8<VPT>(a_vec, a_float);
+
+// Process elements in this chunk
+#pragma unroll
+      for (int k = 0; k < VPT; k++) {
+        float a = a_float[k];
+        float b = b_float[k];
+        acc[m_idx] += a * b;
+      }
+    }
+  }
+
+  // Perform warp-level reduction
+  int const warpSize = 32;
+  int const warpId = tid / warpSize;
+  int const laneId = tid % warpSize;
+
+  // Register for warp-level reduction results
+  float warp_result[kNumTokens];
+
+#pragma unroll
+  for (int m_idx = 0; m_idx < kNumTokens; m_idx++) {
+    warp_result[m_idx] = acc[m_idx];
+  }
+
+// Perform warp-level reduction using optimized butterfly pattern
+#pragma unroll
+  for (int m = 0; m < kNumTokens; m++) {
+    float sum = warp_result[m];
+
+    // Butterfly reduction pattern
+    sum += __shfl_xor_sync(0xffffffff, sum, 16);
+    sum += __shfl_xor_sync(0xffffffff, sum, 8);
+    sum += __shfl_xor_sync(0xffffffff, sum, 4);
+    sum += __shfl_xor_sync(0xffffffff, sum, 2);
+    sum += __shfl_xor_sync(0xffffffff, sum, 1);
+
+    // Only the first thread in each warp stores to shared memory
+    if (laneId == 0) {
+      sm_reduction[m][warpId] = sum;
+    }
+  }
+
+  __syncthreads();
+
+  // Final reduction across warps (only first thread)
+  if (tid == 0) {
+#pragma unroll
+    for (int m = 0; m < kNumTokens; m++) {
+      float final_sum = 0.0f;
+
+// Sum across the kNumWarps
+#pragma unroll
+      for (int w = 0; w < kNumWarps; w++) {
+        final_sum += sm_reduction[m][w];
+      }
+
+      // Write final result
+      out[m * kNumExperts + n_idx] = __float2bfloat16(final_sum);
+    }
+  }
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
+  asm volatile("griddepcontrol.launch_dependents;");
+#endif
+}
+
+template <typename T, int kNumTokens, int kNumExperts, int kHiddenDim>
+void invokeRouterGemmBf16Output(__nv_bfloat16* output, T const* mat_a,
+                                T const* mat_b, cudaStream_t stream) {
+  constexpr int VPT = 16 / sizeof(T);
+  constexpr int kBlockSize = 128;
+  cudaLaunchConfig_t config;
+  config.gridDim = kNumExperts;
+  config.blockDim = kBlockSize;
+  config.dynamicSmemBytes = 0;
+  config.stream = stream;
+  cudaLaunchAttribute attrs[1];
+  attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+  attrs[0].val.programmaticStreamSerializationAllowed = getEnvEnablePDL();
+  config.numAttrs = 1;
+  config.attrs = attrs;
+  cudaLaunchKernelEx(
+      &config,
+      router_gemm_kernel_bf16_output<T, kBlockSize, VPT, kNumTokens,
+                                     kNumExperts, kHiddenDim>,
+      output, mat_a, mat_b);
+}
+
+// Template instantiations for DEFAULT_NUM_EXPERTS experts
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 1, 256, 7168>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 2, 256, 7168>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 3, 256, 7168>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 4, 256, 7168>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 5, 256, 7168>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 6, 256, 7168>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 7, 256, 7168>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 8, 256, 7168>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 9, 256, 7168>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 10, 256, 7168>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 11, 256, 7168>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 12, 256, 7168>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 13, 256, 7168>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 14, 256, 7168>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 15, 256, 7168>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 16, 256, 7168>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+// Template instantiations for KIMI_K2_NUM_EXPERTS experts
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 1, 384, 7168>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 2, 384, 7168>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 3, 384, 7168>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 4, 384, 7168>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 5, 384, 7168>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 6, 384, 7168>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 7, 384, 7168>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 8, 384, 7168>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 9, 384, 7168>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 10, 384, 7168>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 11, 384, 7168>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 12, 384, 7168>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 13, 384, 7168>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 14, 384, 7168>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 15, 384, 7168>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 16, 384, 7168>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
@@ -0,0 +1,163 @@
+/*
+ * Adapted from SGLang's sgl-kernel implementation, which was adapted from
+ * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu
+ * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp
+ *
+ * Copyright (c) 2019-2023, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+
+#include "dsv3_router_gemm_utils.h"
+
+static constexpr int DEFAULT_NUM_EXPERTS = 256;
+static constexpr int KIMI_K2_NUM_EXPERTS = 384;
+static constexpr int DEFAULT_HIDDEN_DIM = 7168;
+
+template <typename T, int kNumTokens, int kNumExperts, int kHiddenDim>
+void invokeRouterGemmFloatOutput(float* output, T const* mat_a, T const* mat_b,
+                                 cudaStream_t stream);
+
+template <typename T, int kNumTokens, int kNumExperts, int kHiddenDim>
+void invokeRouterGemmBf16Output(__nv_bfloat16* output, T const* mat_a,
+                                T const* mat_b, cudaStream_t stream);
+
+template <int kBegin, int kEnd, int kNumExperts, int kHiddenDim>
+struct LoopUnroller {
+  static void unroll_float_output(int num_tokens, float* output,
+                                  __nv_bfloat16 const* input,
+                                  __nv_bfloat16 const* weights,
+                                  cudaStream_t stream) {
+    if (num_tokens == kBegin) {
+      invokeRouterGemmFloatOutput<__nv_bfloat16, kBegin, kNumExperts,
+                                  kHiddenDim>(output, input, weights, stream);
+    } else {
+      LoopUnroller<kBegin + 1, kEnd, kNumExperts,
+                   kHiddenDim>::unroll_float_output(num_tokens, output, input,
+                                                    weights, stream);
+    }
+  }
+
+  static void unroll_bf16_output(int num_tokens, __nv_bfloat16* output,
+                                 __nv_bfloat16 const* input,
+                                 __nv_bfloat16 const* weights,
+                                 cudaStream_t stream) {
+    if (num_tokens == kBegin) {
+      invokeRouterGemmBf16Output<__nv_bfloat16, kBegin, kNumExperts,
+                                 kHiddenDim>(output, input, weights, stream);
+    } else {
+      LoopUnroller<kBegin + 1, kEnd, kNumExperts,
+                   kHiddenDim>::unroll_bf16_output(num_tokens, output, input,
+                                                   weights, stream);
+    }
+  }
+};
+
+template <int kEnd, int kNumExperts, int kHiddenDim>
+struct LoopUnroller<kEnd, kEnd, kNumExperts, kHiddenDim> {
+  static void unroll_float_output(int num_tokens, float* output,
+                                  __nv_bfloat16 const* input,
+                                  __nv_bfloat16 const* weights,
+                                  cudaStream_t stream) {
+    if (num_tokens == kEnd) {
+      invokeRouterGemmFloatOutput<__nv_bfloat16, kEnd, kNumExperts, kHiddenDim>(
+          output, input, weights, stream);
+    } else {
+      throw std::invalid_argument("Invalid num_tokens, only supports 1 to 16");
+    }
+  }
+
+  static void unroll_bf16_output(int num_tokens, __nv_bfloat16* output,
+                                 __nv_bfloat16 const* input,
+                                 __nv_bfloat16 const* weights,
+                                 cudaStream_t stream) {
+    if (num_tokens == kEnd) {
+      invokeRouterGemmBf16Output<__nv_bfloat16, kEnd, kNumExperts, kHiddenDim>(
+          output, input, weights, stream);
+    } else {
+      throw std::invalid_argument("Invalid num_tokens, only supports 1 to 16");
+    }
+  }
+};
+
+void dsv3_router_gemm(at::Tensor& output,       // [num_tokens, num_experts]
+                      const at::Tensor& mat_a,  // [num_tokens, hidden_dim]
+                      const at::Tensor& mat_b   // [num_experts, hidden_dim]
+) {
+  TORCH_CHECK(output.dim() == 2 && mat_a.dim() == 2 && mat_b.dim() == 2);
+
+  const int num_tokens = mat_a.size(0);
+  const int num_experts = mat_b.size(0);
+  const int hidden_dim = mat_a.size(1);
+
+  TORCH_CHECK(mat_a.size(1) == mat_b.size(1),
+              "mat_a and mat_b must have the same hidden_dim");
+  TORCH_CHECK(hidden_dim == DEFAULT_HIDDEN_DIM,
+              "Expected hidden_dim=", DEFAULT_HIDDEN_DIM,
+              ", but got hidden_dim=", hidden_dim);
+  TORCH_CHECK(
+      num_experts == DEFAULT_NUM_EXPERTS || num_experts == KIMI_K2_NUM_EXPERTS,
+      "Expected num_experts=", DEFAULT_NUM_EXPERTS,
+      " or num_experts=", KIMI_K2_NUM_EXPERTS,
+      ", but got num_experts=", num_experts);
+  TORCH_CHECK(num_tokens >= 1 && num_tokens <= 16,
+              "currently num_tokens must be less than or equal to 16 for "
+              "router_gemm");
+  TORCH_CHECK(mat_a.dtype() == at::kBFloat16, "mat_a must be bf16");
+  TORCH_CHECK(mat_b.dtype() == at::kBFloat16, "mat_b must be bf16");
+  TORCH_CHECK(output.dtype() == at::kFloat || output.dtype() == at::kBFloat16,
+              "output must be float32 or bf16");
+
+  auto const sm = getSMVersion();
+  TORCH_CHECK(sm >= 90 && sm <= 103, "required SM_103 >= CUDA ARCH >= SM_90");
+
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+  if (output.dtype() == at::kFloat) {
+    if (num_experts == DEFAULT_NUM_EXPERTS) {
+      LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS, DEFAULT_HIDDEN_DIM>::
+          unroll_float_output(
+              num_tokens, reinterpret_cast<float*>(output.mutable_data_ptr()),
+              reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
+              reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), stream);
+    } else if (num_experts == KIMI_K2_NUM_EXPERTS) {
+      LoopUnroller<1, 16, KIMI_K2_NUM_EXPERTS, DEFAULT_HIDDEN_DIM>::
+          unroll_float_output(
+              num_tokens, reinterpret_cast<float*>(output.mutable_data_ptr()),
+              reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
+              reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), stream);
+    }
+  } else if (output.dtype() == at::kBFloat16) {
+    if (num_experts == DEFAULT_NUM_EXPERTS) {
+      LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS, DEFAULT_HIDDEN_DIM>::
+          unroll_bf16_output(
+              num_tokens,
+              reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()),
+              reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
+              reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), stream);
+    } else if (num_experts == KIMI_K2_NUM_EXPERTS) {
+      LoopUnroller<1, 16, KIMI_K2_NUM_EXPERTS, DEFAULT_HIDDEN_DIM>::
+          unroll_bf16_output(
+              num_tokens,
+              reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()),
+              reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
+              reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), stream);
+    }
+  }
+}
@@ -0,0 +1,291 @@
+/*
+ * Adapted from SGLang's sgl-kernel implementation, which was adapted from
+ * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu
+ * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp
+ *
+ * Copyright (c) 2019-2023, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+
+#include "dsv3_router_gemm_utils.h"
+
+// Custom FMA implementation using PTX assembly instructions
+__device__ __forceinline__ void fma(float2& d, float2 const& a, float2 const& b,
+                                    float2 const& c) {
+  asm volatile("fma.rn.f32x2 %0, %1, %2, %3;\n"
+               : "=l"(reinterpret_cast<uint64_t&>(d))
+               : "l"(reinterpret_cast<uint64_t const&>(a)),
+                 "l"(reinterpret_cast<uint64_t const&>(b)),
+                 "l"(reinterpret_cast<uint64_t const&>(c)));
+}
+
+// Convert 8 bfloat16 values from a uint4 to float array - optimized conversion
+template <int VPT>
+__device__ __forceinline__ void bf16_uint4_to_float8(uint4 const& vec,
+                                                     float* dst) {
+  __nv_bfloat16* bf16_ptr =
+      reinterpret_cast<__nv_bfloat16*>(const_cast<uint4*>(&vec));
+
+#pragma unroll
+  for (int i = 0; i < VPT; i++) {
+    dst[i] = __bfloat162float(bf16_ptr[i]);
+  }
+}
+
+template <typename T, int kBlockSize, int VPT, int kNumTokens, int kNumExperts,
+          int kHiddenDim>
+__global__ __launch_bounds__(128, 1) void router_gemm_kernel_float_output(
+    float* out, T const* mat_a, T const* mat_b) {
+  // Each block handles one expert column
+  int const n_idx = blockIdx.x;
+  int const tid = threadIdx.x;
+  constexpr int kWarpSize = 32;
+  constexpr int kNumWarps = kBlockSize / kWarpSize;
+  // Constants for this kernel
+  constexpr int k_elems_per_k_iteration = VPT * kBlockSize;
+  constexpr int k_iterations =
+      kHiddenDim / k_elems_per_k_iteration;  // Total K iterations
+
+  // Initialize accumulators for all M rows
+  float acc[kNumTokens] = {};
+
+  // Shared memory for warp-level reduction
+  __shared__ float sm_reduction[kNumTokens][kNumWarps];  // kNumWarps
+
+  // B matrix is in column-major order, so we can directly load a column for the
+  // n_idx expert
+  T const* b_col = mat_b + n_idx * kHiddenDim;
+
+  // Pre-compute k_base values for each iteration to help compiler optimize
+  int k_bases[k_iterations];
+#pragma unroll
+  for (int ki = 0; ki < k_iterations; ki++) {
+    k_bases[ki] = ki * k_elems_per_k_iteration + tid * VPT;
+  }
+
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
+  asm volatile("griddepcontrol.wait;");
+#endif
+
+  // Process the GEMM in chunks
+  for (int ki = 0; ki < k_iterations; ki++) {
+    int const k_base = k_bases[ki];
+
+    // Load B matrix values using vector load (8 bf16 values)
+    uint4 b_vec = *reinterpret_cast<uint4 const*>(b_col + k_base);
+
+    // Convert B values to float
+    float b_float[VPT];
+    bf16_uint4_to_float8<VPT>(b_vec, b_float);
+
+// Process each token
+#pragma unroll
+    for (int m_idx = 0; m_idx < kNumTokens; m_idx++) {
+      // Load both rows of A matrix using vector loads
+      uint4 a_vec = *reinterpret_cast<uint4 const*>(
+          mat_a + (m_idx * kHiddenDim) + k_base);
+
+      // Convert A values to float
+      float a_float[VPT];
+      bf16_uint4_to_float8<VPT>(a_vec, a_float);
+
+// Process elements in this chunk
+#pragma unroll
+      for (int k = 0; k < VPT; k++) {
+        float a = a_float[k];
+        float b = b_float[k];
+        acc[m_idx] += a * b;
+      }
+    }
+  }
+
+  // Perform warp-level reduction
+  int const warpSize = 32;
+  int const warpId = tid / warpSize;
+  int const laneId = tid % warpSize;
+
+  // Register for warp-level reduction results
+  float warp_result[kNumTokens];
+
+#pragma unroll
+  for (int m_idx = 0; m_idx < kNumTokens; m_idx++) {
+    warp_result[m_idx] = acc[m_idx];
+  }
+
+// Perform warp-level reduction using optimized butterfly pattern
+#pragma unroll
+  for (int m = 0; m < kNumTokens; m++) {
+    float sum = warp_result[m];
+
+    // Butterfly reduction pattern
+    sum += __shfl_xor_sync(0xffffffff, sum, 16);
+    sum += __shfl_xor_sync(0xffffffff, sum, 8);
+    sum += __shfl_xor_sync(0xffffffff, sum, 4);
+    sum += __shfl_xor_sync(0xffffffff, sum, 2);
+    sum += __shfl_xor_sync(0xffffffff, sum, 1);
+
+    // Only the first thread in each warp stores to shared memory
+    if (laneId == 0) {
+      sm_reduction[m][warpId] = sum;
+    }
+  }
+
+  __syncthreads();
+
+  // Final reduction across warps (only first thread)
+  if (tid == 0) {
+#pragma unroll
+    for (int m = 0; m < kNumTokens; m++) {
+      float final_sum = 0.0f;
+
+// Sum across the kNumWarps
+#pragma unroll
+      for (int w = 0; w < kNumWarps; w++) {
+        final_sum += sm_reduction[m][w];
+      }
+
+      // Write final result
+      out[m * kNumExperts + n_idx] = final_sum;
+    }
+  }
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
+  asm volatile("griddepcontrol.launch_dependents;");
+#endif
+}
+
+template <typename T, int kNumTokens, int kNumExperts, int kHiddenDim>
+void invokeRouterGemmFloatOutput(float* output, T const* mat_a, T const* mat_b,
+                                 cudaStream_t stream) {
+  constexpr int VPT = 16 / sizeof(T);
+  constexpr int kBlockSize = 128;
+  cudaLaunchConfig_t config;
+  config.gridDim = kNumExperts;
+  config.blockDim = kBlockSize;
+  config.dynamicSmemBytes = 0;
+  config.stream = stream;
+  cudaLaunchAttribute attrs[1];
+  attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+  attrs[0].val.programmaticStreamSerializationAllowed = getEnvEnablePDL();
+  config.numAttrs = 1;
+  config.attrs = attrs;
+  cudaLaunchKernelEx(
+      &config,
+      router_gemm_kernel_float_output<T, kBlockSize, VPT, kNumTokens,
+                                      kNumExperts, kHiddenDim>,
+      output, mat_a, mat_b);
+}
+
+// Template instantiations for DEFAULT_NUM_EXPERTS experts
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 1, 256, 7168>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 2, 256, 7168>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 3, 256, 7168>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 4, 256, 7168>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 5, 256, 7168>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 6, 256, 7168>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 7, 256, 7168>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 8, 256, 7168>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 9, 256, 7168>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 10, 256, 7168>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 11, 256, 7168>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 12, 256, 7168>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 13, 256, 7168>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 14, 256, 7168>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 15, 256, 7168>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 16, 256, 7168>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+// Template instantiations for KIMI_K2_NUM_EXPERTS experts
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 1, 384, 7168>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 2, 384, 7168>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 3, 384, 7168>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 4, 384, 7168>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 5, 384, 7168>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 6, 384, 7168>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 7, 384, 7168>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 8, 384, 7168>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 9, 384, 7168>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 10, 384, 7168>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 11, 384, 7168>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 12, 384, 7168>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 13, 384, 7168>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 14, 384, 7168>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 15, 384, 7168>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 16, 384, 7168>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
@@ -0,0 +1,43 @@
+/*
+ * Adapted from SGLang's sgl-kernel implementation, which was adapted from
+ * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu
+ * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp
+ *
+ * Copyright (c) 2019-2023, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <ATen/cuda/CUDAContext.h>
+
+#include <cstdlib>
+#include <mutex>
+
+inline int getSMVersion() {
+  auto* props = at::cuda::getCurrentDeviceProperties();
+  return props->major * 10 + props->minor;
+}
+
+inline bool getEnvEnablePDL() {
+  static std::once_flag flag;
+  static bool enablePDL = false;
+  std::call_once(flag, [&]() {
+    if (getSMVersion() >= 90) {
+      const char* env = std::getenv("TRTLLM_ENABLE_PDL");
+      enablePDL = env && env[0] == '1' && env[1] == '\0';
+    }
+  });
+  return enablePDL;
+}
@@ -1,6 +1,6 @@
 /*
 * Adapted from
- * https://github.com/NVIDIA/TensorRT-LLM/blob/v0.21.0/cpp/tensorrt_llm/kernels/noAuxTcKernels.cu
+ * https://github.com/NVIDIA/TensorRT-LLM/blob/v1.3.0rc2/cpp/tensorrt_llm/kernels/noAuxTcKernels.cu
 * Copyright (c) 2025, The vLLM team.
 * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION &
 * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0
@@ -17,8 +17,10 @@
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
+#include "moeTopKFuncs.cuh"
 #include <c10/cuda/CUDAStream.h>
 #include <torch/all.h>
+#include <cmath>
 #include <cuda_fp16.h>
 #include <cuda_bf16.h>
 #include <cuda/std/limits>
@@ -30,7 +32,17 @@ namespace vllm {
 namespace moe {

 constexpr unsigned FULL_WARP_MASK = 0xffffffff;
-constexpr int32_t WARP_SIZE = 32;
+static constexpr int WARP_SIZE = 32;
+static constexpr int NumNemotronExperts = 512;
+static constexpr int NumKimiK2Experts = 384;
+static constexpr int NumDeepseekExperts = 256;
+static constexpr int MaxSupportedExpertCount =
+    std::max({NumNemotronExperts, NumKimiK2Experts, NumDeepseekExperts});
+static constexpr int MaxNumExpertsUnit = 128;
+static constexpr int NumTopGroupScores = 2;
+static constexpr int DefaultMaxNumTopExperts = 8;
+static constexpr int MaxSupportedTopExperts = 22;
+static constexpr int MaxNumTopGroups = 4;

 namespace warp_topk {

@@ -657,76 +669,335 @@ __global__ void grouped_topk_fused_kernel(
 #endif
 }

-template <typename T, typename BiasT, typename IdxT>
+template <typename T, typename BiasT, typename IdxT, ScoringFunc SF,
+          int MaxNumExperts, bool UseGroups,
+          int MaxNumTopExperts = DefaultMaxNumTopExperts>
+__global__ void grouped_topk_fused_small_expert_count_kernel(
+    T* scores, float* topkValues, IdxT* topkIndices, BiasT const* routingBias,
+    int64_t const numTokens, int64_t const numGroup, int64_t const topkGroup,
+    int64_t const topk, int64_t const numExperts,
+    int64_t const numExpertsPerGroup, bool const renormalize,
+    double const routedScalingFactor) {
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
+  cudaGridDependencySynchronize();
+#endif
+  // declare shared memory structure
+  // number of experts is bounded by number of threads
+  __shared__ float __attribute((aligned(128))) smemScoreSigmoid[MaxNumExperts];
+  __shared__ float __attribute((aligned(128))) smemScoreBias[MaxNumExperts];
+  // number of expert groups is bounded by number of warps
+  int constexpr NumWarps = MaxNumExperts / WARP_SIZE;
+  __shared__ float __attribute((aligned(128))) smemGroupScores[NumWarps];
+
+  // needed for warp reduce
+  auto block = cg::this_thread_block();
+  auto warp = cg::tiled_partition<WARP_SIZE>(block);
+
+  // for the final reduction of weight norm, only some lanes need to participate
+  int32_t laneIdx = threadIdx.x % WARP_SIZE;
+  int32_t warpIdx = __shfl_sync(0xffffffff, threadIdx.x / WARP_SIZE, 0);
+
+  if constexpr (UseGroups) {
+    if (warpIdx >= numGroup) {
+      return;
+    }
+  }
+  // note that for invalid scores, we simply use a negative value:
+  // they work well even with the compacted format used in topK, and
+  // sigmoid / bias activated scores cannot be negative
+  const float invalidScoreFloat = float{-INFINITY};
+
+  // load bias already; each warp represents one expert group
+  auto threadExpert = threadIdx.x;
+  bool expertSelected = threadExpert < numExperts;
+  if constexpr (UseGroups) {
+    threadExpert = warpIdx * numExpertsPerGroup + laneIdx;
+    expertSelected = laneIdx < numExpertsPerGroup;
+  }
+
+  auto scoreIdx = int64_t{blockIdx.x} * int64_t{numExperts} + threadExpert;
+  auto biasVal = expertSelected ? static_cast<float>(routingBias[threadExpert])
+                                : invalidScoreFloat;
+  topkValues += blockIdx.x * topk;
+  topkIndices += blockIdx.x * topk;
+
+  // get our assigned thread score; each warp represents one expert group
+  float score =
+      expertSelected ? static_cast<float>(scores[scoreIdx]) : invalidScoreFloat;
+  auto scoreSigmoid = apply_scoring<SF>(score);
+  // write the sigmoid score to shared for later use
+  if (expertSelected) {
+    smemScoreSigmoid[threadExpert] = scoreSigmoid;
+  }
+
+  // get the score with bias
+  // note that with invalid values, because sigmoid is < 1 and bias is -1,
+  // we must get a negative value, which is smaller than any valid value
+  auto scoreBias = float{scoreSigmoid + float{biasVal}};
+
+  if (expertSelected) {
+    smemScoreBias[threadExpert] = scoreBias;
+  }
+
+  // registers for top group score reduction
+  float topExpGroupScores[NumTopGroupScores];
+  [[maybe_unused]] int32_t topExpGroupIdx[NumTopGroupScores];
+  float topGroups[MaxNumTopGroups];  // bound of numGroup
+  int32_t topGroupIdx[MaxNumTopGroups];
+  float expertScoreGroup[MaxNumTopGroups];
+  int32_t expertIdxGroup[MaxNumTopGroups];
+  float topScores[MaxNumTopExperts];  // bound of topk
+  int32_t topExperts[MaxNumTopExperts];
+
+  if constexpr (UseGroups) {
+    reduce_topk::reduceTopK(warp, topExpGroupScores, topExpGroupIdx, scoreBias,
+                            threadExpert,
+                            /* minValue */ invalidScoreFloat);
+
+    // get the final group score and write it to shared
+    if (warp.thread_rank() == 0) {
+      auto groupScore = topExpGroupScores[0] + topExpGroupScores[1];
+      smemGroupScores[warpIdx] = groupScore;
+    }
+  }
+
+  // make group scores available to all warps
+  __syncthreads();
+
+  if constexpr (UseGroups) {
+    if (warpIdx == 0) {
+      // a single warp performs the selection of top groups, and goes on to
+      // select the final experts
+      float groupScore =
+          laneIdx < numGroup ? smemGroupScores[laneIdx] : invalidScoreFloat;
+
+      reduce_topk::reduceTopK(warp, topGroups, topGroupIdx, groupScore, laneIdx,
+                              /* minValue */ invalidScoreFloat);
+      // final expert selection: get relevant indexes and scores from shared
+#pragma unroll
+      for (int ii = 0; ii < MaxNumTopGroups; ++ii) {  // bound of numGroup
+        auto groupIdx = topGroupIdx[ii];
+        expertIdxGroup[ii] = groupIdx * numExpertsPerGroup + laneIdx;
+
+        expertScoreGroup[ii] = (ii < topkGroup) && expertSelected
+                                   ? smemScoreBias[expertIdxGroup[ii]]
+                                   : invalidScoreFloat;
+      }
+
+      reduce_topk::reduceTopK(warp, topScores, topExperts, expertScoreGroup,
+                              expertIdxGroup, /* minValue */ invalidScoreFloat,
+                              topk);
+    }
+  } else if constexpr (MaxNumExperts > MaxNumExpertsUnit) {
+    // without groups, and the expert number is larger than MaxNumExpertsUnit,
+    // we need to use multiple warps to calculate the intermediate topk results
+
+    int constexpr NumExpertWarps = (MaxNumExperts - 1) / MaxNumExpertsUnit + 1;
+    int constexpr NumInterTopK = NumExpertWarps * MaxNumTopExperts;
+    __shared__ float
+        __attribute((aligned(128))) smemInterTopScores[NumInterTopK];
+    __shared__ int32_t
+        __attribute((aligned(128))) smemInterTopExperts[NumInterTopK];
+    if (warpIdx < NumExpertWarps) {
+      int offset = warpIdx * WARP_SIZE * MaxNumTopGroups;
+#pragma unroll
+      for (int ii = 0; ii < MaxNumTopGroups; ++ii) {
+        auto expertIdx = ii * WARP_SIZE + laneIdx;
+        expertIdxGroup[ii] = offset + expertIdx;
+        expertScoreGroup[ii] = offset + expertIdx < numExperts
+                                   ? smemScoreBias[offset + expertIdx]
+                                   : invalidScoreFloat;
+      }
+      reduce_topk::reduceTopK(warp, topScores, topExperts, expertScoreGroup,
+                              expertIdxGroup,
+                              /* minValue */ invalidScoreFloat, topk);
+
+      if (laneIdx < topk) {
+        smemInterTopScores[warpIdx * MaxNumTopExperts + laneIdx] =
+            topScores[laneIdx];
+        smemInterTopExperts[warpIdx * MaxNumTopExperts + laneIdx] =
+            topExperts[laneIdx];
+      } else if (laneIdx >= topk && laneIdx < MaxNumTopExperts) {
+        smemInterTopScores[warpIdx * MaxNumTopExperts + laneIdx] =
+            invalidScoreFloat;
+        smemInterTopExperts[warpIdx * MaxNumTopExperts + laneIdx] =
+            MaxNumExperts - 1;
+      }
+    }
+    __syncthreads();
+    if (warpIdx == 0) {
+      int constexpr NumInterTopKPerThread = (NumInterTopK - 1) / WARP_SIZE + 1;
+      float intermediateScore[NumInterTopKPerThread];
+      int32_t intermediateExpert[NumInterTopKPerThread];
+      for (int i = laneIdx; i < NumInterTopKPerThread * WARP_SIZE;
+           i += WARP_SIZE) {
+        int ii = i / WARP_SIZE;
+        if (i < NumInterTopK) {
+          intermediateScore[ii] = smemInterTopScores[i];
+          intermediateExpert[ii] = smemInterTopExperts[i];
+        } else {
+          intermediateScore[ii] = invalidScoreFloat;
+          intermediateExpert[ii] = MaxNumExperts - 1;
+        }
+      }
+      reduce_topk::reduceTopK(warp, topScores, topExperts, intermediateScore,
+                              intermediateExpert,
+                              /* minValue */ invalidScoreFloat, topk);
+    }
+  } else {
+    // without groups, and the expert number is smaller than MaxNumExpertsUnit
+    // each thread just takes `MaxNumTopGroups` experts
+    if (warpIdx == 0) {
+#pragma unroll
+      for (int ii = 0; ii < MaxNumTopGroups; ++ii) {
+        auto expertIdx = ii * WARP_SIZE + laneIdx;
+        expertIdxGroup[ii] = expertIdx;
+        expertScoreGroup[ii] = expertIdx < numExperts ? smemScoreBias[expertIdx]
+                                                      : invalidScoreFloat;
+      }
+      reduce_topk::reduceTopK(warp, topScores, topExperts, expertScoreGroup,
+                              expertIdxGroup,
+                              /* minValue */ invalidScoreFloat, topk);
+    }
+  }
+
+  if (warpIdx == 0) {
+    // determine our lane's expert index and write to output
+    int32_t expertIdx =
+        laneIdx < topk ? topExperts[laneIdx] : MaxNumExperts - 1;
+    float scoreNorm = laneIdx < topk ? smemScoreSigmoid[expertIdx] : 0.F;
+    float finalScore = static_cast<float>(scoreNorm * routedScalingFactor);
+    // norm the value
+    if (renormalize) {
+      auto redNorm = cg::reduce(warp, scoreNorm, cg::plus<float>{});
+      finalScore /= (redNorm + 1e-20);
+    }
+    // store the topk scores and experts to output
+    if (laneIdx < topk) {
+      topkValues[laneIdx] = finalScore;
+      topkIndices[laneIdx] = expertIdx;
+    }
+  }
+
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
+  cudaTriggerProgrammaticLaunchCompletion();
+#endif
+}
+
+template <typename T, typename BiasT, typename IdxT, ScoringFunc SF>
 void invokeNoAuxTc(T* scores, float* topk_values, IdxT* topk_indices,
                   BiasT const* bias, int64_t const num_tokens,
                   int64_t const num_experts, int64_t const n_group,
                   int64_t const topk_group, int64_t const topk,
                   bool const renormalize, double const routed_scaling_factor,
-                   int const scoring_func, bool enable_pdl = false,
-                   cudaStream_t const stream = 0) {
+                   bool enable_pdl = false, cudaStream_t const stream = 0) {
  cudaLaunchConfig_t config;
-  // One block per token; one warp per group.
-  config.gridDim = static_cast<uint32_t>(num_tokens);
-  config.blockDim = static_cast<uint32_t>(n_group) * WARP_SIZE;
-  // Dynamic shared memory: WarpSelect staging + per-group topk buffers.
-  int32_t const num_warps = static_cast<int32_t>(n_group);
-  size_t const val_bytes =
-      static_cast<size_t>(num_warps) * WARP_SIZE * sizeof(T);
-  size_t const val_bytes_aligned =
-      warp_topk::round_up_to_multiple_of<256>(val_bytes);
-  size_t const idx_bytes =
-      static_cast<size_t>(num_warps) * WARP_SIZE * sizeof(int32_t);
-  size_t const internal_bytes = val_bytes_aligned + idx_bytes;
-  size_t const extra_bytes = 16 + static_cast<size_t>(n_group) * sizeof(T);
-  config.dynamicSmemBytes = internal_bytes + extra_bytes;
  config.stream = stream;
  cudaLaunchAttribute attrs[1];
  attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
  attrs[0].val.programmaticStreamSerializationAllowed = enable_pdl;
  config.numAttrs = 1;
  config.attrs = attrs;
-  auto const sf = static_cast<ScoringFunc>(scoring_func);
-  switch (sf) {
-    case SCORING_NONE: {
-      auto* kernel_instance =
-          &grouped_topk_fused_kernel<T, BiasT, IdxT, SCORING_NONE>;
-      cudaLaunchKernelEx(&config, kernel_instance, scores, topk_values,
-                         topk_indices, bias, num_tokens, num_experts, n_group,
-                         topk_group, topk, renormalize, routed_scaling_factor);
-      return;
+
+  // Check if we can use the optimized
+  // grouped_topk_fused_small_expert_count_kernel
+  bool const is_single_group =
+      (n_group == 1) && (topk_group == 1) &&
+      (num_experts <= MaxSupportedExpertCount) &&
+      (topk <= DefaultMaxNumTopExperts || topk == MaxSupportedTopExperts);
+
+  int64_t const experts_per_group = num_experts / n_group;
+  bool const is_multi_group =
+      (n_group > 1) && (num_experts <= NumDeepseekExperts) &&
+      (experts_per_group <= WARP_SIZE) &&
+      (experts_per_group * topk_group <= MaxNumExpertsUnit) &&
+      (topk <= DefaultMaxNumTopExperts) && (topk_group <= MaxNumTopGroups);
+
+  if (is_single_group || is_multi_group) {
+    auto* kernel_instance =
+        &grouped_topk_fused_small_expert_count_kernel<T, BiasT, IdxT, SF,
+                                                      NumDeepseekExperts, true>;
+    int num_threads = NumDeepseekExperts;
+    if (is_single_group) {
+      // Special case for Nemotron, which selects top 22 from 512 experts, and 1
+      // group only.
+      if (num_experts == NumNemotronExperts && n_group == 1 &&
+          topk == MaxSupportedTopExperts) {
+        kernel_instance = &grouped_topk_fused_small_expert_count_kernel<
+            T, BiasT, IdxT, SF, NumNemotronExperts, false,
+            MaxSupportedTopExperts>;
+        num_threads = NumNemotronExperts;
+      } else if (num_experts > NumKimiK2Experts &&
+                 num_experts <= MaxSupportedExpertCount) {
+        kernel_instance = &grouped_topk_fused_small_expert_count_kernel<
+            T, BiasT, IdxT, SF, MaxSupportedExpertCount, false>;
+        num_threads = MaxSupportedExpertCount;
+      } else if (num_experts > MaxNumExpertsUnit &&
+                 num_experts <= NumKimiK2Experts) {
+        kernel_instance = &grouped_topk_fused_small_expert_count_kernel<
+            T, BiasT, IdxT, SF, NumKimiK2Experts, false>;
+        num_threads = NumKimiK2Experts;
+      } else {
+        kernel_instance = &grouped_topk_fused_small_expert_count_kernel<
+            T, BiasT, IdxT, SF, MaxNumExpertsUnit, false>;
+        num_threads = MaxNumExpertsUnit;
+      }
    }
-    case SCORING_SIGMOID: {
-      auto* kernel_instance =
-          &grouped_topk_fused_kernel<T, BiasT, IdxT, SCORING_SIGMOID>;
-      cudaLaunchKernelEx(&config, kernel_instance, scores, topk_values,
-                         topk_indices, bias, num_tokens, num_experts, n_group,
-                         topk_group, topk, renormalize, routed_scaling_factor);
-      return;
-    }
-    default:
-      // should be guarded by higher level checks.
-      TORCH_CHECK(false, "Unsupported scoring_func in invokeNoAuxTc");
+    config.gridDim = num_tokens;
+    config.blockDim = num_threads;
+    config.dynamicSmemBytes = 0;
+    cudaLaunchKernelEx(&config, kernel_instance, scores, topk_values,
+                       topk_indices, bias, num_tokens, n_group, topk_group,
+                       topk, num_experts, num_experts / n_group, renormalize,
+                       routed_scaling_factor);
+  } else {
+    auto* kernel_instance = &grouped_topk_fused_kernel<T, BiasT, IdxT, SF>;
+    // One block per token; one warp per group.
+    config.gridDim = static_cast<uint32_t>(num_tokens);
+    config.blockDim = static_cast<uint32_t>(n_group) * WARP_SIZE;
+    // Dynamic shared memory: WarpSelect staging + per-group topk buffers.
+    int32_t const num_warps = static_cast<int32_t>(n_group);
+    size_t const val_bytes =
+        static_cast<size_t>(num_warps) * WARP_SIZE * sizeof(T);
+    size_t const val_bytes_aligned =
+        warp_topk::round_up_to_multiple_of<256>(val_bytes);
+    size_t const idx_bytes =
+        static_cast<size_t>(num_warps) * WARP_SIZE * sizeof(int32_t);
+    size_t const internal_bytes = val_bytes_aligned + idx_bytes;
+    size_t const extra_bytes = 16 + static_cast<size_t>(n_group) * sizeof(T);
+    config.dynamicSmemBytes = internal_bytes + extra_bytes;
+    cudaLaunchKernelEx(&config, kernel_instance, scores, topk_values,
+                       topk_indices, bias, num_tokens, num_experts, n_group,
+                       topk_group, topk, renormalize, routed_scaling_factor);
  }
 }

-#define INSTANTIATE_NOAUX_TC(T, BiasT, IdxT)                                 \
-  template void invokeNoAuxTc<T, BiasT, IdxT>(                               \
+#define INSTANTIATE_NOAUX_TC(T, BiasT, IdxT, SF)                             \
+  template void invokeNoAuxTc<T, BiasT, IdxT, SF>(                           \
      T * scores, float* topk_values, IdxT* topk_indices, BiasT const* bias, \
      int64_t const num_tokens, int64_t const num_experts,                   \
      int64_t const n_group, int64_t const topk_group, int64_t const topk,   \
      bool const renormalize, double const routed_scaling_factor,            \
-      int const scoring_func, bool enable_pdl, cudaStream_t const stream);
+      bool enable_pdl, cudaStream_t const stream);

-INSTANTIATE_NOAUX_TC(float, float, int32_t);
-INSTANTIATE_NOAUX_TC(float, half, int32_t);
-INSTANTIATE_NOAUX_TC(float, __nv_bfloat16, int32_t);
-INSTANTIATE_NOAUX_TC(half, float, int32_t);
-INSTANTIATE_NOAUX_TC(half, half, int32_t);
-INSTANTIATE_NOAUX_TC(half, __nv_bfloat16, int32_t);
-INSTANTIATE_NOAUX_TC(__nv_bfloat16, float, int32_t);
-INSTANTIATE_NOAUX_TC(__nv_bfloat16, half, int32_t);
-INSTANTIATE_NOAUX_TC(__nv_bfloat16, __nv_bfloat16, int32_t);
+INSTANTIATE_NOAUX_TC(float, float, int32_t, SCORING_SIGMOID);
+INSTANTIATE_NOAUX_TC(float, half, int32_t, SCORING_SIGMOID);
+INSTANTIATE_NOAUX_TC(float, __nv_bfloat16, int32_t, SCORING_SIGMOID);
+INSTANTIATE_NOAUX_TC(half, float, int32_t, SCORING_SIGMOID);
+INSTANTIATE_NOAUX_TC(half, half, int32_t, SCORING_SIGMOID);
+INSTANTIATE_NOAUX_TC(half, __nv_bfloat16, int32_t, SCORING_SIGMOID);
+INSTANTIATE_NOAUX_TC(__nv_bfloat16, float, int32_t, SCORING_SIGMOID);
+INSTANTIATE_NOAUX_TC(__nv_bfloat16, half, int32_t, SCORING_SIGMOID);
+INSTANTIATE_NOAUX_TC(__nv_bfloat16, __nv_bfloat16, int32_t, SCORING_SIGMOID);
+INSTANTIATE_NOAUX_TC(float, float, int32_t, SCORING_NONE);
+INSTANTIATE_NOAUX_TC(float, half, int32_t, SCORING_NONE);
+INSTANTIATE_NOAUX_TC(float, __nv_bfloat16, int32_t, SCORING_NONE);
+INSTANTIATE_NOAUX_TC(half, float, int32_t, SCORING_NONE);
+INSTANTIATE_NOAUX_TC(half, half, int32_t, SCORING_NONE);
+INSTANTIATE_NOAUX_TC(half, __nv_bfloat16, int32_t, SCORING_NONE);
+INSTANTIATE_NOAUX_TC(__nv_bfloat16, float, int32_t, SCORING_NONE);
+INSTANTIATE_NOAUX_TC(__nv_bfloat16, half, int32_t, SCORING_NONE);
+INSTANTIATE_NOAUX_TC(__nv_bfloat16, __nv_bfloat16, int32_t, SCORING_NONE);
 }  // end namespace moe
 }  // namespace vllm

@@ -762,46 +1033,53 @@ std::tuple<torch::Tensor, torch::Tensor> grouped_topk(
      {num_tokens, topk}, torch::dtype(torch::kInt32).device(torch::kCUDA));

  auto stream = c10::cuda::getCurrentCUDAStream(scores.get_device());
+  auto const sf = static_cast<vllm::moe::ScoringFunc>(scoring_func);

-#define LAUNCH_KERNEL(T, IdxT)                                               \
-  do {                                                                       \
-    switch (bias_type) {                                                     \
-      case torch::kFloat16:                                                  \
-        vllm::moe::invokeNoAuxTc<T, half, IdxT>(                             \
-            reinterpret_cast<T*>(scores.mutable_data_ptr()),                 \
-            reinterpret_cast<float*>(topk_values.mutable_data_ptr()),        \
-            reinterpret_cast<IdxT*>(topk_indices.mutable_data_ptr()),        \
-            reinterpret_cast<half const*>(bias.data_ptr()), num_tokens,      \
-            num_experts, n_group, topk_group, topk, renormalize,             \
-            routed_scaling_factor, static_cast<int>(scoring_func), false,    \
-            stream);                                                         \
-        break;                                                               \
-      case torch::kFloat32:                                                  \
-        vllm::moe::invokeNoAuxTc<T, float, IdxT>(                            \
-            reinterpret_cast<T*>(scores.mutable_data_ptr()),                 \
-            reinterpret_cast<float*>(topk_values.mutable_data_ptr()),        \
-            reinterpret_cast<IdxT*>(topk_indices.mutable_data_ptr()),        \
-            reinterpret_cast<float const*>(bias.data_ptr()), num_tokens,     \
-            num_experts, n_group, topk_group, topk, renormalize,             \
-            routed_scaling_factor, static_cast<int>(scoring_func), false,    \
-            stream);                                                         \
-        break;                                                               \
-      case torch::kBFloat16:                                                 \
-        vllm::moe::invokeNoAuxTc<T, __nv_bfloat16, IdxT>(                    \
-            reinterpret_cast<T*>(scores.mutable_data_ptr()),                 \
-            reinterpret_cast<float*>(topk_values.mutable_data_ptr()),        \
-            reinterpret_cast<IdxT*>(topk_indices.mutable_data_ptr()),        \
-            reinterpret_cast<__nv_bfloat16 const*>(bias.data_ptr()),         \
-            num_tokens, num_experts, n_group, topk_group, topk, renormalize, \
-            routed_scaling_factor, static_cast<int>(scoring_func), false,    \
-            stream);                                                         \
-        break;                                                               \
-      default:                                                               \
-        throw std::invalid_argument(                                         \
-            "Invalid bias dtype, only supports float16, float32, and "       \
-            "bfloat16");                                                     \
-        break;                                                               \
-    }                                                                        \
+#define LAUNCH_KERNEL_SF(T, BiasT, IdxT)                                      \
+  do {                                                                        \
+    switch (sf) {                                                             \
+      case vllm::moe::SCORING_NONE:                                           \
+        vllm::moe::invokeNoAuxTc<T, BiasT, IdxT, vllm::moe::SCORING_NONE>(    \
+            reinterpret_cast<T*>(scores.mutable_data_ptr()),                  \
+            reinterpret_cast<float*>(topk_values.mutable_data_ptr()),         \
+            reinterpret_cast<IdxT*>(topk_indices.mutable_data_ptr()),         \
+            reinterpret_cast<BiasT const*>(bias.data_ptr()), num_tokens,      \
+            num_experts, n_group, topk_group, topk, renormalize,              \
+            routed_scaling_factor, false, stream);                            \
+        break;                                                                \
+      case vllm::moe::SCORING_SIGMOID:                                        \
+        vllm::moe::invokeNoAuxTc<T, BiasT, IdxT, vllm::moe::SCORING_SIGMOID>( \
+            reinterpret_cast<T*>(scores.mutable_data_ptr()),                  \
+            reinterpret_cast<float*>(topk_values.mutable_data_ptr()),         \
+            reinterpret_cast<IdxT*>(topk_indices.mutable_data_ptr()),         \
+            reinterpret_cast<BiasT const*>(bias.data_ptr()), num_tokens,      \
+            num_experts, n_group, topk_group, topk, renormalize,              \
+            routed_scaling_factor, false, stream);                            \
+        break;                                                                \
+      default:                                                                \
+        throw std::invalid_argument("Unsupported scoring_func");              \
+        break;                                                                \
+    }                                                                         \
+  } while (0)
+
+#define LAUNCH_KERNEL(T, IdxT)                                         \
+  do {                                                                 \
+    switch (bias_type) {                                               \
+      case torch::kFloat16:                                            \
+        LAUNCH_KERNEL_SF(T, half, IdxT);                               \
+        break;                                                         \
+      case torch::kFloat32:                                            \
+        LAUNCH_KERNEL_SF(T, float, IdxT);                              \
+        break;                                                         \
+      case torch::kBFloat16:                                           \
+        LAUNCH_KERNEL_SF(T, __nv_bfloat16, IdxT);                      \
+        break;                                                         \
+      default:                                                         \
+        throw std::invalid_argument(                                   \
+            "Invalid bias dtype, only supports float16, float32, and " \
+            "bfloat16");                                               \
+        break;                                                         \
+    }                                                                  \
  } while (0)

  switch (data_type) {
@@ -824,5 +1102,6 @@ std::tuple<torch::Tensor, torch::Tensor> grouped_topk(
      break;
  }
 #undef LAUNCH_KERNEL
+#undef LAUNCH_KERNEL_SF
  return {topk_values, topk_indices};
 }
@@ -0,0 +1,257 @@
+/*
+ * Adapted from
+ * https://github.com/NVIDIA/TensorRT-LLM/blob/v1.3.0rc2/cpp/tensorrt_llm/kernels/moeTopKFuncs.cuh
+ * Copyright (c) 2026, The vLLM team.
+ * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION. All rights
+ * reserved. SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+#include <cooperative_groups.h>
+#include <cooperative_groups/reduce.h>
+#include <cub/cub.cuh>
+
+namespace vllm {
+namespace moe {
+namespace reduce_topk {
+namespace cg = cooperative_groups;
+static constexpr int kWARP_SIZE = 32;
+
+template <typename T_>
+struct TopKRedType {
+  using T = T_;
+  static_assert(
+      std::is_same_v<T, float> || std::is_same_v<T, half> ||
+          std::is_same_v<T, __nv_bfloat16> || std::is_same_v<T, int>,
+      "Top K reduction only implemented for int, float, float16 and bfloat16");
+
+  using TypeCmp = std::conditional_t<sizeof(T) == 4, uint64_t, uint32_t>;
+  using IdxT = std::conditional_t<sizeof(T) == 4, int32_t, int16_t>;
+
+  static constexpr int kMoveBits = (sizeof(T) == 4) ? 32 : 16;
+  static constexpr int kMaxIdx = 65535;
+  TypeCmp compValIdx;
+
+  static __host__ __device__ inline TypeCmp makeCmpVal(T val, int32_t idx = 0) {
+    auto valueBits = cub::Traits<T>::TwiddleIn(
+        reinterpret_cast<typename cub::Traits<T>::UnsignedBits&>(val));
+    TypeCmp compactTmp = valueBits;
+    compactTmp = (compactTmp << kMoveBits) | (0xFFFF & (kMaxIdx - idx));
+    // Use 65535 minus idx to give higher priority to elements with smaller
+    // indices.
+    return compactTmp;
+  }
+
+  static __host__ __device__ void unpack(T& value, int32_t& index,
+                                         TypeCmp cmp) {
+    // Since “65535-idx” is always smaller than 65536 and positive, we can
+    // directly use it as the lower 16 bits
+    index = kMaxIdx - static_cast<int32_t>((cmp & 0xFFFF));
+
+    auto compactTmp = cmp >> kMoveBits;
+    auto valueBits = cub::Traits<T>::TwiddleOut(
+        reinterpret_cast<typename cub::Traits<T>::UnsignedBits&>(compactTmp));
+    value = reinterpret_cast<T&>(valueBits);
+  }
+
+  __host__ __device__ TopKRedType() = default;
+
+  __host__ __device__ TopKRedType(T val, int32_t idx)
+      : compValIdx(makeCmpVal(val, idx)) {}
+
+  __host__ __device__ operator TypeCmp() const noexcept { return compValIdx; }
+
+  __device__ inline TypeCmp reduce(
+      cg::thread_block_tile<kWARP_SIZE> const& warp) {
+    return cg::reduce(warp, compValIdx, cg::greater<TypeCmp>{});
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <int K_, bool Enable_>
+struct TopKIdx {
+  // by default, empty
+};
+
+template <int K_>
+struct TopKIdx<K_, true> {
+  static constexpr int K = K_;
+  int32_t val[K];
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+#define TOPK_SWAP(I, J)                                         \
+  {                                                             \
+    auto pairMin = min(topK[I].compValIdx, topK[J].compValIdx); \
+    auto pairMax = max(topK[I].compValIdx, topK[J].compValIdx); \
+    topK[I].compValIdx = pairMax;                               \
+    topK[J].compValIdx = pairMin;                               \
+  }
+
+template <int N, typename RedType>
+struct Sort;
+
+template <typename RedType>
+struct Sort<1, RedType> {
+  static __device__ void run(RedType* topK) {}
+};
+
+template <typename RedType>
+struct Sort<2, RedType> {
+  static __device__ void run(RedType* topK) { TOPK_SWAP(0, 1); }
+};
+
+template <typename RedType>
+struct Sort<3, RedType> {
+  static __device__ void run(RedType* topK) {
+    TOPK_SWAP(0, 1);
+    TOPK_SWAP(1, 2);
+    TOPK_SWAP(0, 1);
+  }
+};
+
+template <typename RedType>
+struct Sort<4, RedType> {
+  static __device__ void run(RedType* topK) {
+    TOPK_SWAP(0, 2);
+    TOPK_SWAP(1, 3);
+    TOPK_SWAP(0, 1);
+    TOPK_SWAP(2, 3);
+    TOPK_SWAP(1, 2);
+  }
+};
+
+template <int K, typename Type>
+__forceinline__ __device__ void reduceTopK(
+    cg::thread_block_tile<kWARP_SIZE> const& warp, Type (&out)[K],
+    int32_t (&outIdx)[K], Type value, int32_t idx, Type const minValue,
+    int actualK = K) {
+  static_assert(K > 0, "Top K must have K > 0");
+  static_assert(K < kWARP_SIZE, "Top K must have K < kWARP_SIZE");
+  using RedType = TopKRedType<Type>;
+  RedType topK{value, idx};
+  typename RedType::TypeCmp packedMax{};
+#pragma unroll
+  for (int kk = 0; kk < actualK; ++kk) {
+    topK =
+        kk > 0 && packedMax == topK.compValIdx ? RedType{minValue, idx} : topK;
+    // get the next largest value
+    packedMax = topK.reduce(warp);
+    RedType::unpack(out[kk], outIdx[kk], packedMax);
+  }
+};
+
+template <int K, typename Type, int N, bool IsSorted = false>
+__device__ void reduceTopKFunc(cg::thread_block_tile<kWARP_SIZE> const& warp,
+                               Type (&out)[K], int32_t (&outIdx)[K],
+                               Type (&value)[N], int32_t (&idx)[N],
+                               Type minValue, int actualK = K) {
+  static_assert(K > 0, "Top K must have K > 0");
+  static_assert(K < kWARP_SIZE, "Top K must have K < kWARP_SIZE");
+  static_assert(N > 0, "Top K must have N > 0");
+  static_assert(N < 5,
+                "Only support candidates number less than or equal to 128");
+  using RedType = TopKRedType<Type>;
+  RedType topK[N];
+#pragma unroll
+  for (int nn = 0; nn < N; ++nn) {
+    topK[nn] = RedType{value[nn], idx[nn]};
+  }
+
+  if constexpr (!IsSorted) {
+    Sort<N, RedType>::run(topK);
+  }
+  typename RedType::TypeCmp packedMax{};
+#pragma unroll
+  for (int kk = 0; kk < actualK; ++kk) {
+    bool update = kk > 0 && packedMax == topK[0].compValIdx;
+#pragma unroll
+    for (int nn = 0; nn < N; ++nn) {
+      topK[nn] = update && nn == N - 1 ? RedType{minValue, idx[nn]}
+                 : update              ? topK[nn + 1]
+                                       : topK[nn];
+    }
+    // get the next largest value
+    packedMax = topK[0].reduce(warp);
+    RedType::unpack(out[kk], outIdx[kk], packedMax);
+  }
+};
+
+template <int K, typename Type, int N>
+__forceinline__ __device__ void reduceTopK(
+    cg::thread_block_tile<kWARP_SIZE> const& warp, Type (&out)[K],
+    int32_t (&outIdx)[K], Type (&value)[N], int32_t (&idx)[N],
+    Type const minValue, int actualK = K) {
+  static_assert(K > 0, "Top K must have K > 0");
+  static_assert(K < kWARP_SIZE, "Top K must have K < kWARP_SIZE");
+  static_assert(N > 0, "Top K must have N > 0");
+  static_assert(
+      N <= 16,
+      "Only support candidates number less than or equal to 16*32=512");
+  static_assert(N <= 4 || N % 4 == 0,
+                "Only support candidates number is a multiple of 4*32=128 or "
+                "less than or equal to 4");
+  using RedType = TopKRedType<Type>;
+
+  if constexpr (N <= 4) {
+    reduceTopKFunc<K, Type, N>(warp, out, outIdx, value, idx, minValue,
+                               actualK);
+  } else {
+    constexpr int numLoops = N / 4;
+    constexpr int numResults = (numLoops * K - 1) / kWARP_SIZE + 1;
+
+    Type topKBufferValue[numResults];
+    int32_t topKBufferIdx[numResults];
+    int32_t laneIdx = threadIdx.x % kWARP_SIZE;
+
+    for (int ii = 0; ii < numResults; ++ii) {
+      topKBufferValue[ii] = minValue;
+      topKBufferIdx[ii] = ii * kWARP_SIZE - 1;
+    }
+    for (int loop = 0; loop < numLoops; ++loop) {
+      int start = loop * 4;
+      Type topKValue[K];
+      int32_t topKIdx[K];
+      Type inValue[4];
+      int32_t inIdx[4];
+      for (int i = 0; i < 4; ++i) {
+        inValue[i] = value[start + i];
+        inIdx[i] = idx[start + i];
+      }
+      reduceTopKFunc<K, Type, 4>(warp, topKValue, topKIdx, inValue, inIdx,
+                                 minValue, actualK);
+      int inOffset = laneIdx % K;
+      if (laneIdx >= loop * K && laneIdx < (loop + 1) * K) {
+        topKBufferValue[0] = topKValue[inOffset];
+        topKBufferIdx[0] = topKIdx[inOffset];
+      }
+      if (loop == numLoops - 1 && (laneIdx < (numLoops * K - kWARP_SIZE))) {
+        topKBufferValue[1] = topKValue[inOffset];
+        topKBufferIdx[1] = topKIdx[inOffset];
+      }
+    }
+
+    reduceTopKFunc<K, Type, numResults>(warp, out, outIdx, topKBufferValue,
+                                        topKBufferIdx, minValue, actualK);
+  }
+};
+
+#undef TOPK_SWAP
+
+}  // namespace reduce_topk
+}  // namespace moe
+}  // namespace vllm
@@ -55,4 +55,15 @@ bool moe_permute_unpermute_supported();

 void shuffle_rows(const torch::Tensor& input_tensor,
                  const torch::Tensor& dst2src_map,
-                  torch::Tensor& output_tensor);
+                  torch::Tensor& output_tensor);
+
+#ifndef USE_ROCM
+// DeepSeek V3 optimized router GEMM kernel for SM90+
+// Computes output = mat_a @ mat_b.T where:
+//   mat_a: [num_tokens, hidden_dim] in bf16
+//   mat_b: [num_experts, hidden_dim] in bf16
+//   output: [num_tokens, num_experts] in bf16 or fp32
+// Supports num_tokens in [1, 16], num_experts in {256, 384}, hidden_dim = 7168
+void dsv3_router_gemm(torch::Tensor& output, const torch::Tensor& mat_a,
+                      const torch::Tensor& mat_b);
+#endif
@@ -124,6 +124,10 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
      "routed_scaling_factor, Tensor bias, int scoring_func) -> (Tensor, "
      "Tensor)");
  m.impl("grouped_topk", torch::kCUDA, &grouped_topk);
+
+  // DeepSeek V3 optimized router GEMM for SM90+
+  m.def("dsv3_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
+  m.impl("dsv3_router_gemm", torch::kCUDA, &dsv3_router_gemm);
 #endif
 }

@@ -582,7 +582,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 # This is ~1.1GB and only changes when FlashInfer version bumps
 # https://docs.flashinfer.ai/installation.html
 # From versions.json: .flashinfer.version
-ARG FLASHINFER_VERSION=0.6.3
+ARG FLASHINFER_VERSION=0.6.4
 RUN --mount=type=cache,target=/root/.cache/uv \
    uv pip install --system flashinfer-cubin==${FLASHINFER_VERSION} \
    && uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
@@ -217,13 +217,13 @@ RUN pip install setuptools==75.6.0 packaging==23.2 ninja==1.11.1.3 build==1.2.2.


 # build flashinfer for torch nightly from source around 10 mins
-# release version: v0.6.3
+# release version: v0.6.4
 # todo(elainewy): cache flashinfer build result for faster build
 ENV CCACHE_DIR=/root/.cache/ccache
 RUN --mount=type=cache,target=/root/.cache/ccache \
    --mount=type=cache,target=/root/.cache/uv \
    echo "git clone flashinfer..." \
-    && git clone --depth 1 --branch v0.6.3 --recursive https://github.com/flashinfer-ai/flashinfer.git \
+    && git clone --depth 1 --branch v0.6.4 --recursive https://github.com/flashinfer-ai/flashinfer.git \
    && cd flashinfer \
    && git submodule update --init --recursive \
    && echo "finish git clone flashinfer..." \
@@ -68,7 +68,7 @@
      "default": "true"
    },
    "FLASHINFER_VERSION": {
-      "default": "0.6.3"
+      "default": "0.6.4"
    },
    "GDRCOPY_CUDA_VERSION": {
      "default": "12.8"
@@ -293,21 +293,22 @@ Assuming that the memory usage increases with the number of tokens, the dummy in
            self,
            seq_len: int,
            mm_counts: Mapping[str, int],
-            mm_options: Mapping[str, BaseDummyOptions] | None = None,
+            mm_options: Mapping[str, BaseDummyOptions],
        ) -> MultiModalDataDict:
            num_images = mm_counts.get("image", 0)

            target_width, target_height = \
                self.info.get_image_size_with_most_features()

-            image_overrides = mm_options.get("image") if mm_options else None
+            image_overrides = mm_options.get("image")

            return {
-                "image":
-                self._get_dummy_images(width=target_width,
-                                    height=target_height,
-                                    num_images=num_images,
-                                    overrides=image_overrides)
+                "image": self._get_dummy_images(
+                    width=target_width,
+                    height=target_height,
+                    num_images=num_images,
+                    overrides=image_overrides,
+                )
            }
        ```

@@ -479,17 +480,16 @@ Assuming that the memory usage increases with the number of tokens, the dummy in
            self,
            seq_len: int,
            mm_counts: Mapping[str, int],
-            mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,
+            mm_options: Mapping[str, BaseDummyOptions],
        ) -> MultiModalDataDict:
            target_width, target_height = \
                self.info.get_image_size_with_most_features()
            num_images = mm_counts.get("image", 0)

-            image_overrides = mm_options.get("image") if mm_options else None
+            image_overrides = mm_options.get("image")

            return {
-                "image":
-                self._get_dummy_images(
+                "image": self._get_dummy_images(
                    width=target_width,
                    height=target_height,
                    num_images=num_images,
@@ -36,12 +36,12 @@ th:not(:first-child) {
 }
 </style>

-| Feature | [CP](../configuration/optimization.md#chunked-prefill) | [APC](automatic_prefix_caching.md) | [LoRA](lora.md) | [SD](spec_decode/README.md) | CUDA graph | [pooling](../models/pooling_models.md) | <abbr title="Encoder-Decoder Models">enc-dec</abbr> | <abbr title="Logprobs">logP</abbr> | <abbr title="Prompt Logprobs">prmpt logP</abbr> | <abbr title="Async Output Processing">async output</abbr> | multi-step | <abbr title="Multimodal Inputs">mm</abbr> | best-of | beam-search | [prompt-embeds](prompt_embeds.md) |
+| Feature | [CP](../configuration/optimization.md#chunked-prefill) | [APC](automatic_prefix_caching.md) | [LoRA](lora.md) | [SD](speculative_decoding/README.md) | CUDA graph | [pooling](../models/pooling_models.md) | <abbr title="Encoder-Decoder Models">enc-dec</abbr> | <abbr title="Logprobs">logP</abbr> | <abbr title="Prompt Logprobs">prmpt logP</abbr> | <abbr title="Async Output Processing">async output</abbr> | multi-step | <abbr title="Multimodal Inputs">mm</abbr> | best-of | beam-search | [prompt-embeds](prompt_embeds.md) |
 |---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
 | [CP](../configuration/optimization.md#chunked-prefill) | ✅ | | | | | | | | | | | | | | |
 | [APC](automatic_prefix_caching.md) | ✅ | ✅ | | | | | | | | | | | | | |
 | [LoRA](lora.md) | ✅ | ✅ | ✅ | | | | | | | | | | | | |
-| [SD](spec_decode/README.md) | ✅ | ✅ | ❌ | ✅ | | | | | | | | | | | |
+| [SD](speculative_decoding/README.md) | ✅ | ✅ | ❌ | ✅ | | | | | | | | | | | |
 | CUDA graph | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | | | | | | |
 | [pooling](../models/pooling_models.md) | 🟠\* | 🟠\* | ✅ | ❌ | ✅ | ✅ | | | | | | | | | |
 | <abbr title="Encoder-Decoder Models">enc-dec</abbr> | ❌ | [❌](https://github.com/vllm-project/vllm/issues/7366) | ❌ | [❌](https://github.com/vllm-project/vllm/issues/7366) | ✅ | ✅ | ✅ | | | | | | | | |
@@ -64,7 +64,7 @@ th:not(:first-child) {
 | [CP](../configuration/optimization.md#chunked-prefill)                                     | [❌](https://github.com/vllm-project/vllm/issues/2729) | ✅        | ✅        | ✅     | ✅        | ✅                  | ✅     | ✅        |
 | [APC](automatic_prefix_caching.md)                        | [❌](https://github.com/vllm-project/vllm/issues/3687) | ✅        | ✅        | ✅     | ✅        | ✅                  | ✅     | ✅        |
 | [LoRA](lora.md)                                           | ✅                  | ✅        | ✅        | ✅     | ✅        | ✅                  | ✅     | ✅        |
-| [SD](spec_decode/README.md)                        | ✅                  | ✅        | ✅        | ✅     | ✅        | ❌                  | ✅     | ✅        |
+| [SD](speculative_decoding/README.md)                        | ✅                  | ✅        | ✅        | ✅     | ✅        | ❌                  | ✅     | ✅        |
 | CUDA graph                                                | ✅                  | ✅        | ✅        | ✅     | ✅        | ❌                  | ✅     | [❌](https://github.com/vllm-project/vllm/issues/26970)        |
 | [pooling](../models/pooling_models.md)                    | ✅                  | ✅        | ✅        | ✅     | ✅        | ✅                  | ✅     | ✅        |
 | <abbr title="Encoder-Decoder Models">enc-dec</abbr>       | ✅                  | ✅        | ✅        | ✅     | ✅        | ✅                  | ❌     | ✅        |
@@ -197,8 +197,8 @@ For multi-host DP deployment, only need to provide the host/port of the head ins

 The `kv_load_failure_policy` setting controls how the system handles failures when the decoder instance loads KV cache blocks from the prefiller instance:

- **fail** (recommended): Immediately fail the request with an error when KV load fails. This prevents performance degradation by avoiding recomputation of prefill work on the decode instance.
- **recompute** (default): Recompute failed blocks locally on the decode instance. This may cause performance _jitter_ on decode instances as the scheduled prefill will delay and interfere with other decodes. Furthermore, decode instances are typically configured with low-latency optimizations.
+- **fail** (default): Immediately fail the request with an error when KV load fails. This prevents performance degradation by avoiding recomputation of prefill work on the decode instance.
+- **recompute**: Recompute failed blocks locally on the decode instance. This may cause performance _jitter_ on decode instances as the scheduled prefill will delay and interfere with other decodes. Furthermore, decode instances are typically configured with low-latency optimizations.

 !!! warning
    Using `kv_load_failure_policy="recompute"` can lead to performance degradation in production deployments. When KV loads fail, the decode instance will execute prefill work with decode-optimized configurations, which is inefficient and defeats the purpose of disaggregated prefilling. This also increases tail latency for other ongoing decode requests.
@@ -7,7 +7,7 @@ Compared to other quantization methods, BitsAndBytes eliminates the need for cal
 Below are the steps to utilize BitsAndBytes with vLLM.

 ```bash
-pip install bitsandbytes>=0.46.1
+pip install bitsandbytes>=0.49.2
 ```

 vLLM reads the model's config file and supports both in-flight quantization and pre-quantized checkpoint.
@@ -1,330 +0,0 @@
-# Speculative Decoding
-
-!!! warning
-    Currently, speculative decoding in vLLM is not compatible with pipeline parallelism.
-
-This document shows how to use [Speculative Decoding](https://x.com/karpathy/status/1697318534555336961) with vLLM.
-Speculative decoding is a technique which improves inter-token latency in memory-bound LLM inference.
-
-!!! tip
-    To train your own draft models for speculative decoding, see [Speculators](speculators.md), a library for training draft models that integrates seamlessly with vLLM.
-
-## Speculating with a draft model
-
-The following code configures vLLM in an offline mode to use speculative decoding with a draft model, speculating 5 tokens at a time.
-
-!!! warning
-    In vllm v0.10.0, speculative decoding with a draft model is not supported.
-    If you use the following code, you will get a `NotImplementedError`.
-
-??? code
-
-    ```python
-    from vllm import LLM, SamplingParams
-
-    prompts = [
-        "The future of AI is",
-    ]
-    sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
-
-    llm = LLM(
-        model="facebook/opt-6.7b",
-        tensor_parallel_size=1,
-        speculative_config={
-            "model": "facebook/opt-125m",
-            "num_speculative_tokens": 5,
-        },
-    )
-    outputs = llm.generate(prompts, sampling_params)
-
-    for output in outputs:
-        prompt = output.prompt
-        generated_text = output.outputs[0].text
-        print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
-    ```
-
-To perform the same with an online mode launch the server:
-
-```bash
-vllm serve facebook/opt-6.7b \
-    --host 0.0.0.0 \
-    --port 8000 \
-    --seed 42 \
-    -tp 1 \
-    --gpu_memory_utilization 0.8 \
-    --speculative_config '{"model": "facebook/opt-125m", "num_speculative_tokens": 5}'
-```
-
-!!! warning
-    Note: Please use `--speculative_config` to set all configurations related to speculative decoding. The previous method of specifying the model through `--speculative_model` and adding related parameters (e.g., `--num_speculative_tokens`) separately has been deprecated now.
-
-Then use a client:
-
-??? code
-
-    ```python
-    from openai import OpenAI
-
-    # Modify OpenAI's API key and API base to use vLLM's API server.
-    openai_api_key = "EMPTY"
-    openai_api_base = "http://localhost:8000/v1"
-
-    client = OpenAI(
-        # defaults to os.environ.get("OPENAI_API_KEY")
-        api_key=openai_api_key,
-        base_url=openai_api_base,
-    )
-
-    models = client.models.list()
-    model = models.data[0].id
-
-    # Completion API
-    stream = False
-    completion = client.completions.create(
-        model=model,
-        prompt="The future of AI is",
-        echo=False,
-        n=1,
-        stream=stream,
-    )
-
-    print("Completion results:")
-    if stream:
-        for c in completion:
-            print(c)
-    else:
-        print(completion)
-    ```
-
-## Speculating by matching n-grams in the prompt
-
-The following code configures vLLM to use speculative decoding where proposals are generated by
-matching n-grams in the prompt. For more information read [this thread.](https://x.com/joao_gante/status/1747322413006643259)
-
-??? code
-
-    ```python
-    from vllm import LLM, SamplingParams
-
-    prompts = [
-        "The future of AI is",
-    ]
-    sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
-
-    llm = LLM(
-        model="facebook/opt-6.7b",
-        tensor_parallel_size=1,
-        speculative_config={
-            "method": "ngram",
-            "num_speculative_tokens": 5,
-            "prompt_lookup_max": 4,
-        },
-    )
-    outputs = llm.generate(prompts, sampling_params)
-
-    for output in outputs:
-        prompt = output.prompt
-        generated_text = output.outputs[0].text
-        print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
-    ```
-
-## Speculating using Suffix Decoding
-
-The following code configures vLLM to use speculative decoding where proposals are generated using Suffix Decoding ([technical report](https://arxiv.org/abs/2411.04975)).
-
-Like n-gram, Suffix Decoding can generate draft tokens by pattern-matching using the last `n` generated tokens. Unlike n-gram, Suffix Decoding (1) can pattern-match against both the prompt and previous generations, (2) uses frequency counts to propose the most likely continuations, and (3) speculates an adaptive number of tokens for each request at each iteration to get better acceptance rates.
-
-Suffix Decoding can achieve better performance for tasks with high repetition, such as code-editing, agentic loops (e.g. self-reflection, self-consistency), and RL rollouts.
-
-!!! tip "Install Arctic Inference"
-    Suffix Decoding requires [Arctic Inference](https://github.com/snowflakedb/ArcticInference). You can install it with `pip install arctic-inference`.
-
-!!! tip "Suffix Decoding Speculative Tokens"
-    Suffix Decoding will speculate a dynamic number of tokens for each request at each decoding step, so the `num_speculative_tokens` configuration specifies the *maximum* number of speculative tokens. It is suggested to use a high number such as `16` or `32` (default).
-
-??? code
-
-    ```python
-    from vllm import LLM, SamplingParams
-
-    prompts = [
-        "The future of AI is",
-    ]
-    sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
-
-    llm = LLM(
-        model="facebook/opt-6.7b",
-        tensor_parallel_size=1,
-        speculative_config={
-            "method": "suffix",
-            "num_speculative_tokens": 32,
-        },
-    )
-    outputs = llm.generate(prompts, sampling_params)
-
-    for output in outputs:
-        prompt = output.prompt
-        generated_text = output.outputs[0].text
-        print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
-    ```
-
-## Speculating using MLP speculators
-
-The following code configures vLLM to use speculative decoding where proposals are generated by
-draft models that condition draft predictions on both context vectors and sampled tokens.
-For more information see [this blog](https://pytorch.org/blog/hitchhikers-guide-speculative-decoding/) or
-[this technical report](https://arxiv.org/abs/2404.19124).
-
-??? code
-
-    ```python
-    from vllm import LLM, SamplingParams
-
-    prompts = [
-        "The future of AI is",
-    ]
-    sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
-
-    llm = LLM(
-        model="meta-llama/Meta-Llama-3.1-70B-Instruct",
-        tensor_parallel_size=4,
-        speculative_config={
-            "model": "ibm-ai-platform/llama3-70b-accelerator",
-            "draft_tensor_parallel_size": 1,
-        },
-    )
-    outputs = llm.generate(prompts, sampling_params)
-
-    for output in outputs:
-        prompt = output.prompt
-        generated_text = output.outputs[0].text
-        print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
-    ```
-
-Note that these speculative models currently need to be run without tensor parallelism, although
-it is possible to run the main model using tensor parallelism (see example above). Since the
-speculative models are relatively small, we still see significant speedups. However, this
-limitation will be fixed in a future release.
-
-A variety of speculative models of this type are available on HF hub:
-
- [llama-13b-accelerator](https://huggingface.co/ibm-ai-platform/llama-13b-accelerator)
- [llama3-8b-accelerator](https://huggingface.co/ibm-ai-platform/llama3-8b-accelerator)
- [codellama-34b-accelerator](https://huggingface.co/ibm-ai-platform/codellama-34b-accelerator)
- [llama2-70b-accelerator](https://huggingface.co/ibm-ai-platform/llama2-70b-accelerator)
- [llama3-70b-accelerator](https://huggingface.co/ibm-ai-platform/llama3-70b-accelerator)
- [granite-3b-code-instruct-accelerator](https://huggingface.co/ibm-granite/granite-3b-code-instruct-accelerator)
- [granite-8b-code-instruct-accelerator](https://huggingface.co/ibm-granite/granite-8b-code-instruct-accelerator)
- [granite-7b-instruct-accelerator](https://huggingface.co/ibm-granite/granite-7b-instruct-accelerator)
- [granite-20b-code-instruct-accelerator](https://huggingface.co/ibm-granite/granite-20b-code-instruct-accelerator)
-
-## Speculating using EAGLE based draft models
-
-The following code configures vLLM to use speculative decoding where proposals are generated by
-an [EAGLE (Extrapolation Algorithm for Greater Language-model Efficiency)](https://arxiv.org/pdf/2401.15077) based draft model. A more detailed example for offline mode, including how to extract request level acceptance rate, can be found in [examples/offline_inference/spec_decode.py](../../../examples/offline_inference/spec_decode.py)
-
-??? code
-
-    ```python
-    from vllm import LLM, SamplingParams
-
-    prompts = [
-        "The future of AI is",
-    ]
-    sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
-
-    llm = LLM(
-        model="meta-llama/Meta-Llama-3-8B-Instruct",
-        tensor_parallel_size=4,
-        speculative_config={
-            "model": "yuhuili/EAGLE-LLaMA3-Instruct-8B",
-            "draft_tensor_parallel_size": 1,
-            "num_speculative_tokens": 2,
-            "method": "eagle",
-        },
-    )
-
-    outputs = llm.generate(prompts, sampling_params)
-
-    for output in outputs:
-        prompt = output.prompt
-        generated_text = output.outputs[0].text
-        print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
-
-    ```
-
-A few important things to consider when using the EAGLE based draft models:
-
-1. The EAGLE draft models available in the [HF repository for EAGLE models](https://huggingface.co/yuhuili) should
-   be able to be loaded and used directly by vLLM after <https://github.com/vllm-project/vllm/pull/12304>.
-   If you are using vllm version before <https://github.com/vllm-project/vllm/pull/12304>, please use the
-   [script](https://gist.github.com/abhigoyal1997/1e7a4109ccb7704fbc67f625e86b2d6d) to convert the speculative model,
-   and specify `"model": "path/to/modified/eagle/model"` in `speculative_config`. If weight-loading problems still occur when using the latest version of vLLM, please leave a comment or raise an issue.
-
-2. The EAGLE based draft models need to be run without tensor parallelism
-   (i.e. draft_tensor_parallel_size is set to 1 in `speculative_config`), although
-   it is possible to run the main model using tensor parallelism (see example above).
-
-3. When using EAGLE-based speculators with vLLM, the observed speedup is lower than what is
-   reported in the reference implementation [here](https://github.com/SafeAILab/EAGLE). This issue is under
-   investigation and tracked here: <https://github.com/vllm-project/vllm/issues/9565>.
-
-4. When using EAGLE-3 based draft model, option "method" must be set to "eagle3".
-   That is, to specify `"method": "eagle3"` in `speculative_config`.
-
-A variety of EAGLE draft models are available on the Hugging Face hub:
-
-| Base Model                                                           | EAGLE on Hugging Face                     | # EAGLE Parameters |
-|---------------------------------------------------------------------|-------------------------------------------|--------------------|
-| Vicuna-7B-v1.3                                                       | yuhuili/EAGLE-Vicuna-7B-v1.3             | 0.24B              |
-| Vicuna-13B-v1.3                                                      | yuhuili/EAGLE-Vicuna-13B-v1.3            | 0.37B              |
-| Vicuna-33B-v1.3                                                      | yuhuili/EAGLE-Vicuna-33B-v1.3            | 0.56B              |
-| LLaMA2-Chat 7B                                                       | yuhuili/EAGLE-llama2-chat-7B             | 0.24B              |
-| LLaMA2-Chat 13B                                                      | yuhuili/EAGLE-llama2-chat-13B            | 0.37B              |
-| LLaMA2-Chat 70B                                                      | yuhuili/EAGLE-llama2-chat-70B            | 0.99B              |
-| Mixtral-8x7B-Instruct-v0.1                                           | yuhuili/EAGLE-mixtral-instruct-8x7B      | 0.28B              |
-| LLaMA3-Instruct 8B                                                   | yuhuili/EAGLE-LLaMA3-Instruct-8B         | 0.25B              |
-| LLaMA3-Instruct 70B                                                  | yuhuili/EAGLE-LLaMA3-Instruct-70B        | 0.99B              |
-| Qwen2-7B-Instruct                                                    | yuhuili/EAGLE-Qwen2-7B-Instruct          | 0.26B              |
-| Qwen2-72B-Instruct                                                   | yuhuili/EAGLE-Qwen2-72B-Instruct         | 1.05B              |
-
-## Lossless guarantees of Speculative Decoding
-
-In vLLM, speculative decoding aims to enhance inference efficiency while maintaining accuracy. This section addresses the lossless guarantees of
-speculative decoding, breaking down the guarantees into three key areas:
-
-1. **Theoretical Losslessness**
-   \- Speculative decoding sampling is theoretically lossless up to the precision limits of hardware numerics. Floating-point errors might
-   cause slight variations in output distributions, as discussed
-   in [Accelerating Large Language Model Decoding with Speculative Sampling](https://arxiv.org/pdf/2302.01318)
-
-2. **Algorithmic Losslessness**
-   \- vLLM’s implementation of speculative decoding is algorithmically validated to be lossless. Key validation tests include:
-
-    > - **Rejection Sampler Convergence**: Ensures that samples from vLLM’s rejection sampler align with the target
-    >   distribution. [View Test Code](https://github.com/vllm-project/vllm/blob/47b65a550866c7ffbd076ecb74106714838ce7da/tests/samplers/test_rejection_sampler.py#L252)
-    > - **Greedy Sampling Equality**: Confirms that greedy sampling with speculative decoding matches greedy sampling
-    >   without it. This verifies that vLLM's speculative decoding framework, when integrated with the vLLM forward pass and the vLLM rejection sampler,
-    >   provides a lossless guarantee. Almost all of the tests in [tests/spec_decode/e2e](../../tests/spec_decode/e2e).
-    >   verify this property using [this assertion implementation](https://github.com/vllm-project/vllm/blob/b67ae00cdbbe1a58ffc8ff170f0c8d79044a684a/tests/spec_decode/e2e/conftest.py#L291)
-
-3. **vLLM Logprob Stability**
-   \- vLLM does not currently guarantee stable token log probabilities (logprobs). This can result in different outputs for the
-   same request across runs. For more details, see the FAQ section
-   titled *Can the output of a prompt vary across runs in vLLM?* in the [FAQs](../../usage/faq.md).
-
-While vLLM strives to ensure losslessness in speculative decoding, variations in generated outputs with and without speculative decoding
-can occur due to following factors:
-
- **Floating-Point Precision**: Differences in hardware numerical precision may lead to slight discrepancies in the output distribution.
- **Batch Size and Numerical Stability**: Changes in batch size may cause variations in logprobs and output probabilities, potentially
-  due to non-deterministic behavior in batched operations or numerical instability.
-
-For mitigation strategies, please refer to the FAQ entry *Can the output of a prompt vary across runs in vLLM?* in the [FAQs](../../usage/faq.md).
-
-## Resources for vLLM contributors
-
- [A Hacker's Guide to Speculative Decoding in vLLM](https://www.youtube.com/watch?v=9wNAgpX6z_4)
- [What is Lookahead Scheduling in vLLM?](https://docs.google.com/document/d/1Z9TvqzzBPnh5WHcRwjvK2UEeFeq5zMZb5mFE8jR0HCs/edit#heading=h.1fjfb0donq5a)
- [Information on batch expansion](https://docs.google.com/document/d/1T-JaS2T1NRfdP51qzqpyakoCXxSXTtORppiwaj5asxA/edit#heading=h.kk7dq05lc6q8)
- [Dynamic speculative decoding](https://github.com/vllm-project/vllm/issues/4565)
@@ -0,0 +1,62 @@
+# Speculative Decoding
+
+This document shows how to use [Speculative Decoding](https://arxiv.org/pdf/2302.01318) with vLLM to reduce inter-token latency under medium-to-low QPS (query per second), memory-bound workloads.
+
+To train your own draft models for optimized speculative decoding, see [vllm-project/speculators](speculators.md) for seamless training and integration with vLLM.
+
+## vLLM Speculation Methods
+
+vLLM supports a variety of methods of speculative decoding. Model-based methods such as EAGLE, draft models, and mlp provide the best latency reduction, while simpler methods such as n-gram and and suffix decoding provide modest speedups without increasing workload during peak traffic.
+
+- [EAGLE](eagle.md)
+- [Draft Model](draft_model.md)
+- [Multi-Layer Perceptron](mlp.md)
+- [N-Gram](n_gram.md)
+- [Suffix Decoding](suffix.md)
+
+## Lossless guarantees of Speculative Decoding
+
+In vLLM, speculative decoding aims to enhance inference efficiency while maintaining accuracy. This section addresses the lossless guarantees of
+speculative decoding, breaking down the guarantees into three key areas:
+
+1. **Theoretical Losslessness**
+   \- Speculative decoding sampling is theoretically lossless up to the precision limits of hardware numerics. Floating-point errors might
+   cause slight variations in output distributions, as discussed
+   in [Accelerating Large Language Model Decoding with Speculative Sampling](https://arxiv.org/pdf/2302.01318)
+
+2. **Algorithmic Losslessness**
+   \- vLLM’s implementation of speculative decoding is algorithmically validated to be lossless. Key validation tests include:
+
+    > - **Rejection Sampler Convergence**: Ensures that samples from vLLM’s rejection sampler align with the target
+    >   distribution. [View Test Code](https://github.com/vllm-project/vllm/blob/47b65a550866c7ffbd076ecb74106714838ce7da/tests/samplers/test_rejection_sampler.py#L252)
+    > - **Greedy Sampling Equality**: Confirms that greedy sampling with speculative decoding matches greedy sampling
+    >   without it. This verifies that vLLM's speculative decoding framework, when integrated with the vLLM forward pass and the vLLM rejection sampler,
+    >   provides a lossless guarantee. Almost all of the tests in [tests/spec_decode/e2e](/tests/v1/spec_decode).
+    >   verify this property using [this assertion implementation](https://github.com/vllm-project/vllm/blob/b67ae00cdbbe1a58ffc8ff170f0c8d79044a684a/tests/spec_decode/e2e/conftest.py#L291)
+
+3. **vLLM Logprob Stability**
+   \- vLLM does not currently guarantee stable token log probabilities (logprobs). This can result in different outputs for the
+   same request across runs. For more details, see the FAQ section
+   titled *Can the output of a prompt vary across runs in vLLM?* in the [FAQs](../../usage/faq.md).
+
+While vLLM strives to ensure losslessness in speculative decoding, variations in generated outputs with and without speculative decoding
+can occur due to following factors:
+
+- **Floating-Point Precision**: Differences in hardware numerical precision may lead to slight discrepancies in the output distribution.
+- **Batch Size and Numerical Stability**: Changes in batch size may cause variations in logprobs and output probabilities, potentially
+  due to non-deterministic behavior in batched operations or numerical instability.
+
+For mitigation strategies, please refer to the FAQ entry *Can the output of a prompt vary across runs in vLLM?* in the [FAQs](../../usage/faq.md).
+
+## Known Feature Incompatibility
+
+1. Pipeline parallelism is not composible with speculative decoding as of `vllm<=0.15.0`
+2. Speculative decoding with a draft models is not supported in `vllm<=0.10.0`
+
+## Resources for vLLM contributors
+
+- [[vLLM Office Hours #40] Intro to Speculators](https://www.youtube.com/watch?v=2ISAr_JVGLs)
+- [A Hacker's Guide to Speculative Decoding in vLLM](https://www.youtube.com/watch?v=9wNAgpX6z_4)
+- [What is Lookahead Scheduling in vLLM?](https://docs.google.com/document/d/1Z9TvqzzBPnh5WHcRwjvK2UEeFeq5zMZb5mFE8jR0HCs/edit#heading=h.1fjfb0donq5a)
+- [Information on batch expansion](https://docs.google.com/document/d/1T-JaS2T1NRfdP51qzqpyakoCXxSXTtORppiwaj5asxA/edit#heading=h.kk7dq05lc6q8)
+- [Dynamic speculative decoding](https://github.com/vllm-project/vllm/issues/4565)
@@ -0,0 +1,80 @@
+# Draft Models
+
+The following code configures vLLM in an offline mode to use speculative decoding with a draft model, speculating 5 tokens at a time.
+
+```python
+from vllm import LLM, SamplingParams
+
+prompts = ["The future of AI is"]
+sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
+
+llm = LLM(
+    model="Qwen/Qwen3-8B",
+    tensor_parallel_size=1,
+    speculative_config={
+        "model": "Qwen/Qwen3-0.6B",
+        "num_speculative_tokens": 5,
+        "method": "draft_model",
+    },
+)
+outputs = llm.generate(prompts, sampling_params)
+
+for output in outputs:
+    prompt = output.prompt
+    generated_text = output.outputs[0].text
+    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
+```
+
+To perform the equivalent launch in online mode, use the following server-side code:
+
+```bash
+vllm serve Qwen/Qwen3-4B-Thinking-2507 \
+    --host 0.0.0.0 \
+    --port 8000 \
+    --seed 42 \
+    -tp 1 \
+    --max_model_len 2048 \
+    --gpu_memory_utilization 0.8 \
+    --speculative_config '{"model": "Qwen/Qwen3-0.6B", "num_speculative_tokens": 5, "method": "draft_model"}'
+```
+
+The code used to request as completions as a client remains unchanged:
+
+??? code
+
+    ```python
+    from openai import OpenAI
+
+    # Modify OpenAI's API key and API base to use vLLM's API server.
+    openai_api_key = "EMPTY"
+    openai_api_base = "http://localhost:8000/v1"
+
+    client = OpenAI(
+        # defaults to os.environ.get("OPENAI_API_KEY")
+        api_key=openai_api_key,
+        base_url=openai_api_base,
+    )
+
+    models = client.models.list()
+    model = models.data[0].id
+
+    # Completion API
+    stream = False
+    completion = client.completions.create(
+        model=model,
+        prompt="The future of AI is",
+        echo=False,
+        n=1,
+        stream=stream,
+    )
+
+    print("Completion results:")
+    if stream:
+        for c in completion:
+            print(c)
+    else:
+        print(completion)
+    ```
+
+!!! warning
+    Note: Please use `--speculative_config` to set all configurations related to speculative decoding. The previous method of specifying the model through `--speculative_model` and adding related parameters (e.g., `--num_speculative_tokens`) separately has been deprecated.
@@ -0,0 +1,67 @@
+# EAGLE Draft Models
+
+The following code configures vLLM to use speculative decoding where proposals are generated by an [EAGLE (Extrapolation Algorithm for Greater Language-model Efficiency)](https://arxiv.org/pdf/2401.15077) based draft model. A more detailed example for offline mode, including how to extract request level acceptance rate, can be found in [examples/offline_inference/spec_decode.py](../../../examples/offline_inference/spec_decode.py)
+
+## Eagle Drafter Example
+
+```python
+from vllm import LLM, SamplingParams
+
+prompts = ["The future of AI is"]
+sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
+
+llm = LLM(
+    model="meta-llama/Meta-Llama-3-8B-Instruct",
+    tensor_parallel_size=4,
+    speculative_config={
+        "model": "yuhuili/EAGLE-LLaMA3-Instruct-8B",
+        "draft_tensor_parallel_size": 1,
+        "num_speculative_tokens": 2,
+        "method": "eagle",
+    },
+)
+
+outputs = llm.generate(prompts, sampling_params)
+
+for output in outputs:
+    prompt = output.prompt
+    generated_text = output.outputs[0].text
+    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
+```
+
+## Eagle3 Drafter Example
+
+```python
+from vllm import LLM, SamplingParams
+
+prompts = ["The future of AI is"]
+sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
+
+llm = LLM(
+    model="meta-llama/Meta-Llama-3-8B-Instruct",
+    tensor_parallel_size=2,
+    speculative_config={
+        "model": "RedHatAI/Llama-3.1-8B-Instruct-speculator.eagle3",
+        "draft_tensor_parallel_size": 2,
+        "num_speculative_tokens": 2,
+        "method": "eagle3",
+    },
+)
+
+outputs = llm.generate(prompts, sampling_params)
+
+for output in outputs:
+    prompt = output.prompt
+    generated_text = output.outputs[0].text
+    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
+```
+
+## Pre-Trained Eagle Draft Models
+
+A variety of EAGLE draft models are available on the Hugging Face hub:
+
+* [RedHatAI/speculator-models](https://huggingface.co/collections/RedHatAI/speculator-models)
+* [yuhuili/models](https://huggingface.co/yuhuili/models?search=eagle)
+
+!!! warning
+    If you are using `vllm<0.7.0`, please use [this script](https://gist.github.com/abhigoyal1997/1e7a4109ccb7704fbc67f625e86b2d6d) to convert the speculative model and specify `"model": "path/to/modified/eagle/model"` in `speculative_config`.
@@ -0,0 +1,42 @@
+# MLP Draft Models
+
+The following code configures vLLM to use speculative decoding where proposals are generated by draft models that condition draft predictions on both context vectors and sampled tokens. For more information see [The Hitchhiker's Guide to Speculative Decoding](https://pytorch.org/blog/hitchhikers-guide-speculative-decoding/) and [IBM Research's Technical Report](https://arxiv.org/abs/2404.19124).
+
+## MLP Drafter Example
+
+```python
+from vllm import LLM, SamplingParams
+
+prompts = ["The future of AI is"]
+sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
+
+llm = LLM(
+    model="meta-llama/Meta-Llama-3.1-70B-Instruct",
+    tensor_parallel_size=4,
+    speculative_config={
+        "model": "ibm-ai-platform/llama3-70b-accelerator",
+        "draft_tensor_parallel_size": 1,
+        "method": "mlp_speculator",
+    },
+)
+outputs = llm.generate(prompts, sampling_params)
+
+for output in outputs:
+    prompt = output.prompt
+    generated_text = output.outputs[0].text
+    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
+```
+
+## Pre-Trained MLP Drafter Models
+
+A variety of speculative models of this type are available on HF hub:
+
+- [llama-13b-accelerator](https://huggingface.co/ibm-ai-platform/llama-13b-accelerator)
+- [llama3-8b-accelerator](https://huggingface.co/ibm-ai-platform/llama3-8b-accelerator)
+- [codellama-34b-accelerator](https://huggingface.co/ibm-ai-platform/codellama-34b-accelerator)
+- [llama2-70b-accelerator](https://huggingface.co/ibm-ai-platform/llama2-70b-accelerator)
+- [llama3-70b-accelerator](https://huggingface.co/ibm-ai-platform/llama3-70b-accelerator)
+- [granite-3b-code-instruct-accelerator](https://huggingface.co/ibm-granite/granite-3b-code-instruct-accelerator)
+- [granite-8b-code-instruct-accelerator](https://huggingface.co/ibm-granite/granite-8b-code-instruct-accelerator)
+- [granite-7b-instruct-accelerator](https://huggingface.co/ibm-granite/granite-7b-instruct-accelerator)
+- [granite-20b-code-instruct-accelerator](https://huggingface.co/ibm-granite/granite-20b-code-instruct-accelerator)
@@ -0,0 +1,27 @@
+# N-Gram Speculation
+
+The following code configures vLLM to use speculative decoding where proposals are generated by
+matching n-grams in the prompt. For more information read [this thread.](https://x.com/joao_gante/status/1747322413006643259)
+
+```python
+from vllm import LLM, SamplingParams
+
+prompts = ["The future of AI is"]
+sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
+
+llm = LLM(
+    model="Qwen/Qwen3-8B",
+    tensor_parallel_size=1,
+    speculative_config={
+        "method": "ngram",
+        "num_speculative_tokens": 5,
+        "prompt_lookup_max": 4,
+    },
+)
+outputs = llm.generate(prompts, sampling_params)
+
+for output in outputs:
+    prompt = output.prompt
+    generated_text = output.outputs[0].text
+    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
+```
@@ -1,4 +1,7 @@
-# Speculators
+# vLLM-Project/Speculators
+
+![User Flow Light](../../assets/features/speculative_decoding/speculators-user-flow-light.svg#only-light)
+![User Flow Dark](../../assets/features/speculative_decoding/speculators-user-flow-dark.svg#only-dark)

 [Speculators](https://docs.vllm.ai/projects/speculators/en/latest/) is a library for accelerating LLM inference through speculative decoding, providing efficient draft model training that integrates seamlessly with vLLM to reduce latency and improve throughput.

@@ -0,0 +1,35 @@
+# Suffix Decoding
+
+The following code configures vLLM to use speculative decoding where proposals are generated using Suffix Decoding ([technical report](https://arxiv.org/abs/2411.04975)).
+
+Like n-gram, Suffix Decoding can generate draft tokens by pattern-matching using the last `n` generated tokens. Unlike n-gram, Suffix Decoding (1) can pattern-match against both the prompt and previous generations, (2) uses frequency counts to propose the most likely continuations, and (3) speculates an adaptive number of tokens for each request at each iteration to get better acceptance rates.
+
+Suffix Decoding can achieve better performance for tasks with high repetition, such as code-editing, agentic loops (e.g. self-reflection, self-consistency), and RL rollouts.
+
+!!! tip "Install Arctic Inference"
+    Suffix Decoding requires [Arctic Inference](https://github.com/snowflakedb/ArcticInference). You can install it with `pip install arctic-inference`.
+
+!!! tip "Suffix Decoding Speculative Tokens"
+    Suffix Decoding will speculate a dynamic number of tokens for each request at each decoding step, so the `num_speculative_tokens` configuration specifies the *maximum* number of speculative tokens. It is suggested to use a high number such as `16` or `32` (default).
+
+```python
+from vllm import LLM, SamplingParams
+
+prompts = ["The future of AI is"]
+sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
+
+llm = LLM(
+    model="Qwen/Qwen3-8B",
+    tensor_parallel_size=1,
+    speculative_config={
+        "method": "suffix",
+        "num_speculative_tokens": 32,
+    },
+)
+outputs = llm.generate(prompts, sampling_params)
+
+for output in outputs:
+    prompt = output.prompt
+    generated_text = output.outputs[0].text
+    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
+```
@@ -22,6 +22,7 @@ METRIC_SOURCE_FILES = [
        "path": "vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py",
        "output": "nixl_connector.inc.md",
    },
+    {"path": "vllm/v1/metrics/perf.py", "output": "perf.inc.md"},
 ]


@@ -382,6 +382,7 @@ ColQwen3 is based on [ColPali](https://arxiv.org/abs/2407.01449), which extends
 |---|---|---|
 | `ColQwen3` | Qwen3-VL | `TomoroAI/tomoro-colqwen3-embed-4b`, `TomoroAI/tomoro-colqwen3-embed-8b` |
 | `OpsColQwen3Model` | Qwen3-VL | `OpenSearch-AI/Ops-Colqwen3-4B`, `OpenSearch-AI/Ops-Colqwen3-8B` |
+| `Qwen3VLNemotronEmbedModel` | Qwen3-VL | `nvidia/nemotron-colembed-vl-4b-v2`, `nvidia/nemotron-colembed-vl-8b-v2` |

 Start the server:

@@ -389,7 +390,9 @@ Start the server:
 vllm serve TomoroAI/tomoro-colqwen3-embed-4b --max-model-len 4096
 ```

-Then you can use the rerank endpoint:
+#### Text-only scoring and reranking
+
+Use the `/rerank` endpoint:

 ```shell
 curl -s http://localhost:8000/rerank -H "Content-Type: application/json" -d '{
@@ -403,7 +406,7 @@ curl -s http://localhost:8000/rerank -H "Content-Type: application/json" -d '{
 }'
 ```

-Or the score endpoint:
+Or the `/score` endpoint:

 ```shell
 curl -s http://localhost:8000/score -H "Content-Type: application/json" -d '{
@@ -413,7 +416,57 @@ curl -s http://localhost:8000/score -H "Content-Type: application/json" -d '{
 }'
 ```

-You can also get the raw token embeddings using the pooling endpoint with `token_embed` task:
+#### Multi-modal scoring and reranking (text query × image documents)
+
+The `/score` and `/rerank` endpoints also accept multi-modal inputs directly.
+Pass image documents using the `data_1`/`data_2` (for `/score`) or `documents` (for `/rerank`) fields
+with a `content` list containing `image_url` and `text` parts — the same format used by the
+OpenAI chat completion API:
+
+Score a text query against image documents:
+
+```shell
+curl -s http://localhost:8000/score -H "Content-Type: application/json" -d '{
+    "model": "TomoroAI/tomoro-colqwen3-embed-4b",
+    "data_1": "Retrieve the city of Beijing",
+    "data_2": [
+        {
+            "content": [
+                {"type": "image_url", "image_url": {"url": "data:image/png;base64,<BASE64>"}},
+                {"type": "text", "text": "Describe the image."}
+            ]
+        }
+    ]
+}'
+```
+
+Rerank image documents by a text query:
+
+```shell
+curl -s http://localhost:8000/rerank -H "Content-Type: application/json" -d '{
+    "model": "TomoroAI/tomoro-colqwen3-embed-4b",
+    "query": "Retrieve the city of Beijing",
+    "documents": [
+        {
+            "content": [
+                {"type": "image_url", "image_url": {"url": "data:image/png;base64,<BASE64_1>"}},
+                {"type": "text", "text": "Describe the image."}
+            ]
+        },
+        {
+            "content": [
+                {"type": "image_url", "image_url": {"url": "data:image/png;base64,<BASE64_2>"}},
+                {"type": "text", "text": "Describe the image."}
+            ]
+        }
+    ],
+    "top_n": 2
+}'
+```
+
+#### Raw token embeddings
+
+You can also get the raw token embeddings using the `/pooling` endpoint with `token_embed` task:

 ```shell
 curl -s http://localhost:8000/pooling -H "Content-Type: application/json" -d '{
@@ -423,7 +476,7 @@ curl -s http://localhost:8000/pooling -H "Content-Type: application/json" -d '{
 }'
 ```

-For **image inputs**, use the chat-style `messages` field so that the vLLM multimodal processor handles them correctly:
+For **image inputs** via the pooling endpoint, use the chat-style `messages` field:

 ```shell
 curl -s http://localhost:8000/pooling -H "Content-Type: application/json" -d '{
@@ -440,10 +493,10 @@ curl -s http://localhost:8000/pooling -H "Content-Type: application/json" -d '{
 }'
 ```

-Examples can be found here:
+#### Examples

 - Multi-vector retrieval: [examples/pooling/token_embed/colqwen3_token_embed_online.py](../../examples/pooling/token_embed/colqwen3_token_embed_online.py)
- Reranking: [examples/pooling/score/colqwen3_rerank_online.py](../../examples/pooling/score/colqwen3_rerank_online.py)
+- Reranking (text + multi-modal): [examples/pooling/score/colqwen3_rerank_online.py](../../examples/pooling/score/colqwen3_rerank_online.py)

 ### BAAI/bge-m3

@@ -821,6 +821,7 @@ The following table lists those that are tested in vLLM.
 | Architecture | Models | Inputs | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) |
 |--------------|--------|--------|-------------------|----------------------|---------------------------|
 | `CLIPModel` | CLIP | T / I | `openai/clip-vit-base-patch32`, `openai/clip-vit-large-patch14`, etc. | | |
+| `ColModernVBertForRetrieval` | ColModernVBERT | T / I | `ModernVBERT/colmodernvbert-merged` | | |
 | `LlavaNextForConditionalGeneration`<sup>C</sup> | LLaVA-NeXT-based | T / I | `royokong/e5-v` | | ✅︎ |
 | `Phi3VForCausalLM`<sup>C</sup> | Phi-3-Vision-based | T + I | `TIGER-Lab/VLM2Vec-Full` | | ✅︎ |
 | `Qwen3VLForConditionalGeneration`<sup>C</sup> | Qwen3-VL | T + I + V | `Qwen/Qwen3-VL-Embedding-2B`, etc. | ✅︎ | ✅︎ |
@@ -45,6 +45,12 @@ The following metrics are exposed:

 --8<-- "docs/generated/metrics/nixl_connector.inc.md"

+## Model Flops Utilization (MFU) Performance Metrics
+
+These metrics are available via `--enable-mfu-metrics`:
+
+--8<-- "docs/generated/metrics/perf.inc.md"
+
 ## Deprecation Policy

 Note: when metrics are deprecated in version `X.Y`, they are hidden in version `X.Y+1`
@@ -42,6 +42,7 @@ def main():
                "async_load": args.async_load,
            },
            kv_connector_module_path="load_recovery_example_connector",
+            kv_load_failure_policy="recompute",
        )
        out_file = (
            "async_decode_recovered_output.txt"
@@ -26,14 +26,12 @@ workloads. Residual GPU activity interferes with vLLM memory profiling and
 causes unexpected behavior.
 """

-import os
+import asyncio
 import uuid
 from dataclasses import asdict

 import ray
 import torch
-from ray.util.placement_group import placement_group
-from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
 from transformers import AutoModelForCausalLM, AutoTokenizer

 import vllm
@@ -51,14 +49,15 @@ from vllm.distributed.weight_transfer.nccl_engine import (
 from vllm.utils.network_utils import get_ip, get_open_port
 from vllm.v1.executor import Executor

-MODEL_NAME = "facebook/opt-125m"
+MODEL_NAME_V1 = "Qwen/Qwen3-1.7B-Base"
+MODEL_NAME_V2 = "Qwen/Qwen3-1.7B"
+PAUSE_TOKEN_THRESHOLD = 10


 class MyLLM(vllm.AsyncLLMEngine):
    """Configure the vLLM worker for Ray placement group execution."""

    def __init__(self, **kwargs):
-        os.environ["VLLM_RAY_BUNDLE_INDICES"] = "0,1"
        engine_args = vllm.AsyncEngineArgs(**kwargs)
        vllm_config = engine_args.create_engine_config()
        executor_class = Executor.get_class(vllm_config)
@@ -68,26 +67,44 @@ class MyLLM(vllm.AsyncLLMEngine):
            log_requests=engine_args.enable_log_requests,
            log_stats=not engine_args.disable_log_stats,
        )
+        self._generation_paused = False
+        self._request_pause_flag = False

-    async def generate_with_retry(
+    async def do_generate(
        self, prompt_token_ids: list[int], sampling_params: vllm.SamplingParams
-    ) -> vllm.RequestOutput:
-        finish_reason = "abort"
-        while finish_reason == "abort":
-            async for request_output in self.generate(
-                {"prompt_token_ids": prompt_token_ids},
-                sampling_params,
-                request_id=str(uuid.uuid4()),
+    ) -> tuple[vllm.RequestOutput, int]:
+        """Generate a single request, setting the request pause flag once the
+        token count reaches the threshold.
+
+        Returns (output, pause_token_index). pause_token_index is the number
+        of tokens generated before the weight change, or -1 if no pause.
+        """
+        pause_token_index = -1
+        prev_token_count = 0
+        async for request_output in self.generate(
+            {"prompt_token_ids": prompt_token_ids},
+            sampling_params,
+            request_id=str(uuid.uuid4()),
+        ):
+            output = request_output
+            cur_token_count = len(output.outputs[0].token_ids)
+            if (
+                cur_token_count >= PAUSE_TOKEN_THRESHOLD
+                and not self._request_pause_flag
            ):
-                output = request_output
-            finish_reason = output.outputs[0].finish_reason
-            if finish_reason == "abort":
-                print(
-                    f"ABORT, prompt_token_ids: {prompt_token_ids}, "
-                    f"generated token_ids: {list(output.outputs[0].token_ids)}"
-                )
-            prompt_token_ids = prompt_token_ids + list(output.outputs[0].token_ids)
-        return output
+                self._request_pause_flag = True
+            if self._generation_paused and pause_token_index == -1:
+                pause_token_index = prev_token_count
+            prev_token_count = cur_token_count
+        return output, pause_token_index
+
+    async def pause_after_n_tokens(self):
+        """Wait for any request to set the pause flag, then pause."""
+        while not self._request_pause_flag:
+            await asyncio.sleep(0)
+        await super().pause_generation(mode="keep")
+        await asyncio.sleep(0.2)
+        self._generation_paused = True


@ray.remote(num_gpus=1)
@@ -95,6 +112,14 @@ class TrainModel:
    """Ray actor that wraps the training model on a dedicated GPU."""

    def __init__(self, model_name: str):
+        from vllm.model_executor.layers.batch_invariant import (
+            init_batch_invariance,
+        )
+        from vllm.v1.attention.backends.registry import AttentionBackendEnum
+
+        # need to init all env vars for batch invariance which affect nccl ops
+        init_batch_invariance(AttentionBackendEnum.FLASH_ATTN)
+
        self.model = AutoModelForCausalLM.from_pretrained(
            model_name, dtype=torch.bfloat16
        ).to("cuda:0")
@@ -133,70 +158,80 @@ class TrainModel:
            packed=packed,
        )

+    @torch.inference_mode()
+    def generate(self, token_ids: list[int], max_new_tokens: int) -> list[int]:
+        """Greedy-decode max_new_tokens from the given context."""
+        input_ids = torch.tensor([token_ids], device="cuda:0")
+        output = self.model.generate(
+            input_ids,
+            max_new_tokens=max_new_tokens,
+            do_sample=False,
+        )
+        new_token_ids = output[0, len(token_ids) :].tolist()
+        return new_token_ids

-# Initialize Ray and set the visible devices. The vLLM engine will
-# be placed on GPUs 1 and 2.
-ray.init()
+
+ray.init(
+    runtime_env={
+        "env_vars": {
+            # enable batch invariance for deterministic outputs
+            "VLLM_BATCH_INVARIANT": "1",
+            # prevent ray from setting CUDA_VISIBLE_DEVICES
+            "RAY_EXPERIMENTAL_NOSET_CUDA_ENV_VAR": "1",
+        }
+    }
+)

 # Launch the training model actor. Ray's resource scheduler will allocate
 # 1 GPU (via num_gpus=1 in the decorator), ensuring pg_inference gets different GPUs.
-train_model = TrainModel.remote(MODEL_NAME)
-
-# Create a placement group that reserves GPU 1–2 for the vLLM inference engine.
-# Learn more about Ray placement groups:
-# https://docs.ray.io/en/latest/placement-groups.html
-
-pg_inference = placement_group([{"GPU": 1, "CPU": 0}] * 2)
-ray.get(pg_inference.ready())
-scheduling_inference = PlacementGroupSchedulingStrategy(
-    placement_group=pg_inference,
-    placement_group_capture_child_tasks=True,
-    placement_group_bundle_index=0,
-)
+train_model = TrainModel.remote(MODEL_NAME_V2)

 # Launch the vLLM inference engine. The `enforce_eager` flag reduces
 # start-up latency.
-# Note: Weight transfer APIs (init_weight_transfer_engine, update_weights)
-# are now native to vLLM workers.
+# With data_parallel_backend="ray", vLLM's CoreEngineActorManager creates
+# its own placement groups internally for each DP rank, so we must NOT
+# create an outer placement group (it would reserve GPUs and hide them
+# from the internal DP resource check).
 llm = ray.remote(
    num_cpus=0,
    num_gpus=0,
-    scheduling_strategy=scheduling_inference,
 )(MyLLM).remote(
-    model=MODEL_NAME,
+    model=MODEL_NAME_V1,
    enforce_eager=True,
-    tensor_parallel_size=2,
+    max_model_len=8192,
    distributed_executor_backend="ray",
-    load_format="dummy",
+    attention_backend="FLASH_ATTN",
+    gpu_memory_utilization=0.75,
    weight_transfer_config=WeightTransferConfig(backend="nccl"),
 )

-# Generate text from the prompts.
-prompts = [
-    "My name is",
+PROMPTS = [
    "The president of the United States is",
    "The capital of France is",
-    "The future of AI is",
+    "The largest ocean on Earth is",
+    "The speed of light in a vacuum is",
+    "The chemical formula for water is",
+    "The tallest mountain in the world is",
+    "The first person to walk on the moon was",
+    "The Great Wall of China was built to",
+    "Photosynthesis is the process by which",
+    "The theory of general relativity was proposed by",
+    "The boiling point of water at sea level is",
+    "The largest planet in our solar system is",
+    "DNA stands for deoxyribonucleic acid and it",
 ]

-# Tokenize prompts to token IDs
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-prompt_token_ids_list = [
-    tokenizer.encode(prompt, add_special_tokens=False) for prompt in prompts
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME_V1)
+batch_prompt_token_ids = [
+    tokenizer.encode(prompt, add_special_tokens=False) for prompt in PROMPTS
 ]

-sampling_params = [
-    SamplingParams(temperature=0, max_tokens=2),
-    SamplingParams(temperature=0, max_tokens=32),
-    SamplingParams(temperature=0, max_tokens=32),
-    SamplingParams(temperature=0, max_tokens=32),
-]

 # Set up the communication channel between the training process and the
 # inference engine.
 master_address, master_port = ray.get(train_model.get_master_address_and_port.remote())

-world_size = 3  # 1 trainer + 2 inference workers (tensor_parallel_size=2)
+world_size = 2  # 1 trainer + 1 inference worker
 inference_handle = llm.init_weight_transfer_engine.remote(
    WeightTransferInitRequest(
        init_info=asdict(
@@ -215,22 +250,28 @@ train_handle = train_model.init_weight_transfer_group.remote(world_size)
 ray.get([train_handle, inference_handle])


-generation_futures = [
-    llm.generate_with_retry.remote(prompt_token_ids, params)
-    for prompt_token_ids, params in zip(prompt_token_ids_list, sampling_params)
-]
+N_NEW_TOKENS = 100

-finished, pending = ray.wait(generation_futures, num_returns=1)
-
-# Pause generation in preparation for weight sync
-ray.get(llm.pause_generation.remote(wait_for_inflight_requests=False))
-
-# Synchronize the updated weights to the inference engine using batched API.
-# Collect all weight metadata from the training actor
+# Collect weight metadata once
 names, dtype_names, shapes = ray.get(train_model.get_weight_metadata.remote())

-# Issue update_weights call with NCCL-specific update info
-# packed=True enables efficient batched tensor broadcasting
+# ── Phase 1: concurrent requests with weight sync ───────────────────
+print(f"\n{'=' * 50}")
+print(f"Prompts ({len(PROMPTS)}):")
+for p in PROMPTS:
+    print(f"  - {p!r}")
+print(f"{'=' * 50}")
+
+sampling_params = SamplingParams(
+    temperature=0, max_tokens=PAUSE_TOKEN_THRESHOLD + N_NEW_TOKENS
+)
+
+gen_futures = [
+    llm.do_generate.remote(ptids, sampling_params) for ptids in batch_prompt_token_ids
+]
+
+ray.get(llm.pause_after_n_tokens.remote())
+
 inference_handle = llm.update_weights.remote(
    WeightTransferUpdateRequest(
        update_info=asdict(
@@ -243,41 +284,76 @@ inference_handle = llm.update_weights.remote(
        )
    )
 )
-
-# Broadcast all weights from trainer using the weight transfer API
 train_handle = train_model.broadcast_weights.remote(packed=True)
 ray.get([train_handle, inference_handle])

-# Resume generation since weight sync is complete
 ray.get(llm.resume_generation.remote())
+results = ray.get(gen_futures)

-# Get outputs separately - finished completed before pause, pending were paused/resumed
-finished_outputs = ray.get(finished)
-pending_outputs = ray.get(pending)
+for i, (output, pause_idx) in enumerate(results):
+    all_token_ids = list(output.outputs[0].token_ids)
+    before_text = tokenizer.decode(all_token_ids[:pause_idx])
+    after_text = tokenizer.decode(all_token_ids[pause_idx:])
+    print(f"\n  Request {i} ({PROMPTS[i]!r}):")
+    print(f"    Old weights ({pause_idx} tokens): {before_text!r}")
+    n_after = len(all_token_ids) - pause_idx
+    print(f"    New weights ({n_after} tokens): {after_text!r}")

-# Requests that finished before the pause: all generation used original weights
-print("-" * 50)
-print("Requests that completed BEFORE weight change:")
-print("-" * 50)
-for output in finished_outputs:
-    prompt_text = tokenizer.decode(output.prompt_token_ids)
-    print(f"Prompt: {prompt_text!r}")
-    print(f"Generated (with original weights): {output.outputs[0].text!r}")
-    print("-" * 50)
+# ── Phase 2: validate with a fresh V2 vLLM instance ────────────────
+print(f"\n{'=' * 50}")
+print("VALIDATION: comparing weight-synced vLLM with fresh V2 instance")
+print(f"{'=' * 50}")

-# Requests that were paused mid-generation: some text before, some after weight change
-print("Requests that were PAUSED and RESUMED after weight change:")
-print("-" * 50)
-for output in pending_outputs:
-    # Decode the full prompt token IDs (original + generated before pause)
-    full_prompt_text = tokenizer.decode(output.prompt_token_ids)
-    # Find the original prompt by checking which one this output started with
-    original_prompt = next(p for p in prompts if full_prompt_text.startswith(p))
-    # output.prompt_token_ids contains original prompt + tokens generated before pause
-    # output.outputs[0].text is what was generated after resuming with new weights
-    text_before_pause = full_prompt_text[len(original_prompt) :]
-    text_after_pause = output.outputs[0].text
-    print(f"Original prompt: {original_prompt!r}")
-    print(f"Generated before weight change: {text_before_pause!r}")
-    print(f"Generated after weight change: {text_after_pause!r}")
-    print("-" * 50)
+ray.get(llm.shutdown.remote())
+ray.kill(llm)
+ray.kill(train_model)
+
+llm_v2 = ray.remote(
+    num_cpus=0,
+    num_gpus=0,
+)(MyLLM).remote(
+    model=MODEL_NAME_V2,
+    enforce_eager=True,
+    max_model_len=8192,
+    gpu_memory_utilization=0.75,
+    distributed_executor_backend="ray",
+    attention_backend="FLASH_ATTN",
+)
+
+val_futures = [
+    llm_v2.do_generate.remote(
+        list(output.prompt_token_ids) + list(output.outputs[0].token_ids)[:pause_idx],
+        SamplingParams(
+            temperature=0, max_tokens=len(output.outputs[0].token_ids) - pause_idx
+        ),
+    )
+    for output, pause_idx in results
+]
+val_results = ray.get(val_futures)
+
+all_pass = True
+for i, ((output, pause_idx), (val_output, _)) in enumerate(zip(results, val_results)):
+    expected = list(output.outputs[0].token_ids)[pause_idx:]
+    actual = list(val_output.outputs[0].token_ids)
+    match = actual == expected
+
+    if match:
+        print(f"  [PASS] {PROMPTS[i]!r}")
+    else:
+        all_pass = False
+        print(f"  [FAIL] {PROMPTS[i]!r}")
+        print(f"         weight-synced vLLM: {tokenizer.decode(expected)!r}")
+        print(f"         V2 vLLM:           {tokenizer.decode(actual)!r}")
+        for j, (e, a) in enumerate(zip(expected, actual)):
+            if e != a:
+                print(
+                    f"         first divergence at output token {j}: "
+                    f"expected {e} ({tokenizer.decode([e])!r}) vs "
+                    f"actual {a} ({tokenizer.decode([a])!r})"
+                )
+                break
+
+ray.get(llm_v2.shutdown.remote())
+ray.kill(llm_v2)
+assert all_pass, "Some prompts failed validation, see above for details"
+print("=" * 50)
@@ -37,6 +37,12 @@ class BlockStored(KVCacheEvent):
    medium: str | None
    lora_name: str | None

+    extra_keys: list[tuple[Any, ...] | None] | None = None
+    """Extra keys used in block hash computation, one entry per block in
+    block_hashes. Each entry contains MM identifiers, LoRA name, cache_salt,
+    prompt embeddings data, etc. for that specific block.
+    """
+

 class BlockRemoved(KVCacheEvent):
    block_hashes: list[ExternalBlockHash]
@@ -1,14 +1,6 @@
 # Setup OpenTelemetry POC

-1. Install OpenTelemetry packages:
-
-    ```bash
-    pip install \
-      'opentelemetry-sdk>=1.26.0,<1.27.0' \
-      'opentelemetry-api>=1.26.0,<1.27.0' \
-      'opentelemetry-exporter-otlp>=1.26.0,<1.27.0' \
-      'opentelemetry-semantic-conventions-ai>=0.4.1,<0.5.0'
-    ```
+> **Note:** The core OpenTelemetry packages (`opentelemetry-sdk`, `opentelemetry-api`, `opentelemetry-exporter-otlp`, `opentelemetry-semantic-conventions-ai`) are bundled with vLLM. Manual installation is not required.

 1. Start Jaeger in a docker container:

@@ -0,0 +1,166 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Example of using ColModernVBERT late interaction model for reranking.
+
+ColModernVBERT is a multi-modal ColBERT-style model combining a SigLIP
+vision encoder with a ModernBERT text encoder. It produces per-token
+embeddings and uses MaxSim scoring for retrieval and reranking.
+Supports both text and image inputs.
+
+Start the server with:
+    vllm serve ModernVBERT/colmodernvbert-merged --max-model-len 8192
+
+Then run this script:
+    python colmodernvbert_rerank_online.py
+"""
+
+import requests
+
+MODEL = "ModernVBERT/colmodernvbert-merged"
+BASE_URL = "http://127.0.0.1:8000"
+
+headers = {"accept": "application/json", "Content-Type": "application/json"}
+
+IMAGE_URL = "https://upload.wikimedia.org/wikipedia/commons/thumb/4/47/PNG_transparency_demonstration_1.png/300px-PNG_transparency_demonstration_1.png"  # noqa: E501
+
+
+def rerank_text():
+    """Text-only reranking via /rerank endpoint."""
+    print("=" * 60)
+    print("1. Text reranking (/rerank)")
+    print("=" * 60)
+
+    data = {
+        "model": MODEL,
+        "query": "What is machine learning?",
+        "documents": [
+            "Machine learning is a subset of artificial intelligence.",
+            "Python is a programming language.",
+            "Deep learning uses neural networks for complex tasks.",
+            "The weather today is sunny.",
+        ],
+    }
+
+    response = requests.post(f"{BASE_URL}/rerank", headers=headers, json=data)
+
+    if response.status_code == 200:
+        result = response.json()
+        print("\n  Ranked documents (most relevant first):")
+        for item in result["results"]:
+            doc_idx = item["index"]
+            score = item["relevance_score"]
+            print(f"    [{score:.4f}] {data['documents'][doc_idx]}")
+    else:
+        print(f"  Request failed: {response.status_code}")
+        print(f"  {response.text[:300]}")
+
+
+def score_text():
+    """Text-only scoring via /score endpoint."""
+    print()
+    print("=" * 60)
+    print("2. Text scoring (/score)")
+    print("=" * 60)
+
+    query = "What is the capital of France?"
+    documents = [
+        "The capital of France is Paris.",
+        "Berlin is the capital of Germany.",
+        "Python is a programming language.",
+    ]
+
+    data = {
+        "model": MODEL,
+        "text_1": query,
+        "text_2": documents,
+    }
+
+    response = requests.post(f"{BASE_URL}/score", headers=headers, json=data)
+
+    if response.status_code == 200:
+        result = response.json()
+        print(f"\n  Query: {query}\n")
+        for item in result["data"]:
+            idx = item["index"]
+            score = item["score"]
+            print(f"    Doc {idx} (score={score:.4f}): {documents[idx]}")
+    else:
+        print(f"  Request failed: {response.status_code}")
+        print(f"  {response.text[:300]}")
+
+
+def score_text_top_n():
+    """Text reranking with top_n filtering via /rerank endpoint."""
+    print()
+    print("=" * 60)
+    print("3. Text reranking with top_n=2 (/rerank)")
+    print("=" * 60)
+
+    data = {
+        "model": MODEL,
+        "query": "What is the capital of France?",
+        "documents": [
+            "The capital of France is Paris.",
+            "Berlin is the capital of Germany.",
+            "Python is a programming language.",
+            "The Eiffel Tower is in Paris.",
+        ],
+        "top_n": 2,
+    }
+
+    response = requests.post(f"{BASE_URL}/rerank", headers=headers, json=data)
+
+    if response.status_code == 200:
+        result = response.json()
+        print(f"\n  Top {data['top_n']} results:")
+        for item in result["results"]:
+            doc_idx = item["index"]
+            score = item["relevance_score"]
+            print(f"    [{score:.4f}] {data['documents'][doc_idx]}")
+    else:
+        print(f"  Request failed: {response.status_code}")
+        print(f"  {response.text[:300]}")
+
+
+def rerank_multimodal():
+    """Multimodal reranking with text and image documents via /rerank."""
+    print()
+    print("=" * 60)
+    print("4. Multimodal reranking: text query vs image document (/rerank)")
+    print("=" * 60)
+
+    data = {
+        "model": MODEL,
+        "query": "A colorful logo with transparency",
+        "documents": [
+            {"content": [{"type": "image_url", "image_url": {"url": IMAGE_URL}}]},
+            "Python is a programming language.",
+            "The weather today is sunny.",
+        ],
+    }
+
+    response = requests.post(f"{BASE_URL}/rerank", headers=headers, json=data)
+
+    if response.status_code == 200:
+        result = response.json()
+        print("\n  Ranked documents (most relevant first):")
+        labels = ["[image]", "Python doc", "Weather doc"]
+        for item in result["results"]:
+            doc_idx = item["index"]
+            score = item["relevance_score"]
+            print(f"    [{score:.4f}] {labels[doc_idx]}")
+    else:
+        print(f"  Request failed: {response.status_code}")
+        print(f"  {response.text[:300]}")
+
+
+def main():
+    rerank_text()
+    score_text()
+    score_text_top_n()
+    rerank_multimodal()
+
+
+if __name__ == "__main__":
+    main()
@@ -1,7 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# ruff: noqa: E501
 """
-Example of using ColQwen3 late interaction model for reranking.
+Example of using ColQwen3 late interaction model for reranking and scoring.

 ColQwen3 is a multi-modal ColBERT-style model based on Qwen3-VL.
 It produces per-token embeddings and uses MaxSim scoring for retrieval
@@ -14,13 +15,65 @@ Then run this script:
    python colqwen3_rerank_online.py
 """

+import base64
+from io import BytesIO
+
 import requests
+from PIL import Image

 MODEL = "TomoroAI/tomoro-colqwen3-embed-4b"
 BASE_URL = "http://127.0.0.1:8000"

 headers = {"accept": "application/json", "Content-Type": "application/json"}

+# ── Image helpers ──────────────────────────────────────────
+
+
+def load_image(url: str) -> Image.Image:
+    """Download an image from URL (handles Wikimedia 403)."""
+    for hdrs in (
+        {},
+        {"User-Agent": "Mozilla/5.0 (compatible; ColQwen3-demo/1.0)"},
+    ):
+        resp = requests.get(url, headers=hdrs, timeout=15)
+        if resp.status_code == 403:
+            continue
+        resp.raise_for_status()
+        return Image.open(BytesIO(resp.content)).convert("RGB")
+    raise RuntimeError(f"Could not fetch image from {url}")
+
+
+def encode_image_base64(image: Image.Image) -> str:
+    """Encode a PIL image to a base64 data URI."""
+    buf = BytesIO()
+    image.save(buf, format="PNG")
+    return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()
+
+
+def make_image_content(image_url: str, text: str = "Describe the image.") -> dict:
+    """Build a ScoreMultiModalParam dict from an image URL."""
+    image = load_image(image_url)
+    return {
+        "content": [
+            {
+                "type": "image_url",
+                "image_url": {"url": encode_image_base64(image)},
+            },
+            {"type": "text", "text": text},
+        ]
+    }
+
+
+# ── Sample image URLs ─────────────────────────────────────
+
+IMAGE_URLS = {
+    "beijing": "https://upload.wikimedia.org/wikipedia/commons/6/61/Beijing_skyline_at_night.JPG",
+    "london": "https://upload.wikimedia.org/wikipedia/commons/4/49/London_skyline.jpg",
+    "singapore": "https://upload.wikimedia.org/wikipedia/commons/2/27/Singapore_skyline_2022.jpg",
+}
+
+# ── Text-only examples ────────────────────────────────────
+

 def rerank_text():
    """Text-only reranking via /rerank endpoint."""
@@ -120,11 +173,86 @@ def score_text_top_n():
        print(f"  {response.text[:300]}")


+# ── Multi-modal examples (text query × image documents) ──
+
+
+def score_text_vs_images():
+    """Score a text query against image documents via /score."""
+    print()
+    print("=" * 60)
+    print("4. Multi-modal scoring: text query vs image docs (/score)")
+    print("=" * 60)
+
+    query = "Retrieve the city of Beijing"
+    labels = list(IMAGE_URLS.keys())
+    print(f"\n  Loading {len(labels)} images...")
+    image_contents = [make_image_content(IMAGE_URLS[name]) for name in labels]
+
+    data = {
+        "model": MODEL,
+        "data_1": query,
+        "data_2": image_contents,
+    }
+
+    response = requests.post(f"{BASE_URL}/score", headers=headers, json=data)
+
+    if response.status_code == 200:
+        result = response.json()
+        print(f'\n  Query: "{query}"\n')
+        for item in result["data"]:
+            idx = item["index"]
+            print(f"    Doc {idx} [{labels[idx]}] score={item['score']:.4f}")
+    else:
+        print(f"  Request failed: {response.status_code}")
+        print(f"  {response.text[:300]}")
+
+
+def rerank_text_vs_images():
+    """Rerank image documents by a text query via /rerank."""
+    print()
+    print("=" * 60)
+    print("5. Multi-modal reranking: text query vs image docs (/rerank)")
+    print("=" * 60)
+
+    query = "Retrieve the city of London"
+    labels = list(IMAGE_URLS.keys())
+    print(f"\n  Loading {len(labels)} images...")
+    image_contents = [make_image_content(IMAGE_URLS[name]) for name in labels]
+
+    data = {
+        "model": MODEL,
+        "query": query,
+        "documents": image_contents,
+        "top_n": 2,
+    }
+
+    response = requests.post(f"{BASE_URL}/rerank", headers=headers, json=data)
+
+    if response.status_code == 200:
+        result = response.json()
+        print(f'\n  Query: "{query}"')
+        print(f"  Top {data['top_n']} results:\n")
+        for item in result["results"]:
+            idx = item["index"]
+            print(f"    [{item['relevance_score']:.4f}] {labels[idx]}")
+    else:
+        print(f"  Request failed: {response.status_code}")
+        print(f"  {response.text[:300]}")
+
+
+# ── Main ──────────────────────────────────────────────────
+
+
 def main():
+    # Text-only
    rerank_text()
    score_text()
    score_text_top_n()

+    # Multi-modal (text query × image documents)
+    score_text_vs_images()
+    rerank_text_vs_images()
+

 if __name__ == "__main__":
    main()
@@ -105,6 +105,10 @@ plugins:
          - https://numpy.org/doc/stable/objects.inv
          - https://pytorch.org/docs/stable/objects.inv
          - https://psutil.readthedocs.io/en/stable/objects.inv
+  - redirects:
+      redirect_maps:
+        features/spec_decode/README.md: features/speculative_decoding/README.md
+        features/spec_decode/speculators.md: features/speculative_decoding/speculators.md

 markdown_extensions:
  - attr_list
@@ -53,3 +53,7 @@ model-hosting-container-standards >= 0.1.13, < 1.0.0
 mcp
 grpcio
 grpcio-reflection
+opentelemetry-sdk >= 1.27.0
+opentelemetry-api >= 1.27.0
+opentelemetry-exporter-otlp >= 1.27.0
+opentelemetry-semantic-conventions-ai >= 0.4.1
@@ -10,4 +10,4 @@ torchaudio==2.10.0
 # These must be updated alongside torch
 torchvision==0.25.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
 # FlashInfer should be updated together with the Dockerfile
-flashinfer-python==0.6.3
+flashinfer-python==0.6.4
@@ -1,4 +1,4 @@
-mkdocs
+mkdocs<2.0.0
 mkdocs-api-autonav
 mkdocs-material
 mkdocstrings-python
@@ -7,6 +7,7 @@ mkdocs-awesome-nav
 mkdocs-glightbox
 mkdocs-git-revision-date-localized-plugin
 mkdocs-minify-plugin
+mkdocs-redirects
 regex
 ruff
 pydantic
@@ -28,12 +28,12 @@ num2words # required for smolvlm test
 opencv-python-headless >= 4.13.0 # required for video test
 datamodel_code_generator # required for minicpm3 test
 lm-eval[api]>=0.4.11 # required for model evaluation test
-mteb>=1.38.11, <2 # required for mteb test
+mteb[bm25s]>=2, <3 # required for mteb test
 transformers==4.57.5
 tokenizers==0.22.0
 schemathesis>=3.39.15 # Required for openai schema test.
 # quantization
-bitsandbytes>=0.46.1
+bitsandbytes>=0.49.2
 buildkite-test-collector==0.1.9


@@ -70,7 +70,7 @@ ray[cgraph,default]>=2.48.0
 torchgeo==0.7.0
    # via terratorch
 # MTEB Benchmark Test
-mteb==2.1.2
+mteb[bm25s]>=2, <3

 # Utilities
 num2words==0.5.14
@@ -102,3 +102,7 @@ terratorch==1.2.2
 segmentation-models-pytorch==0.5.0
 # Required for Prithvi tests
 imagehash==4.3.2
+# Required for bitsandbytes quantization test
+bitsandbytes==0.49.2
+# Examples (tensorizer) tests
+tensorizer==2.10.1
@@ -1,6 +1,11 @@
 # Common dependencies
 -r common.txt

+# The version of gRPC libraries should be consistent with each other
+grpcio==1.78.0
+grpcio-reflection==1.78.0
+grpcio-tools==1.78.0
+
 numba == 0.61.2 # Required for N-gram speculative decoding

 # Dependencies for AMD GPUs
@@ -14,5 +19,4 @@ setuptools>=77.0.3,<80.0.0
 setuptools-scm>=8
 runai-model-streamer[s3,gcs]==0.15.3
 conch-triton-kernels==1.2.1
-timm>=1.0.17
-grpcio-tools==1.78.0 # Should match `build.txt`
+timm>=1.0.17
@@ -41,14 +41,18 @@ transformers==4.57.5
 tokenizers==0.22.0
 schemathesis>=3.39.15 # Required for openai schema test.
 # quantization
-bitsandbytes==0.46.1
+bitsandbytes==0.49.2
 buildkite-test-collector==0.1.9


 genai_perf>=0.0.8
 tritonclient>=2.51.0

-grpcio-tools==1.78.0 # Should match `build.txt`
+# The version of gRPC libraries should be consistent with each other
+grpcio==1.78.0
+grpcio-reflection==1.78.0
+grpcio-tools==1.78.0
+
 arctic-inference == 0.1.1 # Required for suffix decoding test
 numba == 0.61.2 # Required for N-gram speculative decoding
 numpy
@@ -66,7 +66,7 @@ backoff==2.2.1
    # via
    #   -r requirements/test.in
    #   schemathesis
-bitsandbytes==0.46.1
+bitsandbytes==0.49.2
    # via
    #   -r requirements/test.in
    #   lightning
@@ -287,9 +287,13 @@ greenlet==3.2.3
    # via sqlalchemy
 grpcio==1.78.0
    # via
+    #   -r requirements/test.in
+    #   grpcio-reflection
    #   grpcio-tools
    #   ray
    #   tensorboard
+grpcio-reflection==1.78.0
+    # via -r requirements/test.in
 grpcio-tools==1.78.0
    # via -r requirements/test.in
 h11==0.14.0
@@ -487,7 +491,7 @@ msgpack==1.1.0
    # via
    #   librosa
    #   ray
-mteb==2.1.2
+mteb==2.8.3
    # via -r requirements/test.in
 multidict==6.1.0
    # via
@@ -653,6 +657,7 @@ orjson==3.11.5
 packaging==24.2
    # via
    #   accelerate
+    #   bitsandbytes
    #   black
    #   datamodel-code-generator
    #   datasets
@@ -757,6 +762,7 @@ protobuf==6.33.2
    # via
    #   google-api-core
    #   googleapis-common-protos
+    #   grpcio-reflection
    #   grpcio-tools
    #   opentelemetry-proto
    #   proto-plus
@@ -33,6 +33,7 @@ def graph_allreduce(
 ):
    with monkeypatch.context() as m:
        m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
+        m.delenv("HIP_VISIBLE_DEVICES", raising=False)
        device = torch.device(f"cuda:{rank}")
        torch.cuda.set_device(device)
        init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
@@ -92,6 +93,7 @@ def eager_allreduce(
 ):
    with monkeypatch.context() as m:
        m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
+        m.delenv("HIP_VISIBLE_DEVICES", raising=False)
        device = torch.device(f"cuda:{rank}")
        torch.cuda.set_device(device)
        init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
@@ -1,7 +1,22 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from __future__ import annotations
+
+import json
+import logging
+from collections.abc import Callable
+from typing import Any
+
 import pytest

+logger = logging.getLogger(__name__)
+
+BASE_TEST_ENV = {
+    # The day vLLM said "hello world" on arxiv 🚀
+    "VLLM_SYSTEM_START_DATE": "2023-09-12",
+}
+DEFAULT_MAX_RETRIES = 3
+

@pytest.fixture
 def pairs_of_event_types() -> dict[str, str]:
@@ -28,3 +43,159 @@ def pairs_of_event_types() -> dict[str, str]:
    }
    # fmt: on
    return event_pairs
+
+
+async def retry_for_tool_call(
+    client,
+    *,
+    model: str,
+    expected_tool_type: str,
+    max_retries: int = DEFAULT_MAX_RETRIES,
+    **create_kwargs: Any,
+):
+    """Call ``client.responses.create`` up to *max_retries* times, returning
+    the first response that contains an output item of *expected_tool_type*.
+
+    Returns the **last** response if none match so the caller's assertions
+    fire with a clear diagnostic.
+    """
+    last_response = None
+    for attempt in range(max_retries):
+        response = await client.responses.create(model=model, **create_kwargs)
+        last_response = response
+        if any(
+            getattr(item, "type", None) == expected_tool_type
+            for item in response.output
+        ):
+            return response
+    assert last_response is not None
+    return last_response
+
+
+async def retry_streaming_for(
+    client,
+    *,
+    model: str,
+    validate_events: Callable[[list], bool],
+    max_retries: int = DEFAULT_MAX_RETRIES,
+    **create_kwargs: Any,
+) -> list:
+    """Call ``client.responses.create(stream=True)`` up to *max_retries*
+    times, returning the first event list where *validate_events* returns
+    ``True``.
+    """
+    last_events: list = []
+    for attempt in range(max_retries):
+        stream = await client.responses.create(
+            model=model, stream=True, **create_kwargs
+        )
+        events: list = []
+        async for event in stream:
+            events.append(event)
+        last_events = events
+        if validate_events(events):
+            return events
+    return last_events
+
+
+def has_output_type(response, type_name: str) -> bool:
+    """Return True if *response* has at least one output item of *type_name*."""
+    return any(getattr(item, "type", None) == type_name for item in response.output)
+
+
+def events_contain_type(events: list, type_substring: str) -> bool:
+    """Return True if any event's type contains *type_substring*."""
+    return any(type_substring in getattr(e, "type", "") for e in events)
+
+
+def validate_streaming_event_stack(
+    events: list, pairs_of_event_types: dict[str, str]
+) -> None:
+    """Validate that streaming events are properly nested/paired."""
+    stack: list[str] = []
+    for event in events:
+        etype = event.type
+        if etype == "response.created":
+            stack.append(etype)
+        elif etype == "response.completed":
+            assert stack and stack[-1] == pairs_of_event_types[etype], (
+                f"Unexpected stack top for {etype}: "
+                f"got {stack[-1] if stack else '<empty>'}"
+            )
+            stack.pop()
+        elif etype.endswith("added") or etype == "response.mcp_call.in_progress":
+            stack.append(etype)
+        elif etype.endswith("delta"):
+            if stack and stack[-1] == etype:
+                continue
+            stack.append(etype)
+        elif etype.endswith("done") or etype == "response.mcp_call.completed":
+            assert etype in pairs_of_event_types, f"Unknown done event: {etype}"
+            expected_start = pairs_of_event_types[etype]
+            assert stack and stack[-1] == expected_start, (
+                f"Stack mismatch for {etype}: "
+                f"expected {expected_start}, "
+                f"got {stack[-1] if stack else '<empty>'}"
+            )
+            stack.pop()
+    assert len(stack) == 0, f"Unclosed events on stack: {stack}"
+
+
+def log_response_diagnostics(
+    response,
+    *,
+    label: str = "Response Diagnostics",
+) -> dict[str, Any]:
+    """Extract and log diagnostic info from a Responses API response.
+
+    Logs reasoning, tool-call attempts, MCP items, and output types so
+    that CI output (``pytest -s`` or ``--log-cli-level=INFO``) gives
+    full visibility into model behaviour even on passing runs.
+
+    Returns the extracted data so callers can make additional assertions
+    if needed.
+    """
+    reasoning_texts = [
+        text
+        for item in response.output
+        if getattr(item, "type", None) == "reasoning"
+        for content in getattr(item, "content", [])
+        if (text := getattr(content, "text", None))
+    ]
+
+    tool_call_attempts = [
+        {
+            "recipient": msg.get("recipient"),
+            "channel": msg.get("channel"),
+        }
+        for msg in response.output_messages
+        if (msg.get("recipient") or "").startswith("python")
+    ]
+
+    mcp_items = [
+        {
+            "name": getattr(item, "name", None),
+            "status": getattr(item, "status", None),
+        }
+        for item in response.output
+        if getattr(item, "type", None) == "mcp_call"
+    ]
+
+    output_types = [getattr(o, "type", None) for o in response.output]
+
+    diagnostics = {
+        "model_attempted_tool_calls": bool(tool_call_attempts),
+        "tool_call_attempts": tool_call_attempts,
+        "mcp_items": mcp_items,
+        "reasoning": reasoning_texts,
+        "output_text": response.output_text,
+        "output_types": output_types,
+    }
+
+    logger.info(
+        "\n====== %s ======\n%s\n==============================",
+        label,
+        json.dumps(diagnostics, indent=2, default=str),
+    )
+
+    return diagnostics
@@ -1,6 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Integration tests for MCP tool support in the Responses API."""

+from __future__ import annotations

 import pytest
 import pytest_asyncio
@@ -10,11 +12,31 @@ from openai_harmony import ToolDescription, ToolNamespaceConfig
 from vllm.entrypoints.mcp.tool_server import MCPToolServer

 from ....utils import RemoteOpenAIServer
+from .conftest import (
+    BASE_TEST_ENV,
+    events_contain_type,
+    log_response_diagnostics,
+    retry_for_tool_call,
+    retry_streaming_for,
+    validate_streaming_event_stack,
+)

 MODEL_NAME = "openai/gpt-oss-20b"

+_BASE_SERVER_ARGS = [
+    "--enforce-eager",
+    "--tool-server",
+    "demo",
+    "--max_model_len",
+    "5000",
+]

-def test_get_tool_description():
+_PYTHON_TOOL_INSTRUCTION = (
+    "You must use the Python tool to execute code. Never simulate execution."
+)
+
+
+class TestMCPToolServerUnit:
    """Test MCPToolServer.get_tool_description filtering logic.

    Note: The wildcard "*" is normalized to None by
@@ -22,283 +44,240 @@ def test_get_tool_description():
    so we only test None and specific tool filtering here.
    See test_serving_responses.py for "*" normalization tests.
    """
-    pytest.importorskip("mcp")

-    server = MCPToolServer()
-    tool1 = ToolDescription.new(
-        name="tool1", description="First", parameters={"type": "object"}
-    )
-    tool2 = ToolDescription.new(
-        name="tool2", description="Second", parameters={"type": "object"}
-    )
-    tool3 = ToolDescription.new(
-        name="tool3", description="Third", parameters={"type": "object"}
-    )
+    def test_get_tool_description(self):
+        pytest.importorskip("mcp")

-    server.harmony_tool_descriptions = {
-        "test_server": ToolNamespaceConfig(
-            name="test_server", description="test", tools=[tool1, tool2, tool3]
+        server = MCPToolServer()
+        tool1 = ToolDescription.new(
+            name="tool1", description="First", parameters={"type": "object"}
+        )
+        tool2 = ToolDescription.new(
+            name="tool2", description="Second", parameters={"type": "object"}
+        )
+        tool3 = ToolDescription.new(
+            name="tool3", description="Third", parameters={"type": "object"}
        )
-    }

-    # Nonexistent server
-    assert server.get_tool_description("nonexistent") is None
+        server.harmony_tool_descriptions = {
+            "test_server": ToolNamespaceConfig(
+                name="test_server",
+                description="test",
+                tools=[tool1, tool2, tool3],
+            )
+        }

-    # None (no filter) - returns all tools
-    result = server.get_tool_description("test_server", allowed_tools=None)
-    assert len(result.tools) == 3
+        # Nonexistent server
+        assert server.get_tool_description("nonexistent") is None

-    # Filter to specific tools
-    result = server.get_tool_description(
-        "test_server", allowed_tools=["tool1", "tool3"]
-    )
-    assert len(result.tools) == 2
-    assert result.tools[0].name == "tool1"
-    assert result.tools[1].name == "tool3"
+        # None (no filter) - returns all tools
+        result = server.get_tool_description("test_server", allowed_tools=None)
+        assert len(result.tools) == 3

-    # Single tool
-    result = server.get_tool_description(
-        "test_server",
-        allowed_tools=["tool2"],
-    )
-    assert len(result.tools) == 1
-    assert result.tools[0].name == "tool2"
+        # Filter to specific tools
+        result = server.get_tool_description(
+            "test_server", allowed_tools=["tool1", "tool3"]
+        )
+        assert len(result.tools) == 2
+        assert result.tools[0].name == "tool1"
+        assert result.tools[1].name == "tool3"

-    # No matching tools - returns None
-    result = server.get_tool_description("test_server", allowed_tools=["nonexistent"])
-    assert result is None
+        # Single tool
+        result = server.get_tool_description("test_server", allowed_tools=["tool2"])
+        assert len(result.tools) == 1
+        assert result.tools[0].name == "tool2"

-    # Empty list - returns None
-    assert server.get_tool_description("test_server", allowed_tools=[]) is None
+        # No matching tools - returns None
+        result = server.get_tool_description(
+            "test_server", allowed_tools=["nonexistent"]
+        )
+        assert result is None
+
+        # Empty list - returns None
+        assert server.get_tool_description("test_server", allowed_tools=[]) is None
+
+    def test_builtin_tools_consistency(self):
+        """MCP_BUILTIN_TOOLS must match _BUILTIN_TOOL_TO_MCP_SERVER_LABEL values."""
+        from vllm.entrypoints.openai.parser.harmony_utils import (
+            _BUILTIN_TOOL_TO_MCP_SERVER_LABEL,
+            MCP_BUILTIN_TOOLS,
+        )
+
+        assert set(_BUILTIN_TOOL_TO_MCP_SERVER_LABEL.values()) == MCP_BUILTIN_TOOLS, (
+            f"MCP_BUILTIN_TOOLS {MCP_BUILTIN_TOOLS} does not match "
+            f"_BUILTIN_TOOL_TO_MCP_SERVER_LABEL values "
+            f"{set(_BUILTIN_TOOL_TO_MCP_SERVER_LABEL.values())}"
+        )


 class TestMCPEnabled:
    """Tests that require MCP tools to be enabled via environment variable."""

    @pytest.fixture(scope="class")
-    def monkeypatch_class(self):
-        from _pytest.monkeypatch import MonkeyPatch
-
-        mpatch = MonkeyPatch()
-        yield mpatch
-        mpatch.undo()
-
-    @pytest.fixture(scope="class")
-    def mcp_enabled_server(self, monkeypatch_class: pytest.MonkeyPatch):
-        args = ["--enforce-eager", "--tool-server", "demo"]
-
-        with monkeypatch_class.context() as m:
-            m.setenv("VLLM_ENABLE_RESPONSES_API_STORE", "1")
-            m.setenv("PYTHON_EXECUTION_BACKEND", "dangerously_use_uv")
-            m.setenv(
-                "VLLM_GPT_OSS_SYSTEM_TOOL_MCP_LABELS", "code_interpreter,container"
-            )
-            # Helps the model follow instructions better
-            m.setenv("VLLM_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS", "1")
-            with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
-                yield remote_server
+    def mcp_enabled_server(self):
+        env_dict = {
+            **BASE_TEST_ENV,
+            "VLLM_ENABLE_RESPONSES_API_STORE": "1",
+            "PYTHON_EXECUTION_BACKEND": "dangerously_use_uv",
+            "VLLM_GPT_OSS_SYSTEM_TOOL_MCP_LABELS": ("code_interpreter,container"),
+            "VLLM_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS": "1",
+        }
+        with RemoteOpenAIServer(
+            MODEL_NAME, list(_BASE_SERVER_ARGS), env_dict=env_dict
+        ) as remote_server:
+            yield remote_server

    @pytest_asyncio.fixture
-    async def mcp_enabled_client(self, mcp_enabled_server):
+    async def client(self, mcp_enabled_server):
        async with mcp_enabled_server.get_async_client() as async_client:
            yield async_client

+    @staticmethod
+    def _mcp_tools_payload(*, allowed_tools: list[str] | None = None) -> list[dict]:
+        tool: dict = {
+            "type": "mcp",
+            "server_label": "code_interpreter",
+            "server_url": "http://localhost:8888",
+        }
+        if allowed_tools is not None:
+            tool["allowed_tools"] = allowed_tools
+        return [tool]
+
+    @staticmethod
+    def _python_exec_input(code: str = "") -> str:
+        if not code:
+            code = "import random; print(random.randint(1, 1000000))"
+        return f"Execute the following code: {code}"
+
    @pytest.mark.asyncio
    @pytest.mark.parametrize("model_name", [MODEL_NAME])
-    async def test_mcp_tool_env_flag_enabled(
-        self, mcp_enabled_client: OpenAI, model_name: str
-    ):
-        response = await mcp_enabled_client.responses.create(
+    async def test_mcp_tool_env_flag_enabled(self, client: OpenAI, model_name: str):
+        response = await retry_for_tool_call(
+            client,
            model=model_name,
-            input=(
-                "Execute the following code: "
-                "import random; print(random.randint(1, 1000000))"
-            ),
-            instructions=(
-                "You must use the Python tool to execute code. "
-                "Never simulate execution."
-            ),
-            tools=[
-                {
-                    "type": "mcp",
-                    "server_label": "code_interpreter",
-                    # URL unused for DemoToolServer
-                    "server_url": "http://localhost:8888",
-                }
-            ],
+            expected_tool_type="mcp_call",
+            input=self._python_exec_input(),
+            instructions=_PYTHON_TOOL_INSTRUCTION,
+            tools=self._mcp_tools_payload(),
+            temperature=0.0,
            extra_body={"enable_response_messages": True},
        )
-        assert response is not None
+
        assert response.status == "completed"
-        # Verify output messages: Tool calls and responses on analysis channel
+        log_response_diagnostics(response, label="MCP Enabled")
+
        tool_call_found = False
        tool_response_found = False
        for message in response.output_messages:
            recipient = message.get("recipient")
            if recipient and recipient.startswith("python"):
                tool_call_found = True
-                assert message.get("channel") == "analysis", (
-                    "Tool call should be on analysis channel"
-                )
+                assert message.get("channel") == "analysis"
            author = message.get("author", {})
-            if (
-                author.get("role") == "tool"
-                and author.get("name")
-                and author.get("name").startswith("python")
+            if author.get("role") == "tool" and (author.get("name") or "").startswith(
+                "python"
            ):
                tool_response_found = True
-                assert message.get("channel") == "analysis", (
-                    "Tool response should be on analysis channel"
-                )
+                assert message.get("channel") == "analysis"

-        assert tool_call_found, "Should have found at least one Python tool call"
-        assert tool_response_found, (
-            "Should have found at least one Python tool response"
+        assert tool_call_found, (
+            f"No Python tool call found. "
+            f"Output types: "
+            f"{[getattr(o, 'type', None) for o in response.output]}"
        )
+        assert tool_response_found, "No Python tool response found"
+
        for message in response.input_messages:
-            assert message.get("author").get("role") != "developer", (
-                "No developer messages should be present with valid mcp tool"
-            )
+            assert message.get("author", {}).get("role") != "developer"

-    @pytest.mark.flaky(reruns=3)
    @pytest.mark.asyncio
    @pytest.mark.parametrize("model_name", [MODEL_NAME])
    async def test_mcp_tool_with_allowed_tools_star(
-        self, mcp_enabled_client: OpenAI, model_name: str
+        self, client: OpenAI, model_name: str
    ):
-        """Test MCP tool with allowed_tools=['*'] to select all available
-        tools.
-
-        This E2E test verifies that the "*" wildcard works end-to-end.
-        See test_serving_responses.py for detailed unit tests of "*"
-        normalization.
-        """
-        response = await mcp_enabled_client.responses.create(
+        response = await retry_for_tool_call(
+            client,
            model=model_name,
-            input=(
-                "Execute the following code: "
-                "import random; print(random.randint(1, 1000000))"
-            ),
-            instructions=(
-                "You must use the Python tool to execute code. "
-                "Never simulate execution."
-            ),
-            tools=[
-                {
-                    "type": "mcp",
-                    "server_label": "code_interpreter",
-                    "server_url": "http://localhost:8888",
-                    # Using "*" to allow all tools from this MCP server
-                    "allowed_tools": ["*"],
-                }
-            ],
+            expected_tool_type="mcp_call",
+            input=self._python_exec_input(),
+            instructions=_PYTHON_TOOL_INSTRUCTION,
+            tools=self._mcp_tools_payload(allowed_tools=["*"]),
+            temperature=0.0,
            extra_body={"enable_response_messages": True},
        )
-        assert response is not None
+
        assert response.status == "completed"
-        # Verify tool calls work with allowed_tools=["*"]
-        tool_call_found = False
-        for message in response.output_messages:
-            recipient = message.get("recipient")
-            if recipient and recipient.startswith("python"):
-                tool_call_found = True
-                break
+        log_response_diagnostics(response, label="MCP Allowed Tools *")
+
+        tool_call_found = any(
+            (msg.get("recipient") or "").startswith("python")
+            for msg in response.output_messages
+        )
        assert tool_call_found, (
-            "Should have found at least one Python tool call with '*'"
+            f"No Python tool call with '*'. "
+            f"Output types: "
+            f"{[getattr(o, 'type', None) for o in response.output]}"
        )

-    @pytest.mark.flaky(reruns=3)
    @pytest.mark.asyncio
    @pytest.mark.parametrize("model_name", [MODEL_NAME])
    async def test_mcp_tool_calling_streaming_types(
        self,
        pairs_of_event_types: dict[str, str],
-        mcp_enabled_client: OpenAI,
+        client: OpenAI,
        model_name: str,
    ):
-        tools = [
-            {
-                "type": "mcp",
-                "server_label": "code_interpreter",
-            }
-        ]
-        input_text = "What is 123 * 456? Use python to calculate the result."
+        def _has_mcp_events(events: list) -> bool:
+            return events_contain_type(events, "mcp_call")

-        stream_response = await mcp_enabled_client.responses.create(
+        events = await retry_streaming_for(
+            client,
            model=model_name,
-            input=input_text,
-            tools=tools,
-            stream=True,
-            instructions=(
-                "You must use the Python tool to execute code. "
-                "Never simulate execution."
-            ),
+            validate_events=_has_mcp_events,
+            input=("What is 123 * 456? Use Python to calculate the result."),
+            tools=[{"type": "mcp", "server_label": "code_interpreter"}],
+            instructions=_PYTHON_TOOL_INSTRUCTION,
+            temperature=0.0,
        )

-        stack_of_event_types = []
-        saw_mcp_type = False
-        async for event in stream_response:
-            if event.type == "response.created":
-                stack_of_event_types.append(event.type)
-            elif event.type == "response.completed":
-                assert stack_of_event_types[-1] == pairs_of_event_types[event.type]
-                stack_of_event_types.pop()
-            elif (
-                event.type.endswith("added")
-                or event.type == "response.mcp_call.in_progress"
-            ):
-                stack_of_event_types.append(event.type)
-            elif event.type.endswith("delta"):
-                if stack_of_event_types[-1] == event.type:
-                    continue
-                stack_of_event_types.append(event.type)
-            elif (
-                event.type.endswith("done")
-                or event.type == "response.mcp_call.completed"
-            ):
-                assert stack_of_event_types[-1] == pairs_of_event_types[event.type]
-                if "mcp_call" in event.type:
-                    saw_mcp_type = True
-                stack_of_event_types.pop()
+        validate_streaming_event_stack(events, pairs_of_event_types)

-        assert len(stack_of_event_types) == 0
-        assert saw_mcp_type, "Should have seen at least one mcp call"
+        assert events_contain_type(events, "mcp_call"), (
+            f"No mcp_call events after retries. "
+            f"Event types: {sorted({e.type for e in events})}"
+        )


 class TestMCPDisabled:
-    """Tests that verify behavior when MCP tools are disabled."""
+    """Tests that MCP tools are not executed when the env flag is unset."""

    @pytest.fixture(scope="class")
-    def monkeypatch_class(self):
-        from _pytest.monkeypatch import MonkeyPatch
-
-        mpatch = MonkeyPatch()
-        yield mpatch
-        mpatch.undo()
-
-    @pytest.fixture(scope="class")
-    def mcp_disabled_server(self, monkeypatch_class: pytest.MonkeyPatch):
-        args = ["--enforce-eager", "--tool-server", "demo"]
-
-        with monkeypatch_class.context() as m:
-            m.setenv("VLLM_ENABLE_RESPONSES_API_STORE", "1")
-            m.setenv("PYTHON_EXECUTION_BACKEND", "dangerously_use_uv")
-            # Helps the model follow instructions better
-            m.setenv("VLLM_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS", "1")
-            with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
-                yield remote_server
+    def mcp_disabled_server(self):
+        env_dict = {
+            **BASE_TEST_ENV,
+            "VLLM_ENABLE_RESPONSES_API_STORE": "1",
+            "PYTHON_EXECUTION_BACKEND": "dangerously_use_uv",
+            "VLLM_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS": "1",
+        }
+        with RemoteOpenAIServer(
+            MODEL_NAME, list(_BASE_SERVER_ARGS), env_dict=env_dict
+        ) as remote_server:
+            yield remote_server

    @pytest_asyncio.fixture
-    async def mcp_disabled_client(self, mcp_disabled_server):
+    async def client(self, mcp_disabled_server):
        async with mcp_disabled_server.get_async_client() as async_client:
            yield async_client

    @pytest.mark.asyncio
    @pytest.mark.parametrize("model_name", [MODEL_NAME])
-    async def test_mcp_tool_env_flag_disabled(
-        self, mcp_disabled_client: OpenAI, model_name: str
+    async def test_mcp_disabled_server_does_not_execute(
+        self, client: OpenAI, model_name: str
    ):
-        response = await mcp_disabled_client.responses.create(
+        """When MCP is disabled the model may still attempt tool calls
+        (tool descriptions can remain in the prompt), but the server
+        must NOT execute them."""
+        response = await client.responses.create(
            model=model_name,
            input=(
                "Execute the following code if the tool is present: "
@@ -308,38 +287,35 @@ class TestMCPDisabled:
                {
                    "type": "mcp",
                    "server_label": "code_interpreter",
-                    # URL unused for DemoToolServer
                    "server_url": "http://localhost:8888",
                }
            ],
+            temperature=0.0,
            extra_body={"enable_response_messages": True},
        )
        assert response is not None
        assert response.status == "completed"
-        # Verify output messages: No tool calls and responses
-        tool_call_found = False
-        tool_response_found = False
+
+        log_response_diagnostics(response, label="MCP Disabled")
+
+        # Server must not have executed any tool calls
        for message in response.output_messages:
-            recipient = message.get("recipient")
-            if recipient and recipient.startswith("python"):
-                tool_call_found = True
-                assert message.get("channel") == "analysis", (
-                    "Tool call should be on analysis channel"
-                )
            author = message.get("author", {})
-            if (
+            assert not (
                author.get("role") == "tool"
-                and author.get("name")
-                and author.get("name").startswith("python")
-            ):
-                tool_response_found = True
-                assert message.get("channel") == "analysis", (
-                    "Tool response should be on analysis channel"
+                and (author.get("name") or "").startswith("python")
+            ), (
+                "Server executed a python tool call even though MCP is "
+                f"disabled. Message: {message}"
+            )
+
+        # No completed mcp_call output items
+        for item in response.output:
+            if getattr(item, "type", None) == "mcp_call":
+                assert getattr(item, "status", None) != "completed", (
+                    "MCP call should not be completed when MCP is disabled"
                )

-        assert not tool_call_found, "Should not have a python call"
-        assert not tool_response_found, "Should not have a tool response"
+        # No developer messages injected
        for message in response.input_messages:
-            assert message.get("author").get("role") != "developer", (
-                "No developer messages should be present without a valid tool"
-            )
+            assert message.get("author", {}).get("role") != "developer"
@@ -3,15 +3,29 @@

 import importlib.util
 import json
+import logging

 import pytest
 import pytest_asyncio
 from openai import OpenAI

 from ....utils import RemoteOpenAIServer
+from .conftest import (
+    BASE_TEST_ENV,
+    has_output_type,
+    log_response_diagnostics,
+    retry_for_tool_call,
+)
+
+logger = logging.getLogger(__name__)

 MODEL_NAME = "Qwen/Qwen3-8B"

+_PYTHON_TOOL_INSTRUCTION = (
+    "You must use the Python tool to execute code. "
+    "Never simulate execution. You must print the final answer."
+)
+

@pytest.fixture(scope="module")
 def server():
@@ -32,12 +46,12 @@ def server():
        "--tool-server",
        "demo",
    ]
-    env_dict = dict(
-        VLLM_ENABLE_RESPONSES_API_STORE="1",
-        VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT="1",
-        PYTHON_EXECUTION_BACKEND="dangerously_use_uv",
-    )
-
+    env_dict = {
+        **BASE_TEST_ENV,
+        "VLLM_ENABLE_RESPONSES_API_STORE": "1",
+        "VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT": "1",
+        "PYTHON_EXECUTION_BACKEND": "dangerously_use_uv",
+    }
    with RemoteOpenAIServer(MODEL_NAME, args, env_dict=env_dict) as remote_server:
        yield remote_server

@@ -54,6 +68,7 @@ async def test_basic(client: OpenAI, model_name: str):
    response = await client.responses.create(
        model=model_name,
        input="What is 123 * 456?",
+        temperature=0.0,
    )
    assert response is not None
    print("response: ", response)
@@ -99,10 +114,15 @@ async def test_reasoning_and_function_items(client: OpenAI, model_name: str):
    )
    assert response is not None
    assert response.status == "completed"
-    # make sure we get a reasoning and text output
-    assert response.output[0].type == "reasoning"
-    assert response.output[1].type == "message"
-    assert type(response.output[1].content[0].text) is str
+
+    output_types = [getattr(o, "type", None) for o in response.output]
+    assert "reasoning" in output_types, (
+        f"Expected reasoning in output, got: {output_types}"
+    )
+    assert "message" in output_types, f"Expected message in output, got: {output_types}"
+
+    msg = next(o for o in response.output if o.type == "message")
+    assert type(msg.content[0].text) is str


 def get_horoscope(sign):
@@ -110,10 +130,10 @@ def get_horoscope(sign):


 def call_function(name, args):
+    logger.info("Calling function %s with args %s", name, args)
    if name == "get_horoscope":
        return get_horoscope(**args)
-    else:
-        raise ValueError(f"Unknown function: {name}")
+    raise ValueError(f"Unknown function: {name}")


@pytest.mark.asyncio
@@ -136,61 +156,111 @@ async def test_function_call_first_turn(client: OpenAI, model_name: str):
        }
    ]

-    response = await client.responses.create(
+    response = await retry_for_tool_call(
+        client,
        model=model_name,
+        expected_tool_type="function_call",
        input="What is the horoscope for Aquarius today?",
        tools=tools,
        temperature=0.0,
    )
    assert response is not None
    assert response.status == "completed"
-    assert len(response.output) == 2
-    assert response.output[0].type == "reasoning"
-    assert response.output[1].type == "function_call"

-    function_call = response.output[1]
+    output_types = [getattr(o, "type", None) for o in response.output]
+    assert "reasoning" in output_types, (
+        f"Expected reasoning in output, got: {output_types}"
+    )
+    assert has_output_type(response, "function_call"), (
+        f"Expected function_call in output, got: {output_types}"
+    )
+
+    function_call = next(o for o in response.output if o.type == "function_call")
    assert function_call.name == "get_horoscope"
    assert function_call.call_id is not None

    args = json.loads(function_call.arguments)
    assert "sign" in args

-    # the multi turn function call is tested above in
-    # test_reasoning_and_function_items
-

@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
 async def test_mcp_tool_call(client: OpenAI, model_name: str):
-    response = await client.responses.create(
+    """MCP tool calling with code_interpreter.
+
+    The model may make one or more tool calls before producing a final
+    message.  We validate server invariants (mcp_call items have correct
+    fields) with hard assertions.  Output indices are never hardcoded
+    since the model can produce multiple tool-call rounds.
+    """
+    # MCP + container init + code execution can be slow
+    client_with_timeout = client.with_options(timeout=client.timeout * 3)
+
+    response = await retry_for_tool_call(
+        client_with_timeout,
        model=model_name,
-        input="What is 123 * 456? Use python to calculate the result.",
+        expected_tool_type="mcp_call",
+        input=(
+            "What is 123 * 456? Use python to calculate the result. "
+            "Print the result with print()."
+        ),
        tools=[{"type": "code_interpreter", "container": {"type": "auto"}}],
-        extra_body={"enable_response_messages": True},
+        instructions=_PYTHON_TOOL_INSTRUCTION,
        temperature=0.0,
+        extra_body={"enable_response_messages": True},
    )

    assert response is not None
-    assert response.status == "completed"

-    # The model may produce multiple reasoning/mcp_call rounds before the
-    # final message, so validate structurally rather than by exact index.
-    output_types = [o.type for o in response.output]
-    assert "reasoning" in output_types
-    mcp_calls = [o for o in response.output if o.type == "mcp_call"]
-    assert len(mcp_calls) >= 1
-    assert type(mcp_calls[0].arguments) is str
-    assert type(mcp_calls[0].output) is str
+    output_types = [getattr(o, "type", None) for o in response.output]
+    log_response_diagnostics(response, label="test_mcp_tool_call")

-    # The final output should be a message containing the correct answer
-    assert response.output[-1].type == "message"
-    assert any(s in response.output[-1].content[0].text for s in ("56088", "56,088"))
+    assert response.status == "completed", (
+        f"Response status={response.status} "
+        f"(details={getattr(response, 'incomplete_details', None)}). "
+        f"Output types: {output_types}."
+    )

-    # Test raw input_messages / output_messages
-    assert len(response.input_messages) == 1
-    assert len(response.output_messages) >= 3
+    assert "reasoning" in output_types, (
+        f"Expected reasoning in output, got: {output_types}"
+    )
+    assert "mcp_call" in output_types, (
+        f"Expected mcp_call in output, got: {output_types}"
+    )
+
+    # Every mcp_call item must have well-typed fields
+    for item in response.output:
+        if getattr(item, "type", None) == "mcp_call":
+            assert type(item.arguments) is str, (
+                f"mcp_call.arguments should be str, got {type(item.arguments)}"
+            )
+            assert type(item.output) is str, (
+                f"mcp_call.output should be str, got {type(item.output)}"
+            )
+
+    # The model may make 1+ tool-call rounds but must still produce
+    # a final message for a trivial calculation like 123 * 456.
+    message_outputs = [
+        o for o in response.output if getattr(o, "type", None) == "message"
+    ]
+    assert message_outputs, (
+        f"Model did not produce a final message. Output types: {output_types}"
+    )
+
+    final_message = message_outputs[-1]
+    assert any(s in final_message.content[0].text for s in ("56088", "56,088")), (
+        f"Expected 56088 in final message, got: {final_message.content[0].text!r}"
+    )
+
+    # Validate raw input_messages / output_messages
+    assert len(response.input_messages) >= 1, "Expected at least 1 input message"
+    assert len(response.output_messages) >= 1, "Expected at least 1 output message"
    assert any(
-        s in response.output_messages[-1]["message"] for s in ("56088", "56,088")
+        any(s in str(msg) for s in ("56088", "56,088"))
+        for msg in response.output_messages
+    ), (
+        f"Expected 56088 in at least one output_message, "
+        f"got {len(response.output_messages)} messages"
    )


@@ -202,6 +272,7 @@ async def test_max_tokens(client: OpenAI, model_name: str):
        input="What is the first paragraph of Moby Dick?",
        reasoning={"effort": "low"},
        max_output_tokens=30,
+        temperature=0.0,
    )
    assert response is not None
    assert response.status == "incomplete"
@@ -12,13 +12,15 @@ MODEL_NAME = "Qwen/Qwen3-8B"

@pytest.fixture(scope="module")
 def server():
-    args = ["--reasoning-parser", "qwen3", "--max_model_len", "5000"]
-    env_dict = dict(
-        VLLM_ENABLE_RESPONSES_API_STORE="1",
-        # uncomment for tool calling
-        # PYTHON_EXECUTION_BACKEND="dangerously_use_uv",
-    )
+    from .conftest import BASE_TEST_ENV

+    args = ["--reasoning-parser", "qwen3", "--max_model_len", "5000"]
+    env_dict = {
+        **BASE_TEST_ENV,
+        "VLLM_ENABLE_RESPONSES_API_STORE": "1",
+        # uncomment for tool calling
+        # PYTHON_EXECUTION_BACKEND: "dangerously_use_uv",
+    }
    with RemoteOpenAIServer(MODEL_NAME, args, env_dict=env_dict) as remote_server:
        yield remote_server

@@ -134,6 +136,53 @@ async def test_streaming_output_consistency(client: OpenAI, model_name: str):
    )


+@pytest.mark.asyncio
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
+async def test_streaming_reasoning_tokens_e2e(client: OpenAI, model_name: str):
+    """Verify final usage includes reasoning_tokens in streaming mode."""
+    response = await client.responses.create(
+        model=model_name,
+        input="Compute 17 * 19 and explain briefly.",
+        reasoning={"effort": "low"},
+        temperature=0.0,
+        stream=True,
+    )
+
+    completed_event = None
+    async for event in response:
+        if event.type == "response.completed":
+            completed_event = event
+
+    assert completed_event is not None
+    assert completed_event.response.status == "completed"
+    assert completed_event.response.usage is not None
+    assert completed_event.response.usage.output_tokens_details is not None
+    assert completed_event.response.usage.output_tokens_details.reasoning_tokens > 0, (
+        "Expected reasoning_tokens > 0 for streamed Qwen3 response."
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
+async def test_non_streaming_reasoning_tokens_e2e(client: OpenAI, model_name: str):
+    """Verify usage includes reasoning_tokens in non-streaming mode."""
+    response = await client.responses.create(
+        model=model_name,
+        input="Compute 23 * 17 and explain briefly.",
+        reasoning={"effort": "low"},
+        temperature=0.0,
+        stream=False,
+    )
+
+    assert response is not None
+    assert response.status == "completed"
+    assert response.usage is not None
+    assert response.usage.output_tokens_details is not None
+    assert response.usage.output_tokens_details.reasoning_tokens > 0, (
+        "Expected reasoning_tokens > 0 for non-streamed Qwen3 response."
+    )
+
+
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
 async def test_max_tokens(client: OpenAI, model_name: str):
@@ -4,7 +4,7 @@
 from dataclasses import dataclass, field
 from http import HTTPStatus
 from typing import Any
-from unittest.mock import AsyncMock, MagicMock
+from unittest.mock import AsyncMock, MagicMock, patch

 import pytest

@@ -233,3 +233,140 @@ async def test_chat_error_stream():
        f"Expected error message in chunks: {chunks}"
    )
    assert chunks[-1] == "data: [DONE]\n\n"
+
+
+@pytest.mark.parametrize(
+    "image_content",
+    [
+        [{"type": "image_url", "image_url": {"url": "https://example.com/image.jpg"}}],
+        [{"image_url": {"url": "https://example.com/image.jpg"}}],
+    ],
+)
+def test_system_message_warns_on_image(image_content):
+    """Test that system messages with image content trigger a warning."""
+    with patch(
+        "vllm.entrypoints.openai.chat_completion.protocol.logger"
+    ) as mock_logger:
+        ChatCompletionRequest(
+            model=MODEL_NAME,
+            messages=[
+                {
+                    "role": "system",
+                    "content": image_content,
+                }
+            ],
+        )
+
+    mock_logger.warning_once.assert_called()
+    call_args = str(mock_logger.warning_once.call_args)
+    assert "System messages should only contain text" in call_args
+    assert "image_url" in call_args
+
+
+def test_system_message_accepts_text():
+    """Test that system messages can contain text content."""
+    # Should not raise an exception
+    request = ChatCompletionRequest(
+        model=MODEL_NAME,
+        messages=[
+            {"role": "system", "content": "You are a helpful assistant."},
+        ],
+    )
+    assert request.messages[0]["role"] == "system"
+
+
+def test_system_message_accepts_text_array():
+    """Test that system messages can contain an array with text content."""
+    # Should not raise an exception
+    request = ChatCompletionRequest(
+        model=MODEL_NAME,
+        messages=[
+            {
+                "role": "system",
+                "content": [{"type": "text", "text": "You are a helpful assistant."}],
+            },
+        ],
+    )
+    assert request.messages[0]["role"] == "system"
+
+
+def test_user_message_accepts_image():
+    """Test that user messages can still contain image content."""
+    # Should not raise an exception
+    request = ChatCompletionRequest(
+        model=MODEL_NAME,
+        messages=[
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": "What's in this image?"},
+                    {
+                        "type": "image_url",
+                        "image_url": {"url": "https://example.com/image.jpg"},
+                    },
+                ],
+            },
+        ],
+    )
+    assert request.messages[0]["role"] == "user"
+
+
+@pytest.mark.parametrize(
+    "audio_content",
+    [
+        [
+            {
+                "type": "input_audio",
+                "input_audio": {"data": "base64data", "format": "wav"},
+            }
+        ],
+        [{"input_audio": {"data": "base64data", "format": "wav"}}],
+    ],
+)
+def test_system_message_warns_on_audio(audio_content):
+    """Test that system messages with audio content trigger a warning."""
+    with patch(
+        "vllm.entrypoints.openai.chat_completion.protocol.logger"
+    ) as mock_logger:
+        ChatCompletionRequest(
+            model=MODEL_NAME,
+            messages=[
+                {
+                    "role": "system",
+                    "content": audio_content,
+                }
+            ],
+        )
+
+    mock_logger.warning_once.assert_called()
+    call_args = str(mock_logger.warning_once.call_args)
+    assert "System messages should only contain text" in call_args
+    assert "input_audio" in call_args
+
+
+@pytest.mark.parametrize(
+    "video_content",
+    [
+        [{"type": "video_url", "video_url": {"url": "https://example.com/video.mp4"}}],
+        [{"video_url": {"url": "https://example.com/video.mp4"}}],
+    ],
+)
+def test_system_message_warns_on_video(video_content):
+    """Test that system messages with video content trigger a warning."""
+    with patch(
+        "vllm.entrypoints.openai.chat_completion.protocol.logger"
+    ) as mock_logger:
+        ChatCompletionRequest(
+            model=MODEL_NAME,
+            messages=[
+                {
+                    "role": "system",
+                    "content": video_content,
+                }
+            ],
+        )
+
+    mock_logger.warning_once.assert_called()
+    call_args = str(mock_logger.warning_once.call_args)
+    assert "System messages should only contain text" in call_args
+    assert "video_url" in call_args
@@ -20,10 +20,22 @@ CHATML_JINJA_PATH = VLLM_PATH / "examples/template_chatml.jinja"
 assert CHATML_JINJA_PATH.exists()


+def _build_vllm_parsers():
+    vllm_parser = FlexibleArgumentParser()
+    subparsers = vllm_parser.add_subparsers()
+    serve_parser = subparsers.add_parser("serve")
+    make_arg_parser(serve_parser)
+    return {"vllm": vllm_parser, "vllm serve": serve_parser}
+
+
+@pytest.fixture
+def vllm_parser():
+    return _build_vllm_parsers()["vllm"]
+
+
@pytest.fixture
 def serve_parser():
-    parser = FlexibleArgumentParser(description="vLLM's remote OpenAI server.")
-    return make_arg_parser(parser)
+    return _build_vllm_parsers()["vllm serve"]


 ### Test config parsing
@@ -241,3 +253,41 @@ def test_default_chat_template_kwargs_invalid_json(serve_parser):
        serve_parser.parse_args(
            args=["--default-chat-template-kwargs", "not valid json"]
        )
+
+
+@pytest.mark.parametrize(
+    "args, raises",
+    [
+        (["user/model"], None),
+        (["user/model", "--served-model-name", "model"], None),
+        (["--served-model-name", "model", "user/model"], ValueError),
+        (["--served-model-name", "model", "--config", "config.yaml"], None),
+        (["--served-model-name", "model", "--config", "config.yaml"], ValueError),
+    ],
+    ids=[
+        "model_tag_only",
+        "model_tag_with_served_model_name",
+        "served_model_name_before_model_tag",
+        "served_model_name_with_model_in_config",
+        "served_model_name_with_no_model_in_config",
+    ],
+)
+def test_served_model_name_parsing(tmp_path, vllm_parser, args, raises):
+    """Ensure that users don't misuse --served-model-name and end up with the default
+    model tag instead of the one they intended to serve."""
+    # Call the serve subparser
+    args.insert(0, "serve")
+    # Create a dummy config file if the test case includes it
+    if "config.yaml" in args:
+        # Create a dummy config file if the test case includes it
+        config_path = tmp_path / "config.yaml"
+        config_path.write_text("model: user/model" if raises is None else "port: 8000")
+        args[args.index("config.yaml")] = config_path.as_posix()
+    # Do the parsing and check for expected exceptions or values
+    if raises is None:
+        parsed_args = vllm_parser.parse_args(args=args)
+        expected = "user/model"
+        assert parsed_args.model_tag == expected or parsed_args.model == expected
+    else:
+        with pytest.raises(raises):
+            vllm_parser.parse_args(args=args)
@@ -4,6 +4,7 @@
 import asyncio
 import base64
 import json
+import warnings

 import librosa
 import numpy as np
@@ -85,7 +86,41 @@ async def test_multi_chunk_streaming(

            await send_event(ws, {"type": "session.update", "model": model_name})

-            # Send commit to start transcription
+            # Wait for the server to acknowledge the session update.
+            try:
+                while True:
+                    event = await receive_event(ws, timeout=5.0)
+                    if event["type"] == "session.updated":
+                        break
+            except TimeoutError:
+                warnings.warn(
+                    f"session.updated not received within {5.0}s after "
+                    "session.update. The server may not implement this event.",
+                    stacklevel=2,
+                )
+
+            # (ROCm) Warm-up: send a non-final commit (required to start
+            # transcription) with a small audio chunk to trigger aiter
+            # compilation on first use.
+            await send_event(ws, {"type": "input_audio_buffer.commit"})
+            await send_event(
+                ws,
+                {
+                    "type": "input_audio_buffer.append",
+                    "audio": mary_had_lamb_audio_chunks[0],
+                },
+            )
+            await send_event(ws, {"type": "input_audio_buffer.commit", "final": True})
+
+            # (ROCm) Drain all warm-up responses with generous timeout for
+            # JIT compilation
+            warmup_done = False
+            while not warmup_done:
+                event = await receive_event(ws, timeout=360.0)
+                if event["type"] in ("transcription.done", "error"):
+                    warmup_done = True
+
+            # Now send the real test audio
            await send_event(ws, {"type": "input_audio_buffer.commit"})

            # Send multiple audio chunks
@@ -121,3 +156,101 @@ async def test_multi_chunk_streaming(
                " it sleeps with quite a flow, and everywhere that Mary went,"
                " the lamb was sure to go."
            )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
+async def test_empty_commit_does_not_crash_engine(
+    model_name, mary_had_lamb_audio_chunks, rocm_aiter_fa_attention
+):
+    """Test that committing without audio does not crash the engine.
+
+    Regression test for https://github.com/vllm-project/vllm/issues/34532.
+    An empty commit (no prior input_audio_buffer.append) used to trigger
+    ``AssertionError: For realtime you must provide a multimodal_embedding
+    at every step`` which killed the entire engine process, disconnecting
+    every connected client.
+    """
+    server_args = ["--enforce-eager", "--max-model-len", "2048"]
+
+    if model_name.startswith("mistralai"):
+        server_args += MISTRAL_FORMAT_ARGS
+
+    add_attention_backend(server_args, rocm_aiter_fa_attention)
+
+    with RemoteOpenAIServer(model_name, server_args) as remote_server:
+        ws_url = _get_websocket_url(remote_server)
+
+        # --- First connection: empty commit (no audio appended) ----------
+        async with websockets.connect(ws_url) as ws:
+            event = await receive_event(ws, timeout=30.0)
+            assert event["type"] == "session.created"
+
+            await send_event(ws, {"type": "session.update", "model": model_name})
+
+            try:
+                while True:
+                    event = await receive_event(ws, timeout=5.0)
+                    if event["type"] == "session.updated":
+                        break
+            except TimeoutError:
+                warnings.warn(
+                    f"session.updated not received within {5.0}s after "
+                    "session.update. The server may not implement this event.",
+                    stacklevel=2,
+                )
+
+            # Start generation without sending any audio
+            await send_event(ws, {"type": "input_audio_buffer.commit"})
+
+            # Immediately signal end-of-audio
+            await send_event(ws, {"type": "input_audio_buffer.commit", "final": True})
+
+            # We should get *some* response (error or empty transcription),
+            # but the engine must NOT crash.
+            # (ROCm) Use generous timeout for first request (aiter JIT compilation)
+            event = await receive_event(ws, timeout=360.0)
+            assert event["type"] in (
+                "error",
+                "transcription.done",
+                "transcription.delta",
+            )
+
+        # --- Second connection: normal transcription ---------------------
+        # Verifies the engine is still alive after the empty commit above.
+        async with websockets.connect(ws_url) as ws:
+            event = await receive_event(ws, timeout=30.0)
+            assert event["type"] == "session.created"
+
+            await send_event(ws, {"type": "session.update", "model": model_name})
+
+            try:
+                while True:
+                    event = await receive_event(ws, timeout=5.0)
+                    if event["type"] == "session.updated":
+                        break
+            except TimeoutError:
+                warnings.warn(
+                    f"session.updated not received within {5.0}s after "
+                    "session.update. The server may not implement this event.",
+                    stacklevel=2,
+                )
+
+            # Start transcription
+            await send_event(ws, {"type": "input_audio_buffer.commit"})
+
+            for chunk in mary_had_lamb_audio_chunks:
+                await send_event(
+                    ws, {"type": "input_audio_buffer.append", "audio": chunk}
+                )
+
+            await send_event(ws, {"type": "input_audio_buffer.commit", "final": True})
+
+            done_received = False
+            while not done_received:
+                event = await receive_event(ws, timeout=60.0)
+                if event["type"] == "transcription.done":
+                    done_received = True
+                elif event["type"] == "error":
+                    pytest.fail(f"Engine error after empty commit: {event}")
+            assert done_received
@@ -126,7 +126,7 @@ def gptoss_speculative_server(default_server_args: list[str]):
    if is_aiter_found_and_supported():
        env_dict = {"VLLM_ROCM_USE_AITER": "1"}
    with RemoteOpenAIServer(
-        GPT_OSS_MODEL_NAME, server_args, env_dict=env_dict
+        GPT_OSS_MODEL_NAME, server_args, env_dict=env_dict, max_wait_seconds=480
    ) as remote_server:
        yield remote_server

@@ -13,9 +13,13 @@ from openai.types.responses.tool import (
    Tool,
 )

+import vllm.envs as envs
 from vllm.entrypoints.mcp.tool_server import ToolServer
-from vllm.entrypoints.openai.engine.protocol import ErrorResponse
-from vllm.entrypoints.openai.responses.context import ConversationContext
+from vllm.entrypoints.openai.engine.protocol import (
+    ErrorResponse,
+    RequestResponseMetadata,
+)
+from vllm.entrypoints.openai.responses.context import ConversationContext, SimpleContext
 from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
 from vllm.entrypoints.openai.responses.serving import (
    OpenAIServingResponses,
@@ -23,6 +27,8 @@ from vllm.entrypoints.openai.responses.serving import (
    extract_tool_types,
 )
 from vllm.inputs.data import TokensPrompt
+from vllm.outputs import CompletionOutput, RequestOutput
+from vllm.sampling_params import SamplingParams


 class MockConversationContext(ConversationContext):
@@ -259,6 +265,87 @@ class TestValidateGeneratorInput:
        assert isinstance(result, ErrorResponse)


+@pytest.mark.asyncio
+async def test_reasoning_tokens_counted_for_text_reasoning_model(monkeypatch):
+    """Ensure reasoning_tokens usage is derived from thinking token spans."""
+
+    class FakeTokenizer:
+        def __init__(self):
+            self._vocab = {"<think>": 1, "</think>": 2, "reason": 3, "final": 4}
+
+        def get_vocab(self):
+            return self._vocab
+
+    # Force non-harmony, SimpleContext path
+    monkeypatch.setattr(envs, "VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT", False)
+
+    engine_client = MagicMock()
+    model_config = MagicMock()
+    model_config.hf_config.model_type = "test"
+    model_config.hf_text_config = MagicMock()
+    model_config.get_diff_sampling_param.return_value = {}
+    engine_client.model_config = model_config
+    engine_client.input_processor = MagicMock()
+    engine_client.io_processor = MagicMock()
+    engine_client.renderer = MagicMock()
+
+    tokenizer = FakeTokenizer()
+    engine_client.renderer.get_tokenizer.return_value = tokenizer
+
+    models = MagicMock()
+
+    serving = OpenAIServingResponses(
+        engine_client=engine_client,
+        models=models,
+        request_logger=None,
+        chat_template=None,
+        chat_template_content_format="auto",
+        reasoning_parser="qwen3",
+    )
+
+    # Build a SimpleContext with thinking tokens in the output.
+    context = SimpleContext()
+    token_ids = [1, 10, 2, 20]  # <think> 10 </think> 20 -> reasoning token count = 1
+    completion = CompletionOutput(
+        index=0,
+        text="<think>reason</think>final",
+        token_ids=token_ids,
+        cumulative_logprob=0.0,
+        logprobs=None,
+        finish_reason="stop",
+        stop_reason=None,
+    )
+    req_output = RequestOutput(
+        request_id="req",
+        prompt="hi",
+        prompt_token_ids=[7, 8],
+        prompt_logprobs=None,
+        outputs=[completion],
+        finished=True,
+        num_cached_tokens=0,
+    )
+    context.append_output(req_output)
+
+    async def dummy_result_generator():
+        yield None
+
+    request = ResponsesRequest(input="hi", tools=[], stream=False)
+    sampling_params = SamplingParams(max_tokens=16)
+    metadata = RequestResponseMetadata(request_id="req")
+
+    response = await serving.responses_full_generator(
+        request=request,
+        sampling_params=sampling_params,
+        result_generator=dummy_result_generator(),
+        context=context,
+        model_name="test-model",
+        tokenizer=tokenizer,
+        request_metadata=metadata,
+    )
+
+    assert response.usage.output_tokens_details.reasoning_tokens == 1
+
+
 class TestExtractAllowedToolsFromMcpRequests:
    """Test class for _extract_allowed_tools_from_mcp_requests function"""

@@ -273,3 +273,30 @@ async def test_audio_with_max_tokens(whisper_client, mary_had_lamb):
    out_text = out["text"]
    out_tokens = tok(out_text, add_special_tokens=False)["input_ids"]
    assert len(out_tokens) < 450  # ~Whisper max output len
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("fixture_name", "expected_lang", "expected_text"),
+    [
+        ("mary_had_lamb", "en", ["Mary had a little lamb"]),
+        ("foscolo", "it", ["zacinto", "sacre"]),
+    ],
+    ids=["english", "italian"],
+)
+async def test_language_auto_detect(
+    whisper_client, fixture_name, expected_lang, expected_text, request
+):
+    """Auto-detect language when no language param is provided."""
+    audio_file = request.getfixturevalue(fixture_name)
+    transcription = await whisper_client.audio.transcriptions.create(
+        model=MODEL_NAME,
+        file=audio_file,
+        response_format="verbose_json",
+        temperature=0.0,
+    )
+    assert transcription.language == expected_lang
+    text_lower = transcription.text.lower()
+    assert any(word.lower() in text_lower for word in expected_text), (
+        f"Expected {expected_lang} text but got: {transcription.text}"
+    )
@@ -58,13 +58,19 @@ if current_platform.is_rocm():
    torch.backends.cuda.enable_mem_efficient_sdp(False)
    torch.backends.cuda.enable_math_sdp(True)

+# On ROCm, floating-point reductions in attention and GEMM kernels are
+# non-associative and sensitive to batch geometry. Force LLM instances
+# into an identical, deterministic execution mode:
+ROCM_DETERMINISM_ARGS: list[str] = (
+    ["--max-num-seqs", "1"] if current_platform.is_rocm() else []
+)
+

@pytest.fixture(scope="module")
 def server():
    args = [
        "--runner",
        "pooling",
-        # use half precision for speed and memory savings in CI environment
        "--dtype",
        DTYPE,
        "--enforce-eager",
@@ -72,12 +78,9 @@ def server():
        "512",
        "--chat-template",
        DUMMY_CHAT_TEMPLATE,
+        *ROCM_DETERMINISM_ARGS,
    ]

-    # ROCm: Use Flex Attention to support encoder-only self-attention.
-    if current_platform.is_rocm():
-        args.extend(["--attention-backend", "FLEX_ATTENTION"])
-
    with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
        yield remote_server

@@ -343,8 +346,15 @@ async def test_chat_request(
    assert chat_embeddings.id is not None
    assert completion_embeddings.id is not None
    assert chat_embeddings.created <= completion_embeddings.created
-    assert chat_embeddings.model_dump(exclude={"id", "created"}) == (
-        completion_embeddings.model_dump(exclude={"id", "created"})
+    # Use tolerance-based comparison for embeddings
+    check_embeddings_close(
+        embeddings_0_lst=[d.embedding for d in chat_embeddings.data],
+        embeddings_1_lst=[d.embedding for d in completion_embeddings.data],
+        name_0="chat",
+        name_1="completion",
+    )
+    assert chat_embeddings.model_dump(exclude={"id", "created", "data"}) == (
+        completion_embeddings.model_dump(exclude={"id", "created", "data"})
    )

    # test add_generation_prompt
@@ -124,6 +124,8 @@ def test_init_weight_transfer_engine_calls_engine():
    if torch.cuda.device_count() < 1:
        pytest.skip("Need at least 1 GPU for this test")

+    # Run in-process so mock.patch works (spawn won't inherit the mock)
+    os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
    # Enable insecure serialization to allow pickling functions for collective_rpc
    os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"

@@ -171,6 +173,8 @@ def test_update_weights_calls_engine():
    if torch.cuda.device_count() < 1:
        pytest.skip("Need at least 1 GPU for this test")

+    # Run in-process so mock.patch works (spawn won't inherit the mock)
+    os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
    # Enable insecure serialization to allow pickling functions for collective_rpc
    os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"

@@ -228,6 +232,8 @@ def test_full_weight_transfer_flow():
    if torch.cuda.device_count() < 1:
        pytest.skip("Need at least 1 GPU for this test")

+    # Run in-process so mock.patch works (spawn won't inherit the mock)
+    os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
    # Enable insecure serialization to allow pickling functions for collective_rpc
    os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"

@@ -54,8 +54,8 @@ def reset_config_manager_singleton():
 class TestSiluMulFp8ConfigPicker:
    def test_config_picker_exact_match(self):
        config_keys = [
-            "intermediate_2048_batchsize_256",
-            "intermediate_4096_batchsize_256",
+            "intermediate_2048_numtokens_256",
+            "intermediate_4096_numtokens_256",
        ]

        input_tensor = torch.randn(32, 4096, dtype=torch.bfloat16, device="cuda")
@@ -63,12 +63,12 @@ class TestSiluMulFp8ConfigPicker:
        args = (input_tensor, scale)

        selected_key = pick_silu_mul_fp8_config(args, config_keys)
-        assert selected_key == "intermediate_2048_batchsize_256"
+        assert selected_key == "intermediate_2048_numtokens_256"

    def test_config_picker_closest_match(self):
        config_keys = [
-            "intermediate_2048_batchsize_256",
-            "intermediate_4096_batchsize_256",
+            "intermediate_2048_numtokens_256",
+            "intermediate_4096_numtokens_256",
        ]
        # Use 7000 (intermediate_size=3500) which is closer to 4096 than 2048
        input_tensor = torch.randn(32, 7000, dtype=torch.bfloat16, device="cuda")
@@ -76,10 +76,10 @@ class TestSiluMulFp8ConfigPicker:
        args = (input_tensor, scale)

        selected_key = pick_silu_mul_fp8_config(args, config_keys)
-        assert selected_key == "intermediate_4096_batchsize_256"
+        assert selected_key == "intermediate_4096_numtokens_256"

    def test_config_picker_fallback_to_default(self):
-        config_keys = ["default", "some_other_key"]
+        config_keys = ["default"]

        input_tensor = torch.randn(32, 4096, dtype=torch.bfloat16, device="cuda")
        scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
@@ -101,9 +101,9 @@ class TestSiluMulFp8ConfigPicker:
    @pytest.mark.parametrize("intermediate_size", [2048, 4096, 5120])
    def test_config_picker_different_sizes(self, intermediate_size):
        config_keys = [
-            "intermediate_2048_batchsize_256",
-            "intermediate_4096_batchsize_256",
-            "intermediate_5120_batchsize_256",
+            "intermediate_2048_numtokens_256",
+            "intermediate_4096_numtokens_256",
+            "intermediate_5120_numtokens_256",
        ]

        input_tensor = torch.randn(
@@ -113,9 +113,73 @@ class TestSiluMulFp8ConfigPicker:
        args = (input_tensor, scale)

        selected_key = pick_silu_mul_fp8_config(args, config_keys)
-        expected_key = f"intermediate_{intermediate_size}_batchsize_256"
+        expected_key = f"intermediate_{intermediate_size}_numtokens_256"
        assert selected_key == expected_key

+    def test_config_picker_numtokens_ceiling(self):
+        """Pick the smallest numtokens >= input num_tokens."""
+        config_keys = [
+            "intermediate_4096_numtokens_8",
+            "intermediate_4096_numtokens_32",
+            "intermediate_4096_numtokens_128",
+            "intermediate_4096_numtokens_256",
+        ]
+        # 20 tokens -> should pick numtokens_32 (smallest >= 20)
+        input_tensor = torch.randn(20, 8192, dtype=torch.bfloat16, device="cuda")
+        scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
+
+        selected_key = pick_silu_mul_fp8_config((input_tensor, scale), config_keys)
+        assert selected_key == "intermediate_4096_numtokens_32"
+
+    def test_config_picker_numtokens_exact(self):
+        """Exact num_tokens match is preferred over ceiling."""
+        config_keys = [
+            "intermediate_4096_numtokens_8",
+            "intermediate_4096_numtokens_32",
+            "intermediate_4096_numtokens_128",
+        ]
+        input_tensor = torch.randn(32, 8192, dtype=torch.bfloat16, device="cuda")
+        scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
+
+        selected_key = pick_silu_mul_fp8_config((input_tensor, scale), config_keys)
+        assert selected_key == "intermediate_4096_numtokens_32"
+
+    def test_config_picker_numtokens_fallback_to_largest(self):
+        """Fall back to the largest numtokens when input exceeds all."""
+        config_keys = [
+            "intermediate_4096_numtokens_8",
+            "intermediate_4096_numtokens_32",
+            "intermediate_4096_numtokens_128",
+        ]
+        # 512 tokens -> exceeds all available, should pick largest (128)
+        input_tensor = torch.randn(512, 8192, dtype=torch.bfloat16, device="cuda")
+        scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
+
+        selected_key = pick_silu_mul_fp8_config((input_tensor, scale), config_keys)
+        assert selected_key == "intermediate_4096_numtokens_128"
+
+    def test_config_picker_malformed_key_raises(self):
+        """Malformed config keys should raise ValueError."""
+        config_keys = ["intermediate_4096_badformat_256"]
+        input_tensor = torch.randn(32, 8192, dtype=torch.bfloat16, device="cuda")
+        scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
+
+        with pytest.raises(ValueError, match="Malformed config key"):
+            pick_silu_mul_fp8_config((input_tensor, scale), config_keys)
+
+    def test_config_picker_default_ignored_when_valid_keys_exist(self):
+        """'default' is skipped in favor of a real match."""
+        config_keys = [
+            "default",
+            "intermediate_4096_numtokens_32",
+            "intermediate_4096_numtokens_128",
+        ]
+        input_tensor = torch.randn(64, 8192, dtype=torch.bfloat16, device="cuda")
+        scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
+
+        selected_key = pick_silu_mul_fp8_config((input_tensor, scale), config_keys)
+        assert selected_key == "intermediate_4096_numtokens_128"
+

 class TestSiluMulFp8Correctness:
    @pytest.mark.parametrize("batch_size", [1, 8, 32, 128])
@@ -11,11 +11,13 @@ from vllm.kernels.helion.utils import canonicalize_gpu_name
    "driver_reported_name,expected",
    [
        ("NVIDIA H200", "nvidia_h200"),
-        ("NVIDIA A100-SXM4-80GB", "nvidia_a100_sxm4_80gb"),
-        ("NVIDIA H100 80GB HBM3", "nvidia_h100_80gb_hbm3"),
+        ("NVIDIA A100-SXM4-80GB", "nvidia_a100"),
+        ("NVIDIA H100 80GB HBM3", "nvidia_h100"),
+        ("NVIDIA H100 PCIe", "nvidia_h100"),
+        ("NVIDIA H100 SXM5", "nvidia_h100"),
        ("NVIDIA GeForce RTX 4090", "nvidia_geforce_rtx_4090"),
        ("AMD Instinct MI300X", "amd_instinct_mi300x"),
-        ("Tesla V100-SXM2-32GB", "tesla_v100_sxm2_32gb"),
+        ("Tesla V100-SXM2-32GB", "tesla_v100"),
    ],
 )
 def test_canonicalize_gpu_name(driver_reported_name, expected):
@@ -398,80 +398,3 @@ def test_convert_moe_weights_to_flashinfer_trtllm_block_layout(

    assert w13_converted.shape[0] == num_experts
    assert w2_converted.shape[0] == num_experts
-
-
-def test_flashinfer_blockscale_fp8_none_expert_group(monkeypatch):
-    """Test that flashinfer_fused_moe_blockscale_fp8 handles num_expert_group=None.
-
-    Regression test for https://github.com/vllm-project/vllm/issues/34477
-    MiniMax-M2.1 uses sigmoid scoring with e_score_correction_bias but no
-    grouped top-k, resulting in num_expert_group=None. This triggered a crash
-    in the flashinfer kernel when DeepSeekV3 routing was selected.
-    """
-    if not current_platform.has_device_capability(100):
-        pytest.skip("Test requires SM >= 100 (Blackwell)")
-
-    import vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe  # noqa: E501, F401
-    from tests.kernels.quant_utils import native_per_token_group_quant_fp8
-
-    set_random_seed(7)
-    monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", "8192")
-
-    e = 16  # num_experts (must be divisible by 4)
-    topk = 6  # top_k > 1 triggers DeepSeekV3 routing with sigmoid
-    m, n, k = 10, 4096, 5120
-    block_shape = [128, 128]
-    block_k = block_shape[1]
-
-    with set_current_vllm_config(vllm_config):
-        # Create BF16 hidden states
-        x = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) / 10
-
-        # Create FP8 block-scale quantized weights
-        w13_bf16 = torch.randn((e, 2 * n, k), device="cuda", dtype=torch.bfloat16) / 10
-        w2_bf16 = torch.randn((e, k, n), device="cuda", dtype=torch.bfloat16) / 10
-
-        # Quantize weights per-block to FP8
-        w13_fp8_list, w13_scale_list = [], []
-        w2_fp8_list, w2_scale_list = [], []
-        for i in range(e):
-            wq, ws = native_per_token_group_quant_fp8(w13_bf16[i], block_k)
-            w13_fp8_list.append(wq)
-            w13_scale_list.append(ws)
-
-            wq, ws = native_per_token_group_quant_fp8(w2_bf16[i], block_k)
-            w2_fp8_list.append(wq)
-            w2_scale_list.append(ws)
-
-        w13_fp8 = torch.stack(w13_fp8_list)
-        w13_scale = torch.stack(w13_scale_list)
-        w2_fp8 = torch.stack(w2_fp8_list)
-        w2_scale = torch.stack(w2_scale_list)
-
-        # DeepSeekV3 routing uses float32 logits + optional bias
-        routing_logits = torch.randn((m, e), device="cuda", dtype=torch.float32)
-        routing_bias = torch.randn(e, device="cuda", dtype=torch.float32)
-
-        # This should NOT crash with num_expert_group=None
-        output = torch.ops.vllm.flashinfer_fused_moe_blockscale_fp8(
-            routing_logits=routing_logits,
-            routing_bias=routing_bias,
-            x=x,
-            w13_weight=w13_fp8,
-            w13_weight_scale_inv=w13_scale,
-            w2_weight=w2_fp8,
-            w2_weight_scale_inv=w2_scale,
-            global_num_experts=e,
-            top_k=topk,
-            num_expert_group=None,
-            topk_group=None,
-            intermediate_size=n,
-            expert_offset=0,
-            local_num_experts=e,
-            block_shape=block_shape,
-            routing_method_type=RoutingMethodType.DeepSeekV3,
-            routed_scaling=1.0,
-        )
-
-        assert output is not None
-        assert output.shape == (m, k)
@@ -8,6 +8,7 @@ Run `pytest tests/kernels/moe/test_grouped_topk.py`.
 import pytest
 import torch

+import vllm.model_executor.layers.batch_invariant as batch_invariant
 from vllm.config import (
    CompilationConfig,
    VllmConfig,
@@ -27,11 +28,17 @@ from vllm.utils.torch_utils import set_random_seed
 )
@pytest.mark.parametrize("n_token", [1, 33, 64])
@pytest.mark.parametrize("n_hidden", [1024, 2048])
-@pytest.mark.parametrize("n_expert", [16])
-@pytest.mark.parametrize("topk", [2])
+@pytest.mark.parametrize(
+    "n_expert,topk,num_expert_group,topk_group",
+    [
+        (16, 2, 8, 2),
+        (128, 2, 8, 2),
+        (256, 8, 8, 4),
+        (384, 8, 1, 1),
+        (512, 22, 1, 1),
+    ],
+)
@pytest.mark.parametrize("renormalize", [True, False])
-@pytest.mark.parametrize("num_expert_group", [8])
-@pytest.mark.parametrize("topk_group", [2])
@pytest.mark.parametrize("scoring_func", ["softmax", "sigmoid"])
@pytest.mark.parametrize("routed_scaling_factor", [1.0, 2.5])
@pytest.mark.parametrize("input_dtype", [torch.bfloat16, torch.float32])
@@ -42,9 +49,9 @@ def test_grouped_topk(
    n_hidden: int,
    n_expert: int,
    topk: int,
-    renormalize: bool,
    num_expert_group: int,
    topk_group: int,
+    renormalize: bool,
    scoring_func: str,
    routed_scaling_factor: float,
    input_dtype: torch.dtype,
@@ -62,6 +69,7 @@ def test_grouped_topk(

    with set_current_vllm_config(vllm_config), monkeypatch.context() as m:
        m.setenv("VLLM_USE_FUSED_MOE_GROUPED_TOPK", "0")
+        m.setattr(batch_invariant, "VLLM_BATCH_INVARIANT", True)
        grouped_topk = GroupedTopk(
            topk=topk,
            renormalize=renormalize,
@@ -89,8 +97,7 @@ def test_grouped_topk(
            e_score_correction_bias=e_score_correction_bias,
        )

-        if renormalize:
-            torch.testing.assert_close(
-                baseline_topk_weights, test_topk_weights, atol=2e-2, rtol=0
-            )
+        torch.testing.assert_close(
+            baseline_topk_weights, test_topk_weights, atol=2e-2, rtol=0
+        )
        torch.testing.assert_close(baseline_topk_ids, test_topk_ids, atol=0, rtol=0)
@@ -14,6 +14,7 @@ from tests.kernels.utils import torch_moe
 from vllm import _custom_ops as ops
 from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
 from vllm.model_executor.layers.fused_moe import fused_topk
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
 from vllm.model_executor.layers.fused_moe.config import nvfp4_moe_quant_config
 from vllm.model_executor.layers.fused_moe.cutlass_moe import (
    CutlassExpertsFp4,
@@ -147,5 +148,130 @@ def test_cutlass_fp4_moe_no_graph(
        torch.testing.assert_close(torch_output, cutlass_output, atol=1e-1, rtol=1e-1)


+# step3.5-flash uses swiglustep activation (clipped SwiGLU with limit=7.0)
+# for MoE layers 43-44. This tests the non-fused activation fallback path
+# in run_cutlass_moe_fp4 (apply_moe_activation + separate fp4 quantization).
+# Model dims: e=288, topk=8, n=1280 (moe_intermediate_size), k=4096 (hidden)
+SWIGLUSTEP_MNK_FACTORS = [
+    (2, 1280, 4096),
+    (64, 1280, 4096),
+    (224, 1280, 4096),
+]
+
+
+@pytest.mark.parametrize("m,n,k", SWIGLUSTEP_MNK_FACTORS)
+@pytest.mark.parametrize("e", [64, 288])
+@pytest.mark.parametrize("topk", [1, 8])
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
+@torch.inference_mode()
+def test_cutlass_fp4_moe_swiglustep(
+    m: int, n: int, k: int, e: int, topk: int, dtype: torch.dtype, workspace_init
+):
+    set_random_seed(7)
+    with set_current_vllm_config(
+        VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1))
+    ):
+        quant_blocksize = 16
+
+        a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
+
+        (_, w1_q, w1_blockscale, w1_gs), (_, w2_q, w2_blockscale, w2_gs) = (
+            make_test_weights(
+                e,
+                n,
+                k,
+                in_dtype=dtype,
+                quant_dtype="nvfp4",
+                block_shape=None,
+                per_out_ch_quant=False,
+            )
+        )
+
+        score = torch.randn((m, e), device="cuda", dtype=dtype)
+        topk_weights, topk_ids, _ = fused_topk(a, score, topk, renormalize=False)
+
+        a1_gs = torch.ones((e,), device="cuda", dtype=torch.float32)
+        a2_gs = torch.ones((e,), device="cuda", dtype=torch.float32)
+
+        assert w1_gs is not None
+        assert w2_gs is not None
+        assert w1_blockscale is not None
+        assert w2_blockscale is not None
+
+        quant_config = nvfp4_moe_quant_config(
+            g1_alphas=(1 / w1_gs),
+            g2_alphas=(1 / w2_gs),
+            a1_gscale=a1_gs,
+            a2_gscale=a2_gs,
+            w1_scale=w1_blockscale,
+            w2_scale=w2_blockscale,
+        )
+
+        kernel = mk.FusedMoEModularKernel(
+            MoEPrepareAndFinalizeNoEP(),
+            CutlassExpertsFp4(
+                moe_config=make_dummy_moe_config(),
+                quant_config=quant_config,
+            ),
+            inplace=False,
+        )
+
+        cutlass_output = kernel(
+            hidden_states=a,
+            w1=w1_q,
+            w2=w2_q,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            activation=MoEActivation.SWIGLUSTEP,
+        )
+
+        # Reference: dequantize everything and run torch_moe with swiglustep
+        a_global_scale = (
+            (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(a.flatten(), dim=-1)
+        ).to(torch.float32)
+        a_fp4, a_scale_interleaved = ops.scaled_fp4_quant(a, a_global_scale)
+
+        a_in_dtype = dequantize_nvfp4_to_dtype(
+            a_fp4,
+            a_scale_interleaved,
+            a_global_scale,
+            dtype=a.dtype,
+            device=a.device,
+            block_size=quant_blocksize,
+        )
+
+        w1_d = torch.empty((e, 2 * n, k), device="cuda", dtype=dtype)
+        w2_d = torch.empty((e, k, n), device="cuda", dtype=dtype)
+
+        for idx in range(0, e):
+            w1_d[idx] = dequantize_nvfp4_to_dtype(
+                w1_q[idx],
+                w1_blockscale[idx],
+                w1_gs[idx],
+                dtype=dtype,
+                device=w1_q.device,
+                block_size=quant_blocksize,
+            )
+            w2_d[idx] = dequantize_nvfp4_to_dtype(
+                w2_q[idx],
+                w2_blockscale[idx],
+                w2_gs[idx],
+                dtype=dtype,
+                device=w2_q.device,
+                block_size=quant_blocksize,
+            )
+
+        torch_output = torch_moe(
+            a_in_dtype,
+            w1_d,
+            w2_d,
+            score,
+            topk,
+            activation=MoEActivation.SWIGLUSTEP,
+        )
+
+        torch.testing.assert_close(torch_output, cutlass_output, atol=1e-1, rtol=1e-1)
+
+
 if __name__ == "__main__":
    test_cutlass_fp4_moe_no_graph((2, 1024, 1024), 40, 1, torch.half)
@@ -155,9 +155,9 @@ def test_rocm_wvsplitkrc_kernel(xnorm, n, k, m, dtype, seed, bias_mode):
    out = ops.wvSplitKrc(B, A.view(-1, A.size(-1)), cu_count, BIAS)

    if xnorm:
-        assert torch.allclose(out, ref_out, atol=1e-3, rtol=1e-8)
+        torch.testing.assert_close(out, ref_out, atol=1e-3, rtol=1e-8)
    else:
-        assert torch.allclose(out, ref_out, atol=1e-3, rtol=1e-2)
+        torch.testing.assert_close(out, ref_out, atol=1e-3, rtol=1e-2)


@pytest.mark.parametrize("n,k,m", NKM_FACTORS_LLMM1)
@@ -177,7 +177,7 @@ def test_rocm_llmm1_kernel(n, k, m, dtype, rows_per_block, seed):
    ref_out = torch.matmul(A, B.t())
    out = ops.LLMM1(B, A, rows_per_block)

-    assert torch.allclose(out, ref_out, rtol=0.01)
+    torch.testing.assert_close(out, ref_out, atol=1e-8, rtol=1e-2)


@pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITK)
@@ -194,7 +194,7 @@ def test_rocm_wvsplitk_kernel(n, k, m, dtype, seed):
    ref_out = torch.nn.functional.linear(A, B)
    out = ops.wvSplitK(B, A.view(-1, A.size(-1)), cu_count)

-    assert torch.allclose(out, ref_out, rtol=0.01)
+    torch.testing.assert_close(out, ref_out, atol=1e-8, rtol=1e-2)


@pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITK)
@@ -213,7 +213,7 @@ def test_rocm_wvsplitk_bias1D_kernel(n, k, m, dtype, seed):
    ref_out = torch.nn.functional.linear(A, B, BIAS)
    out = ops.wvSplitK(B, A.view(-1, A.size(-1)), cu_count, BIAS)

-    assert torch.allclose(out, ref_out, rtol=0.01)
+    torch.testing.assert_close(out, ref_out, atol=1e-8, rtol=1e-2)


@pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITK)
@@ -232,7 +232,7 @@ def test_rocm_wvsplitk_bias2D_kernel(n, k, m, dtype, seed):
    ref_out = torch.nn.functional.linear(A, B, BIAS)
    out = ops.wvSplitK(B, A.view(-1, A.size(-1)), cu_count, BIAS)

-    assert torch.allclose(out, ref_out, rtol=0.01)
+    torch.testing.assert_close(out, ref_out, atol=1e-8, rtol=1e-2)


@pytest.mark.parametrize("xnorm", [False, True])
@@ -275,4 +275,4 @@ def test_rocm_wvsplitk_fp8_kernel(
        # wider pytrch thresh for large-K & no xnorm
        torch.testing.assert_close(out, ref_out, atol=0.07, rtol=5e-2)
    else:
-        torch.testing.assert_close(out, ref_out, atol=0.01, rtol=0.01)
+        torch.testing.assert_close(out, ref_out, atol=1e-2, rtol=1e-2)
@@ -153,5 +153,5 @@ def test_default_mm_lora_does_not_expand_string_reqs(vllm_runner):
        # Then check to make sure the submitted lora request
        # and text prompt were zipped together correctly
        engine_args, engine_kwargs = mock_add_request.call_args
+        assert engine_args[1]["prompt"] == AUDIO_PROMPT
        assert engine_kwargs["lora_request"] is None
-        assert engine_kwargs["prompt_text"] == AUDIO_PROMPT
@@ -88,9 +88,8 @@ class Qwen2VLTester:
        # Validate outputs
        for generated, expected in zip(generated_texts, expected_outputs):
            assert expected.startswith(generated), (
-                f"Generated text {generated} doesn't "
+                f"Generated text {generated} doesn't match expected pattern {expected}"
            )
-            f"match expected pattern {expected}"

    def run_beam_search_test(
        self,
@@ -118,11 +117,14 @@ class Qwen2VLTester:
            inputs, beam_search_params, lora_request=lora_request
        )

-        for output_obj, expected_outs in zip(outputs, expected_outputs):
+        for output_obj, expected_texts in zip(outputs, expected_outputs):
            output_texts = [seq.text for seq in output_obj.sequences]
-            assert output_texts == expected_outs, (
-                f"Generated texts {output_texts} do not match expected {expected_outs}"
-            )  # noqa: E501
+
+            for output_text, expected_text in zip(output_texts, expected_texts):
+                # NOTE beam search .text contains the whole text including inputs
+                assert output_text.endswith(expected_text), (
+                    f"Generated {output_text} does not match expected {expected_text}"
+                )


 TEST_IMAGES = [
@@ -151,11 +153,10 @@ EXPECTED_OUTPUTS_VISION_NO_CONNECTOR = [
    "A closeup shot of the Tokyo Skytree with pink flowers in the foreground.",
 ]

-# NOTE - beam search .text contains the whole text
 EXPECTED_BEAM_SEARCH_OUTPUTS = [
    [
-        "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>What is in the image?<|im_end|>\n<|im_start|>assistant\nA majestic skyscraper stands",  # noqa: E501
-        "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>What is in the image?<|im_end|>\n<|im_start|>assistant\nA majestic tower stands tall",  # noqa: E501
+        "A majestic skyscraper stands",
+        "A majestic tower stands tall",
    ],
 ]

@@ -20,6 +20,12 @@ COLBERT_MODELS = {
        "colbert_dim": 96,
        "max_model_len": 512,
        "extra_kwargs": {},
+        "hf_comparison": {
+            "weights_file": "model.safetensors",
+            "weights_key": "linear.weight",
+            "trust_remote_code": False,
+            "model_cls": "BertModel",
+        },
    },
    "modernbert": {
        "model": "lightonai/GTE-ModernColBERT-v1",
@@ -30,6 +36,12 @@ COLBERT_MODELS = {
                "architectures": ["ColBERTModernBertModel"],
            },
        },
+        "hf_comparison": {
+            "weights_file": "1_Dense/model.safetensors",
+            "weights_key": "linear.weight",
+            "trust_remote_code": False,
+            "model_cls": "AutoModel",
+        },
    },
    "jina": {
        "model": "jinaai/jina-colbert-v2",
@@ -40,9 +52,16 @@ COLBERT_MODELS = {
                "architectures": ["ColBERTJinaRobertaModel"],
            },
        },
+        "hf_comparison": {
+            "weights_file": "model.safetensors",
+            "weights_key": "linear.weight",
+            "trust_remote_code": True,
+            "model_cls": "AutoModel",
+        },
    },
 }

+
 TEXTS_1 = [
    "What is the capital of France?",
    "What is the capital of Germany?",
@@ -56,9 +75,68 @@ TEXTS_2 = [
 DTYPE = "half"


-# -----------------------------------------------------------------------
-# Fixtures
-# -----------------------------------------------------------------------
+def _load_hf_model(model_name: str, hf_spec: dict, device: torch.device):
+    """Load HF model on the given device with a compatible attention impl."""
+    from transformers import AutoModel, BertModel
+
+    cls = BertModel if hf_spec["model_cls"] == "BertModel" else AutoModel
+    trust = hf_spec.get("trust_remote_code", False)
+
+    # Flash / Triton kernels require GPU tensors; fall back to eager on CPU.
+    extra = {}
+    if device.type == "cpu":
+        extra["attn_implementation"] = "eager"
+
+    model = cls.from_pretrained(
+        model_name,
+        trust_remote_code=trust,
+        **extra,
+    ).to(device)
+    model.eval()
+    return model
+
+
+def _load_projection_weight(model_name: str, hf_spec: dict, device: torch.device):
+    """Download and return the ColBERT linear projection weight."""
+    from huggingface_hub import hf_hub_download
+    from safetensors.torch import load_file
+
+    path = hf_hub_download(model_name, filename=hf_spec["weights_file"])
+    weights = load_file(path)
+    return weights[hf_spec["weights_key"]].to(device)
+
+
+def _compute_hf_colbert_embeddings(model, tokenizer, linear_weight, texts, device):
+    """Run HF model + projection and return L2-normalised token embeddings."""
+    import torch.nn.functional as F
+
+    embeddings = []
+    for text in texts:
+        inputs = tokenizer(text, return_tensors="pt").to(device)
+        with torch.no_grad():
+            hidden = model(**inputs).last_hidden_state.float()
+            projected = F.linear(hidden, linear_weight.float())
+            normalised = F.normalize(projected, p=2, dim=-1)
+            embeddings.append(normalised.squeeze(0).cpu())
+    return embeddings
+
+
+def _assert_embeddings_close(vllm_outputs, hf_embeddings):
+    """Assert that vLLM and HuggingFace embeddings match."""
+    for i, (hf_emb, vllm_out) in enumerate(zip(hf_embeddings, vllm_outputs)):
+        vllm_emb = torch.as_tensor(vllm_out).float()
+
+        assert hf_emb.shape == vllm_emb.shape, (
+            f"Shape mismatch for text {i}: HF {hf_emb.shape} vs vLLM {vllm_emb.shape}"
+        )
+
+        torch.testing.assert_close(
+            vllm_emb,
+            hf_emb,
+            rtol=1e-2,
+            atol=1e-2,
+            msg=f"Embedding mismatch for text {i}",
+        )


@pytest.fixture(params=list(COLBERT_MODELS.keys()), scope="module")
@@ -87,11 +165,6 @@ def colbert_extra_kwargs(colbert_spec):
    return colbert_spec["extra_kwargs"]


-# -----------------------------------------------------------------------
-# Tests
-# -----------------------------------------------------------------------
-
-
 def test_colbert_token_embed(
    vllm_runner,
    colbert_model_name,
@@ -111,7 +184,7 @@ def test_colbert_token_embed(
        outputs = vllm_model.token_embed([TEXTS_1[0]])

        assert len(outputs) == 1
-        emb = torch.tensor(outputs[0])
+        emb = torch.as_tensor(outputs[0])
        assert emb.dim() == 2
        assert emb.shape[1] == colbert_dim
        assert emb.shape[0] > 1
@@ -135,8 +208,8 @@ def test_colbert_late_interaction_1_to_1(
        q_outputs = vllm_model.token_embed([TEXTS_1[0]])
        d_outputs = vllm_model.token_embed([TEXTS_2[0]])

-        q_emb = torch.tensor(q_outputs[0])
-        d_emb = torch.tensor(d_outputs[0])
+        q_emb = torch.as_tensor(q_outputs[0])
+        d_emb = torch.as_tensor(d_outputs[0])

        manual_score = compute_maxsim_score(q_emb, d_emb).item()

@@ -164,11 +237,11 @@ def test_colbert_late_interaction_1_to_N(
        q_outputs = vllm_model.token_embed([TEXTS_1[0]])
        d_outputs = vllm_model.token_embed(TEXTS_2)

-        q_emb = torch.tensor(q_outputs[0])
+        q_emb = torch.as_tensor(q_outputs[0])

        manual_scores = []
        for d_out in d_outputs:
-            d_emb = torch.tensor(d_out)
+            d_emb = torch.as_tensor(d_out)
            manual_scores.append(compute_maxsim_score(q_emb, d_emb).item())

        vllm_scores = vllm_model.score(TEXTS_1[0], TEXTS_2)
@@ -198,8 +271,8 @@ def test_colbert_late_interaction_N_to_N(

        manual_scores = []
        for q_out, d_out in zip(q_outputs, d_outputs):
-            q_emb = torch.tensor(q_out)
-            d_emb = torch.tensor(d_out)
+            q_emb = torch.as_tensor(q_out)
+            d_emb = torch.as_tensor(d_out)
            manual_scores.append(compute_maxsim_score(q_emb, d_emb).item())

        vllm_scores = vllm_model.score(TEXTS_1, TEXTS_2)
@@ -259,79 +332,16 @@ def test_colbert_embed_not_supported(
        vllm_model.embed([TEXTS_1[0]])


-# -----------------------------------------------------------------------
-# Per-model HuggingFace comparison tests
-# -----------------------------------------------------------------------
+@pytest.mark.parametrize("backend", list(COLBERT_MODELS.keys()))
+def test_colbert_hf_comparison(vllm_runner, backend):
+    """Test that vLLM ColBERT embeddings match HuggingFace for each backend."""
+    from transformers import AutoTokenizer

-
-def _assert_embeddings_close(vllm_outputs, hf_embeddings):
-    """Assert that vLLM and HuggingFace embeddings match."""
-    for i, (hf_emb, vllm_out) in enumerate(zip(hf_embeddings, vllm_outputs)):
-        vllm_emb = torch.tensor(vllm_out).float()
-
-        assert hf_emb.shape == vllm_emb.shape, (
-            f"Shape mismatch for text {i}: HF {hf_emb.shape} vs vLLM {vllm_emb.shape}"
-        )
-
-        torch.testing.assert_close(
-            vllm_emb,
-            hf_emb,
-            rtol=1e-2,
-            atol=1e-2,
-            msg=f"Embedding mismatch for text {i}",
-        )
-
-
-def test_colbert_hf_comparison_bert(vllm_runner):
-    """Test that vLLM ColBERT produces same embeddings as HuggingFace (BERT)."""
-    import torch.nn.functional as F
-    from huggingface_hub import hf_hub_download
-    from safetensors.torch import load_file
-    from transformers import AutoTokenizer, BertModel
-
-    model_name = COLBERT_MODELS["bert"]["model"]
-    test_texts = [TEXTS_1[0], TEXTS_2[0]]
-
-    with vllm_runner(
-        model_name,
-        runner="pooling",
-        dtype="float32",
-        max_model_len=512,
-        enforce_eager=True,
-    ) as vllm_model:
-        vllm_outputs = vllm_model.token_embed(test_texts)
-
-    hf_tokenizer = AutoTokenizer.from_pretrained(model_name)
-    hf_bert = BertModel.from_pretrained(model_name)
-    hf_bert.eval()
-
-    weights_path = hf_hub_download(model_name, filename="model.safetensors")
-    weights = load_file(weights_path)
-    linear_weight = weights["linear.weight"]  # [96, 384]
-
-    hf_embeddings = []
-    for text in test_texts:
-        inputs = hf_tokenizer(text, return_tensors="pt")
-        with torch.no_grad():
-            outputs = hf_bert(**inputs)
-            hidden_states = outputs.last_hidden_state
-            token_emb = F.linear(hidden_states, linear_weight)
-            token_emb = F.normalize(token_emb, p=2, dim=-1)
-            hf_embeddings.append(token_emb.squeeze(0).float())
-
-    _assert_embeddings_close(vllm_outputs, hf_embeddings)
-
-
-def test_colbert_hf_comparison_modernbert(vllm_runner):
-    """Test that vLLM ColBERT produces same embeddings as HuggingFace
-    (ModernBERT)."""
-    import torch.nn.functional as F
-    from huggingface_hub import hf_hub_download
-    from safetensors.torch import load_file
-    from transformers import AutoModel, AutoTokenizer
-
-    spec = COLBERT_MODELS["modernbert"]
+    spec = COLBERT_MODELS[backend]
+    hf_spec = spec["hf_comparison"]
    model_name = spec["model"]
+    assert isinstance(model_name, str)
+    assert isinstance(hf_spec, dict)
    test_texts = [TEXTS_1[0], TEXTS_2[0]]

    with vllm_runner(
@@ -344,73 +354,21 @@ def test_colbert_hf_comparison_modernbert(vllm_runner):
    ) as vllm_model:
        vllm_outputs = vllm_model.token_embed(test_texts)

-    hf_tokenizer = AutoTokenizer.from_pretrained(model_name)
-    hf_model = AutoModel.from_pretrained(model_name)
-    hf_model.eval()
-
-    # Load projection from sentence-transformers 1_Dense layer
-    dense_path = hf_hub_download(model_name, filename="1_Dense/model.safetensors")
-    dense_weights = load_file(dense_path)
-    linear_weight = dense_weights["linear.weight"]  # [128, 768]
-
-    hf_embeddings = []
-    for text in test_texts:
-        inputs = hf_tokenizer(text, return_tensors="pt")
-        with torch.no_grad():
-            outputs = hf_model(**inputs)
-            hidden_states = outputs.last_hidden_state
-            token_emb = F.linear(hidden_states, linear_weight)
-            token_emb = F.normalize(token_emb, p=2, dim=-1)
-            hf_embeddings.append(token_emb.squeeze(0).float())
-
-    _assert_embeddings_close(vllm_outputs, hf_embeddings)
-
-
-def test_colbert_hf_comparison_jina(vllm_runner):
-    """Test that vLLM ColBERT produces same embeddings as HuggingFace
-    (Jina XLM-RoBERTa)."""
-    import torch.nn.functional as F
-    from huggingface_hub import hf_hub_download
-    from safetensors.torch import load_file
-    from transformers import AutoModel, AutoTokenizer
-
-    spec = COLBERT_MODELS["jina"]
-    model_name = spec["model"]
-    test_texts = [TEXTS_1[0], TEXTS_2[0]]
-
-    with vllm_runner(
-        model_name,
-        runner="pooling",
-        dtype="float32",
-        max_model_len=spec["max_model_len"],
-        enforce_eager=True,
-        **spec["extra_kwargs"],
-    ) as vllm_model:
-        vllm_outputs = vllm_model.token_embed(test_texts)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    hf_tokenizer = AutoTokenizer.from_pretrained(
        model_name,
-        trust_remote_code=True,
+        trust_remote_code=hf_spec.get("trust_remote_code", False),
    )
-    hf_model = AutoModel.from_pretrained(
-        model_name,
-        trust_remote_code=True,
+    hf_model = _load_hf_model(model_name, hf_spec, device)
+    linear_weight = _load_projection_weight(model_name, hf_spec, device)
+
+    hf_embeddings = _compute_hf_colbert_embeddings(
+        hf_model,
+        hf_tokenizer,
+        linear_weight,
+        test_texts,
+        device,
    )
-    hf_model.eval()
-
-    # Load projection from main checkpoint
-    weights_path = hf_hub_download(model_name, filename="model.safetensors")
-    weights = load_file(weights_path)
-    linear_weight = weights["linear.weight"]  # [128, 1024]
-
-    hf_embeddings = []
-    for text in test_texts:
-        inputs = hf_tokenizer(text, return_tensors="pt")
-        with torch.no_grad():
-            outputs = hf_model(**inputs)
-            hidden_states = outputs.last_hidden_state
-            token_emb = F.linear(hidden_states.float(), linear_weight.float())
-            token_emb = F.normalize(token_emb, p=2, dim=-1)
-            hf_embeddings.append(token_emb.squeeze(0).float())

    _assert_embeddings_close(vllm_outputs, hf_embeddings)
@@ -191,6 +191,9 @@ def run_mteb_rerank(cross_encoder: mteb.CrossEncoderProtocol, tasks, languages):
        mteb_tasks: list[mteb.abstasks.AbsTaskRetrieval] = mteb.get_tasks(
            tasks=tasks, languages=languages, eval_splits=eval_splits
        )
+        for task in mteb_tasks:
+            if not task.data_loaded:
+                task.load_data()

        mteb.evaluate(
            bm25s,
@@ -305,10 +305,10 @@ def create_text_model_weights(text_config: dict[str, Any]) -> dict[str, torch.Te

        # Self-attention weights (separate q, k, v projections)
        weights[f"{layer_prefix}.self_attn.q_proj.weight"] = torch.randn(
-            hidden_size, num_attention_heads * head_dim, dtype=torch.bfloat16
+            num_attention_heads * head_dim, hidden_size, dtype=torch.bfloat16
        )
        weights[f"{layer_prefix}.self_attn.k_proj.weight"] = torch.randn(
-            hidden_size, num_key_value_heads * head_dim, dtype=torch.bfloat16
+            num_key_value_heads * head_dim, hidden_size, dtype=torch.bfloat16
        )
        weights[f"{layer_prefix}.self_attn.v_proj.weight"] = torch.randn(
            num_key_value_heads * head_dim, hidden_size, dtype=torch.bfloat16
@@ -111,6 +111,47 @@ def check_model_available(model: str) -> None:
    model_info.check_transformers_version(on_fail="skip")


+def test_parse_language_detection_output():
+    """Unit test for WhisperForConditionalGeneration.parse_language_detection_output.
+
+    No GPU or model loading required.
+    """
+    from unittest.mock import MagicMock
+
+    from vllm.model_executor.models.whisper import (
+        WhisperForConditionalGeneration,
+    )
+
+    cls = WhisperForConditionalGeneration
+
+    def make_tokenizer(return_value: str) -> MagicMock:
+        tok = MagicMock()
+        tok.decode = MagicMock(return_value=return_value)
+        return tok
+
+    # English
+    assert (
+        cls.parse_language_detection_output([50259], make_tokenizer("<|en|>")) == "en"
+    )
+
+    # German
+    assert (
+        cls.parse_language_detection_output([50261], make_tokenizer("<|de|>")) == "de"
+    )
+
+    # Unsupported language code
+    with pytest.raises(AssertionError):
+        cls.parse_language_detection_output([99999], make_tokenizer("<|xx|>"))
+
+    # No special token format
+    with pytest.raises(AssertionError):
+        cls.parse_language_detection_output([1], make_tokenizer("hello"))
+
+    # Empty token_ids
+    with pytest.raises((AssertionError, IndexError)):
+        cls.parse_language_detection_output([], make_tokenizer("anything"))
+
+
@pytest.mark.core_model
@pytest.mark.cpu_model
@pytest.mark.parametrize("model", ["openai/whisper-large-v3-turbo"])
@@ -0,0 +1,115 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for ColModernVBERT multimodal late-interaction model.
+
+ColModernVBERT combines SigLIP vision encoder + ModernBERT text encoder
+with a pixel shuffle connector and ColBERT-style 128-dim per-token
+embeddings for visual document retrieval.
+"""
+
+import pytest
+import torch
+
+from vllm.entrypoints.pooling.score.utils import compute_maxsim_score
+
+MODEL_NAME = "ModernVBERT/colmodernvbert-merged"
+COLBERT_DIM = 128
+DTYPE = "half"
+
+
+# -----------------------------------------------------------------------
+# Text-only tests
+# -----------------------------------------------------------------------
+
+
+def test_colmodernvbert_text_token_embed(vllm_runner):
+    """Text query produces per-token embeddings with shape (seq_len, 128)."""
+    with vllm_runner(
+        MODEL_NAME,
+        runner="pooling",
+        dtype=DTYPE,
+        enforce_eager=True,
+    ) as vllm_model:
+        outputs = vllm_model.token_embed(["What is machine learning?"])
+
+        assert len(outputs) == 1
+        emb = torch.tensor(outputs[0])
+        assert emb.dim() == 2
+        assert emb.shape[1] == COLBERT_DIM
+        assert emb.shape[0] > 1
+
+
+def test_colmodernvbert_text_relevance_ordering(vllm_runner):
+    """Relevant documents score higher than irrelevant ones."""
+    query = "What is machine learning?"
+    documents = [
+        "Machine learning is a subset of artificial intelligence.",
+        "The weather in Paris is mild in spring.",
+    ]
+
+    with vllm_runner(
+        MODEL_NAME,
+        runner="pooling",
+        dtype=DTYPE,
+        enforce_eager=True,
+    ) as vllm_model:
+        scores = vllm_model.score(query, documents)
+
+        assert len(scores) == 2
+        assert scores[0] > scores[1], "ML doc should score higher than weather doc"
+
+
+def test_colmodernvbert_text_late_interaction(vllm_runner):
+    """MaxSim scoring via vLLM matches manual computation."""
+    query = "What is the capital of France?"
+    doc = "The capital of France is Paris."
+
+    with vllm_runner(
+        MODEL_NAME,
+        runner="pooling",
+        dtype=DTYPE,
+        enforce_eager=True,
+    ) as vllm_model:
+        q_out = vllm_model.token_embed([query])
+        d_out = vllm_model.token_embed([doc])
+
+        q_emb = torch.tensor(q_out[0])
+        d_emb = torch.tensor(d_out[0])
+        manual_score = compute_maxsim_score(q_emb, d_emb).item()
+
+        vllm_scores = vllm_model.score(query, doc)
+
+        assert len(vllm_scores) == 1
+        assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01)
+
+
+# -----------------------------------------------------------------------
+# Image tests
+# -----------------------------------------------------------------------
+
+
+def test_colmodernvbert_image_token_embed(vllm_runner, image_assets):
+    """Image input produces per-token embeddings including vision tokens."""
+    with vllm_runner(
+        MODEL_NAME,
+        runner="pooling",
+        dtype=DTYPE,
+        enforce_eager=True,
+    ) as vllm_model:
+        image = image_assets[0].pil_image
+        inputs = vllm_model.get_inputs(
+            [""],
+            images=[image],
+        )
+        req_outputs = vllm_model.llm.encode(
+            inputs,
+            pooling_task="token_embed",
+        )
+        outputs = [req_output.outputs.data for req_output in req_outputs]
+
+        assert len(outputs) == 1
+        emb = torch.tensor(outputs[0])
+        assert emb.dim() == 2
+        assert emb.shape[1] == COLBERT_DIM
+        # Should have at least the image tokens (64 after pixel shuffle)
+        assert emb.shape[0] >= 64
@@ -7,19 +7,31 @@ ColBERT-style late interaction scoring (MaxSim). It produces per-token
 embeddings for both text and image inputs.
 """

+import base64
+from io import BytesIO
+
 import pytest
 import torch
+from PIL import Image
+
+from vllm.entrypoints.chat_utils import (
+    ChatCompletionContentPartImageParam,
+    ChatCompletionContentPartTextParam,
+)
+from vllm.entrypoints.pooling.score.utils import ScoreMultiModalParam

 from ....conftest import VllmRunner

 MODELS = [
    "TomoroAI/tomoro-colqwen3-embed-4b",
    "OpenSearch-AI/Ops-Colqwen3-4B",
+    "nvidia/nemotron-colembed-vl-4b-v2",
 ]

 EMBED_DIMS = {
    "TomoroAI/tomoro-colqwen3-embed-4b": 320,
    "OpenSearch-AI/Ops-Colqwen3-4B": 2560,
+    "nvidia/nemotron-colembed-vl-4b-v2": 2560,
 }

 TEXT_QUERIES = [
@@ -33,6 +45,43 @@ TEXT_DOCUMENTS = [
 ]

 DTYPE = "half"
+GPU_MEMORY_UTILIZATION = 0.7
+
+
+def _make_base64_image(
+    width: int = 64, height: int = 64, color: tuple[int, int, int] = (255, 0, 0)
+) -> str:
+    """Create a small solid-color PNG image and return its base64 data URI."""
+    img = Image.new("RGB", (width, height), color)
+    buf = BytesIO()
+    img.save(buf, format="PNG")
+    b64 = base64.b64encode(buf.getvalue()).decode()
+    return f"data:image/png;base64,{b64}"
+
+
+def _make_image_mm_param(
+    image_uri: str,
+    text: str | None = None,
+) -> ScoreMultiModalParam:
+    """Build a ScoreMultiModalParam containing an image (and optional text)."""
+    content: list = [
+        ChatCompletionContentPartImageParam(
+            type="image_url",
+            image_url={"url": image_uri},
+        ),
+    ]
+    if text is not None:
+        content.append(
+            ChatCompletionContentPartTextParam(type="text", text=text),
+        )
+    return ScoreMultiModalParam(content=content)
+
+
+def _make_text_mm_param(text: str) -> ScoreMultiModalParam:
+    """Build a ScoreMultiModalParam containing only text."""
+    return ScoreMultiModalParam(
+        content=[ChatCompletionContentPartTextParam(type="text", text=text)],
+    )


 def _run_token_embed_test(
@@ -48,6 +97,7 @@ def _run_token_embed_test(
        dtype=dtype,
        max_model_len=4096,
        enforce_eager=True,
+        gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
    ) as vllm_model:
        outputs = vllm_model.token_embed([TEXT_QUERIES[0]])

@@ -83,6 +133,7 @@ def _run_late_interaction_test(
        dtype=dtype,
        max_model_len=4096,
        enforce_eager=True,
+        gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
    ) as vllm_model:
        q_outputs = vllm_model.token_embed([TEXT_QUERIES[0]])
        d_outputs = vllm_model.token_embed([TEXT_DOCUMENTS[0]])
@@ -118,6 +169,7 @@ def _run_relevance_test(
        dtype=dtype,
        max_model_len=4096,
        enforce_eager=True,
+        gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
    ) as vllm_model:
        scores = vllm_model.score(query, documents)

@@ -154,3 +206,142 @@ def test_colqwen3_relevance_ordering(
    dtype: str,
 ) -> None:
    _run_relevance_test(vllm_runner, model, dtype=dtype)
+
+
+# ── Multimodal scoring tests ────────────────────────────────
+
+
+def _run_multimodal_text_query_image_docs_test(
+    vllm_runner: type[VllmRunner],
+    model: str,
+    *,
+    dtype: str,
+) -> None:
+    """Score a text query against image documents via the multimodal path.
+
+    Verifies that score_data_to_prompts correctly handles image content
+    and produces valid MaxSim scores.
+    """
+    red_image = _make_base64_image(64, 64, color=(255, 0, 0))
+    blue_image = _make_base64_image(64, 64, color=(0, 0, 255))
+
+    query = "Describe the red object"
+    image_docs = [
+        _make_image_mm_param(red_image),
+        _make_image_mm_param(blue_image),
+    ]
+
+    with vllm_runner(
+        model,
+        runner="pooling",
+        dtype=dtype,
+        max_model_len=4096,
+        enforce_eager=True,
+        gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
+    ) as vllm_model:
+        scores = vllm_model.llm.score(query, image_docs)
+
+        assert len(scores) == 2
+        for s in scores:
+            assert isinstance(s.outputs.score, float)
+
+
+def _run_multimodal_mixed_docs_test(
+    vllm_runner: type[VllmRunner],
+    model: str,
+    *,
+    dtype: str,
+) -> None:
+    """Score a text query against a mix of text and image documents.
+
+    Ensures the late-interaction path handles heterogeneous document
+    types (plain strings alongside ScoreMultiModalParam images) in
+    a single call.
+    """
+    red_image = _make_base64_image(64, 64, color=(255, 0, 0))
+
+    query = "What is the capital of France?"
+    documents: list = [
+        "The capital of France is Paris.",
+        _make_image_mm_param(red_image),
+    ]
+
+    with vllm_runner(
+        model,
+        runner="pooling",
+        dtype=dtype,
+        max_model_len=4096,
+        enforce_eager=True,
+        gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
+    ) as vllm_model:
+        scores = vllm_model.llm.score(query, documents)
+
+        assert len(scores) == 2
+        for s in scores:
+            assert isinstance(s.outputs.score, float)
+        # Text document about France should score higher than a random image
+        assert scores[0].outputs.score > scores[1].outputs.score
+
+
+def _run_multimodal_image_query_text_docs_test(
+    vllm_runner: type[VllmRunner],
+    model: str,
+    *,
+    dtype: str,
+) -> None:
+    """Score an image query against text documents.
+
+    Verifies the reverse direction: multimodal query with text-only
+    documents through the late-interaction scoring path.
+    """
+    red_image = _make_base64_image(64, 64, color=(255, 0, 0))
+    image_query = _make_image_mm_param(red_image, text="red color")
+
+    documents = [
+        "A bright red sports car.",
+        "The weather forecast shows rain tomorrow.",
+    ]
+
+    with vllm_runner(
+        model,
+        runner="pooling",
+        dtype=dtype,
+        max_model_len=4096,
+        enforce_eager=True,
+        gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
+    ) as vllm_model:
+        scores = vllm_model.llm.score(image_query, documents)
+
+        assert len(scores) == 2
+        for s in scores:
+            assert isinstance(s.outputs.score, float)
+
+
+@pytest.mark.parametrize("model", MODELS)
+@pytest.mark.parametrize("dtype", [DTYPE])
+def test_colqwen3_multimodal_text_query_image_docs(
+    vllm_runner,
+    model: str,
+    dtype: str,
+) -> None:
+    _run_multimodal_text_query_image_docs_test(vllm_runner, model, dtype=dtype)
+
+
+@pytest.mark.parametrize("model", MODELS)
+@pytest.mark.parametrize("dtype", [DTYPE])
+def test_colqwen3_multimodal_mixed_docs(
+    vllm_runner,
+    model: str,
+    dtype: str,
+) -> None:
+    _run_multimodal_mixed_docs_test(vllm_runner, model, dtype=dtype)
+
+
+@pytest.mark.parametrize("model", MODELS)
+@pytest.mark.parametrize("dtype", [DTYPE])
+def test_colqwen3_multimodal_image_query_text_docs(
+    vllm_runner,
+    model: str,
+    dtype: str,
+) -> None:
+    _run_multimodal_image_query_text_docs_test(vllm_runner, model, dtype=dtype)
@@ -116,7 +116,7 @@ def test_dummy_data_generation(mock_ctx):
    builder = AudioFlamingo3DummyInputsBuilder(info)

    mm_counts = {"audio": 2}
-    dummy_data = builder.get_dummy_mm_data(100, mm_counts, None)
+    dummy_data = builder.get_dummy_mm_data(100, mm_counts, {})

    assert "audio" in dummy_data
    assert len(dummy_data["audio"]) == 2
--- a/Show More
+++ b/Show More