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
339 changed files with 39286 additions and 6602 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,
@@ -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.
@@ -44,7 +44,7 @@ llm = LLM(
        "model": "RedHatAI/Llama-3.1-8B-Instruct-speculator.eagle3",
        "draft_tensor_parallel_size": 2,
        "num_speculative_tokens": 2,
-        "method": "eagle",
+        "method": "eagle3",
    },
 )

@@ -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()
@@ -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
@@ -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
@@ -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):
@@ -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)
@@ -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
@@ -23,7 +23,7 @@ from vllm.multimodal.cache import MultiModalProcessorOnlyCache
 from vllm.multimodal.inputs import MultiModalInputs, batched_tensors_equal
 from vllm.multimodal.processing import BaseMultiModalProcessor, InputProcessingContext
 from vllm.tokenizers import TokenizerLike, cached_tokenizer_from_config
-from vllm.tokenizers.mistral import MistralTokenizer
+from vllm.utils.mistral import is_mistral_tokenizer

 from ....multimodal.utils import random_audio, random_image, random_video
 from ...registry import (
@@ -183,7 +183,7 @@ def get_text_token_prompts(

    text_prompt: str | None
    token_prompt: list[int]
-    if isinstance(tokenizer, MistralTokenizer):
+    if is_mistral_tokenizer(tokenizer):
        # ChatCompletionRequest only supports ImageChunk natively;
        # for other modalities (e.g. audio), fall back to the model's
        # own dummy inputs builder which knows the right placeholders.
@@ -195,6 +195,7 @@ def get_text_token_prompts(
            inputs = dummy_inputs.get_dummy_processor_inputs(
                model_config.max_model_len,
                mm_counts,
+                mm_options={},
            )
            text_prompt = None
            token_prompt = (
@@ -224,6 +225,7 @@ def get_text_token_prompts(
        inputs = dummy_inputs.get_dummy_processor_inputs(
            model_config.max_model_len,
            mm_counts,
+            mm_options={},
        )
        assert isinstance(inputs.prompt, str)

@@ -387,13 +389,13 @@ def _test_processing_correctness_one(
    mm_items = baseline_processor.info.parse_mm_data(mm_data)
    ignore_mm_keys = _IGNORE_MM_KEYS.get(model_type, set[str]())

-    baseline_tokenized_result = baseline_processor.apply(
+    baseline_tokenized_result = baseline_processor(
        token_prompt,
        mm_items=mm_items,
        hf_processor_mm_kwargs={},
    )

-    cached_tokenized_result = cached_processor.apply(
+    cached_tokenized_result = cached_processor(
        token_prompt,
        mm_items=mm_items,
        hf_processor_mm_kwargs={},
@@ -407,12 +409,12 @@ def _test_processing_correctness_one(
    )

    if text_prompt is not None:
-        baseline_text_result = baseline_processor.apply(
+        baseline_text_result = baseline_processor(
            text_prompt,
            mm_items=mm_items,
            hf_processor_mm_kwargs={},
        )
-        cached_text_result = cached_processor.apply(
+        cached_text_result = cached_processor(
            text_prompt,
            mm_items=mm_items,
            hf_processor_mm_kwargs={},
@@ -176,7 +176,7 @@ def test_get_image_size_with_most_features(

    for asset in image_assets:
        mm_data = {"image": [asset.pil_image]}
-        processed_inputs = processor.apply(
+        processed_inputs = processor(
            prompt,
            mm_items=processor.info.parse_mm_data(mm_data),
            hf_processor_mm_kwargs=hf_processor_mm_kwargs,
@@ -52,7 +52,7 @@ def test_processor_override(
    metadata["fps"] = fps
    mm_data = {"video": [(video, metadata)]}

-    processed_inputs = processor.apply(
+    processed_inputs = processor(
        prompt,
        mm_items=processor.info.parse_mm_data(mm_data),
        hf_processor_mm_kwargs=hf_processor_mm_kwargs,
@@ -104,12 +104,12 @@ def test_video_loader_consistency(
    static_mm_data = {"video": [(static_video, static_metadata)]}
    dynamic_mm_data = {"video": [(dynamic_video, dynamic_metadata)]}

-    static_outputs = processor.apply(
+    static_outputs = processor(
        prompt,
        mm_items=processor.info.parse_mm_data(static_mm_data),
        hf_processor_mm_kwargs=hf_processor_mm_kwargs,
    )
-    dynamic_outputs = processor.apply(
+    dynamic_outputs = processor(
        prompt,
        mm_items=processor.info.parse_mm_data(dynamic_mm_data),
        hf_processor_mm_kwargs=hf_processor_mm_kwargs,
@@ -106,7 +106,7 @@ def _run_check(
        for image in images
    )

-    processed_inputs = processor.apply(
+    processed_inputs = processor(
        prompt,
        mm_items=processor.info.parse_mm_data(mm_data),
        hf_processor_mm_kwargs=mm_processor_kwargs,
@@ -61,7 +61,7 @@ def test_processor_override(
    dummy_image = image_assets[0].pil_image.resize(dummy_image_size)
    mm_data = {"image": [dummy_image] * num_imgs}

-    processed_inputs = processor.apply(
+    processed_inputs = processor(
        prompt,
        mm_items=processor.info.parse_mm_data(mm_data),
        hf_processor_mm_kwargs=hf_processor_mm_kwargs,
@@ -66,7 +66,7 @@ def _run_check(
        for image in images
    )

-    processed_inputs = processor.apply(
+    processed_inputs = processor(
        prompt,
        mm_items=processor.info.parse_mm_data(mm_data),
        hf_processor_mm_kwargs=mm_processor_kwargs,
@@ -49,7 +49,7 @@ def test_processor_override(
    if tokenized_prompt:
        prompt = tokenizer.encode(prompt)

-    processed_inputs = processor.apply(
+    processed_inputs = processor(
        prompt,
        mm_items=processor.info.parse_mm_data(mm_data),
        hf_processor_mm_kwargs=mm_processor_kwargs,
@@ -87,7 +87,7 @@ def _validate_image_prompt_replacements_one(
    try:
        # The processor will throw an error if there is a mismatch
        # in the prompt replacements
-        processed_inputs = processor.apply(
+        processed_inputs = processor(
            prompt,
            mm_items=processor.info.parse_mm_data(mm_data),
            hf_processor_mm_kwargs={},
@@ -87,7 +87,7 @@ def _validate_image_prompt_replacements_one(
    try:
        # The processor will throw an error if there is a mismatch
        # in the prompt replacements
-        processed_inputs = processor.apply(
+        processed_inputs = processor(
            prompt,
            mm_items=processor.info.parse_mm_data(mm_data),
            hf_processor_mm_kwargs={},
@@ -29,7 +29,7 @@ def test_processor_override(
    image = Image.new("RGB", size=(364, 364))
    mm_data = {"image": [image] * num_imgs}

-    processed_inputs = processor.apply(
+    processed_inputs = processor(
        prompt,
        mm_items=processor.info.parse_mm_data(mm_data),
        hf_processor_mm_kwargs={},
@@ -50,7 +50,7 @@ def _validate_image_prompt_replacements_one(
    mm_data = {"image": [image] * num_imgs}

    try:
-        processed_inputs = processor.apply(
+        processed_inputs = processor(
            prompt,
            mm_items=processor.info.parse_mm_data(mm_data),
            hf_processor_mm_kwargs={},
@@ -68,7 +68,7 @@ def _run_check(
        for image in images
    )
    print(total_expected_num_patches)
-    processed_inputs = processor.apply(
+    processed_inputs = processor(
        prompt,
        mm_items=processor.info.parse_mm_data(mm_data),
        hf_processor_mm_kwargs=mm_processor_kwargs,
@@ -47,7 +47,7 @@ def test_processor_override(
    prompt = f"<|user|>\n{img_str}<|end|>\n<|assistant|>\n"
    mm_data = {"image": [image_assets[0].pil_image] * num_imgs}

-    processed_inputs = processor.apply(
+    processed_inputs = processor(
        prompt,
        mm_items=processor.info.parse_mm_data(mm_data),
        hf_processor_mm_kwargs=hf_processor_mm_kwargs,
@@ -51,7 +51,7 @@ def test_processor_override(
    dummy_image = image_assets[0].pil_image.resize(dummy_image_size)
    mm_data = {"image": [dummy_image] * num_imgs}

-    processed_inputs = processor.apply(
+    processed_inputs = processor(
        prompt,
        mm_items=processor.info.parse_mm_data(mm_data),
        hf_processor_mm_kwargs=hf_processor_mm_kwargs,
@@ -42,7 +42,7 @@ def test_processor_override(
    prompt = "<|vision_start|><|image_pad|><|vision_end|>" * num_imgs
    mm_data = {"image": [image_assets[0].pil_image] * num_imgs}

-    processed_inputs = processor.apply(
+    processed_inputs = processor(
        prompt,
        mm_items=processor.info.parse_mm_data(mm_data),
        hf_processor_mm_kwargs=hf_processor_mm_kwargs,
@@ -88,7 +88,7 @@ def test_get_image_size_with_most_features(
    prompt = "<|vision_start|><|image_pad|><|vision_end|>"
    for asset in image_assets:
        mm_data = {"image": [asset.pil_image]}
-        processed_inputs = processor.apply(
+        processed_inputs = processor(
            prompt,
            mm_items=processor.info.parse_mm_data(mm_data),
            hf_processor_mm_kwargs=hf_processor_mm_kwargs,
@@ -51,7 +51,7 @@ def test_processor_with_audio_sample_rate(
    hf_processor_mm_kwargs: dict[str, Any] = {
        "audio_sample_rate": audio_sample_rate,
    }
-    processed_inputs = processor.apply(
+    processed_inputs = processor(
        prompt,
        mm_items=processor.info.parse_mm_data(mm_data),
        hf_processor_mm_kwargs=hf_processor_mm_kwargs,
@@ -94,7 +94,7 @@ def test_longer_audio_generates_more_tokens(model_id: str) -> None:
        hf_processor_mm_kwargs: dict[str, Any] = {
            "audio_sample_rate": audio_sample_rate,
        }
-        processed = processor.apply(
+        processed = processor(
            prompt,
            mm_items=processor.info.parse_mm_data(mm_data),
            hf_processor_mm_kwargs=hf_processor_mm_kwargs,
@@ -61,7 +61,7 @@ def test_processor_override(
    dummy_image = image_assets[0].pil_image.resize(dummy_image_size)
    mm_data = {"image": [dummy_image] * num_imgs}

-    processed_inputs = processor.apply(
+    processed_inputs = processor(
        prompt,
        mm_items=processor.info.parse_mm_data(mm_data),
        hf_processor_mm_kwargs=hf_processor_mm_kwargs,
--- a/Show More
+++ b/Show More