Simplify output assignment condition

Refactor conditional check for output assignment. Signed-off-by: Zhuohan Li <zhuohan123@gmail.com>
Bump Flashinfer Version and Re-enable DeepSeek NVFP4 AR+Norm Fusion (#34899 )
2026-02-20 15:27:56 -08:00 · 2026-02-20 13:37:31 -08:00 · 2026-02-20 21:33:04 +00:00 · 2026-02-20 13:25:50 -08:00 · 2026-02-20 18:51:58 +00:00 · 2026-02-20 10:17:42 -08:00
174 changed files with 61419 additions and 4912 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
@@ -67,7 +67,7 @@ start_nodes() {
        # 3. map the huggingface cache directory to the container
        # 3. assign ip addresses to the containers (head node: 192.168.10.10, worker nodes:
        #    starting from 192.168.10.11)
-        docker run -d "$GPU_DEVICES" --shm-size=10.24gb -e HF_TOKEN \
+        docker run -d $GPU_DEVICES --shm-size=10.24gb -e HF_TOKEN \
            -v ~/.cache/huggingface:/root/.cache/huggingface --name "node$node" \
            --network docker-net --ip 192.168.10.$((10 + $node)) --rm "$DOCKER_IMAGE" \
            /bin/bash -c "tail -f /dev/null"
@@ -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!"
@@ -55,9 +55,11 @@ steps:
  grade: Blocking
  source_file_dependencies:
  - vllm/
+  - tests/detokenizer
  - tests/multimodal
  - tests/utils_
  commands:
+  - pytest -v -s detokenizer
  - pytest -v -s -m 'not cpu_test' multimodal
  - pytest -v -s utils_

@@ -547,7 +549,7 @@ steps:
  - tests/samplers
  - tests/conftest.py
  commands:
-    - pytest -v -s -m 'not skip_v1' samplers
+    - pytest -v -s samplers

 - label: LoRA Test %N # 20min each
  timeout_in_minutes: 30
@@ -1105,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
@@ -1632,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]
@@ -1707,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
@@ -1718,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
@@ -1731,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
@@ -1761,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
@@ -1772,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:
@@ -1789,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
@@ -1804,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
@@ -1822,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
@@ -1839,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
@@ -1858,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
@@ -1873,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
@@ -1888,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:
@@ -1950,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"
@@ -1971,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:
@@ -1983,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
@@ -1998,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/
@@ -2018,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
@@ -2031,7 +1989,6 @@ steps:
  timeout_in_minutes: 15
  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi355_1
-  # grade: Blocking
  source_file_dependencies:
  - vllm/
  - tests/engine
@@ -2048,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
@@ -2062,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
@@ -2073,7 +2028,6 @@ steps:
  timeout_in_minutes: 60
  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_1
-  # grade: Blocking
  source_file_dependencies:
    - vllm/
    - tests/v1
@@ -2101,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:
@@ -2141,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
@@ -2159,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
@@ -2194,7 +2145,6 @@ steps:
  timeout_in_minutes: 15
  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi355_1
-  # grade: Blocking
  source_file_dependencies:
  - vllm/
  - tests/cuda
@@ -2206,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 'not skip_v1' 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
@@ -2240,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/
@@ -2257,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/
@@ -2303,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
@@ -2315,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
@@ -2330,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
@@ -2343,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/
@@ -2360,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
@@ -2404,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
@@ -2421,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/
@@ -2432,7 +2371,6 @@ steps:
  timeout_in_minutes: 20
  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi355_1
-  # grade: Blocking
  source_file_dependencies:
  - vllm/
  - tests/benchmarks/
@@ -2443,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
@@ -2464,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
@@ -2476,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/
@@ -2493,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/
@@ -2506,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/
@@ -2526,7 +2459,6 @@ steps:
  timeout_in_minutes: 45
  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_1
-  # grade: Blocking
  torch_nightly: true
  source_file_dependencies:
  - vllm/
@@ -2539,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:
@@ -2554,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/
@@ -2568,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/
@@ -2589,7 +2518,6 @@ steps:
  timeout_in_minutes: 75
  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_1
-  # grade: Blocking
  torch_nightly: true
  source_file_dependencies:
  - vllm/
@@ -2610,7 +2538,6 @@ steps:
  timeout_in_minutes: 110
  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_1
-  # grade: Blocking
  optional: true
  source_file_dependencies:
  - vllm/
@@ -2626,7 +2553,6 @@ steps:
  timeout_in_minutes: 110
  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_1
-  # grade: Blocking
  optional: true
  source_file_dependencies:
  - vllm/
@@ -2638,7 +2564,6 @@ steps:
  timeout_in_minutes: 50
  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_1
-  # grade: Blocking
  optional: true
  source_file_dependencies:
  - vllm/
@@ -2674,7 +2599,6 @@ steps:
  timeout_in_minutes: 60
  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_1
-  # grade: Blocking
  source_file_dependencies:
  - vllm/
  - tests/models/multimodal
@@ -2686,7 +2610,6 @@ steps:
  timeout_in_minutes: 100
  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_1
-  # grade: Blocking
  torch_nightly: true
  source_file_dependencies:
  - vllm/
@@ -2704,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/
@@ -2719,7 +2641,6 @@ steps:
  timeout_in_minutes: 120
  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_1
-  # grade: Blocking
  optional: true
  source_file_dependencies:
  - vllm/
@@ -2734,7 +2655,6 @@ steps:
  timeout_in_minutes: 120
  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_1
-  # grade: Blocking
  optional: true
  source_file_dependencies:
  - vllm/
@@ -2749,7 +2669,6 @@ steps:
  timeout_in_minutes: 150
  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_1
-  # grade: Blocking
  optional: true
  source_file_dependencies:
  - vllm/
@@ -2764,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:
@@ -2925,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:
@@ -2941,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
@@ -2968,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:
@@ -3008,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:
@@ -3030,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:
@@ -3064,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:
@@ -3081,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
@@ -3106,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
@@ -3119,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
@@ -3132,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
@@ -3146,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
@@ -3163,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
@@ -3186,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:
@@ -3202,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:
@@ -3217,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/"
@@ -3252,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
@@ -3262,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
@@ -3298,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
@@ -3329,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
@@ -3352,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"
@@ -165,6 +165,7 @@ steps:
  num_devices: 2
  num_nodes: 2
  no_plugin: true
+  optional: true # TODO: revert once infra issue solved
  source_file_dependencies:
  - vllm/distributed/
  - vllm/engine/
@@ -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
@@ -108,9 +108,11 @@ steps:
  timeout_in_minutes: 50
  source_file_dependencies:
  - vllm/
+  - tests/detokenizer
  - tests/multimodal
  - tests/utils_
  commands:
+  - pytest -v -s detokenizer
  - pytest -v -s -m 'not cpu_test' multimodal
  - pytest -v -s utils_

@@ -145,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 'not skip_v1' samplers
+      - pytest -v -s samplers
@@ -771,6 +771,25 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
    endif()
  endif()

+  # DeepSeek V3 fused A GEMM kernel (requires SM 9.0+, Hopper and later)
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
+    cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0f;11.0f" "${CUDA_ARCHS}")
+  else()
+    cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
+  endif()
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND DSV3_FUSED_A_GEMM_ARCHS)
+    set(DSV3_FUSED_A_GEMM_SRC "csrc/dsv3_fused_a_gemm.cu")
+    set_gencode_flags_for_srcs(
+      SRCS "${DSV3_FUSED_A_GEMM_SRC}"
+      CUDA_ARCHS "${DSV3_FUSED_A_GEMM_ARCHS}")
+    list(APPEND VLLM_EXT_SRC ${DSV3_FUSED_A_GEMM_SRC})
+    list(APPEND VLLM_GPU_FLAGS "-DENABLE_DSV3_FUSED_A_GEMM=1")
+    message(STATUS "Building dsv3_fused_a_gemm for archs: ${DSV3_FUSED_A_GEMM_ARCHS}")
+  else()
+    message(STATUS "Not building dsv3_fused_a_gemm as no compatible archs found "
+                   "in CUDA target architectures.")
+  endif()
+
  # moe_data.cu is used by all CUTLASS MoE kernels.
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
    cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0f;11.0f;12.0f" "${CUDA_ARCHS}")
@@ -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,
@@ -1305,7 +1305,8 @@ void indexer_k_quant_and_cache(
  const at::cuda::OptionalCUDAGuard device_guard(device_of(k));
  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();

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

@@ -0,0 +1,747 @@
+/*
+ * Adapted from
+ * https://github.com/sgl-project/sglang/blob/main/sgl-kernel/csrc/gemm/dsv3_fused_a_gemm.cu
+ * which was adapted from
+ * https://github.com/NVIDIA/TensorRT-LLM/blob/619709fc33bd5dc268f19d6a741fe7ed51c0f8f5/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3FusedAGemm.cu
+ *
+ * Copyright (c) 2019-2024, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2021, NAVER Corp.  Authored by CLOVA.
+ *
+ * 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 <torch/all.h>
+
+#include "core/registration.h"
+
+#include <cstdlib>
+#include <mutex>
+
+namespace {
+
+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) {
+      char const* env = std::getenv("TRTLLM_ENABLE_PDL");
+      enablePDL = env && env[0] == '1' && env[1] == '\0';
+    }
+  });
+  return enablePDL;
+}
+
+}  // namespace
+
+using bf16_t = __nv_bfloat16;
+
+__device__ void hmma_16_8_16_f32acc_bf16ab(float (&d_reg)[4],
+                                           const bf16_t (&a_reg)[8],
+                                           const bf16_t (&b_reg)[4],
+                                           float const (&c_reg)[4]) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+  uint32_t a0 = *reinterpret_cast<uint32_t const*>(a_reg + 0);
+  uint32_t a1 = *reinterpret_cast<uint32_t const*>(a_reg + 2);
+  uint32_t a2 = *reinterpret_cast<uint32_t const*>(a_reg + 4);
+  uint32_t a3 = *reinterpret_cast<uint32_t const*>(a_reg + 6);
+  uint32_t b0 = *reinterpret_cast<uint32_t const*>(b_reg + 0);
+  uint32_t b1 = *reinterpret_cast<uint32_t const*>(b_reg + 2);
+  asm volatile(
+      "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
+      "{%0,  %1,  %2,  %3},"
+      "{%4,  %5,  %6,  %7},"
+      "{%8,  %9},"
+      "{%10, %11, %12, %13};\n"
+      : "=f"(d_reg[0]), "=f"(d_reg[1]), "=f"(d_reg[2]), "=f"(d_reg[3])
+      : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1), "f"(d_reg[0]),
+        "f"(d_reg[1]), "f"(d_reg[2]), "f"(d_reg[3]));
+#endif
+}
+
+extern "C" {
+__device__ uint32_t __nvvm_get_smem_pointer(void*);
+}
+
+__device__ void ldgsts_128(void const* gPtr, void* sPtr, uint32_t pred) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+  if (pred) {
+    uint32_t smemPtrAsUint32 = __nvvm_get_smem_pointer(sPtr);
+    asm volatile("cp.async.cg.shared.global.L2::128B [%0], [%1], %2;\n" ::"r"(
+                     smemPtrAsUint32),
+                 "l"(gPtr), "n"(16));
+  }
+#endif
+}
+
+__device__ void ldsm_x4(void* smem_ptr, uint32_t* reg_ptr) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+  asm volatile(
+      "ldmatrix.sync.aligned.x4.m8n8.shared.b16 {%0, %1, %2, %3}, [%4];\n"
+      : "=r"(reg_ptr[0]), "=r"(reg_ptr[1]), "=r"(reg_ptr[2]), "=r"(reg_ptr[3])
+      : "r"(__nvvm_get_smem_pointer(smem_ptr)));
+#endif
+}
+
+template <class Type>
+__device__ int apply_swizzle_343_on_elem_row_col(int row_idx_, int col_idx_) {
+  uint32_t row_idx = *reinterpret_cast<uint32_t*>(&row_idx_);
+  uint32_t col_idx = *reinterpret_cast<uint32_t*>(&col_idx_);
+  row_idx = row_idx % 8;
+  row_idx = row_idx * (16 / sizeof(Type));
+  col_idx = col_idx ^ row_idx;
+  return *reinterpret_cast<int*>(&col_idx);
+}
+
+__device__ void initialize_barrier(
+    uint64_t* smem_barrier,  // 64 bits user-manged barrier in smem
+    int thread_count =
+        1)  // Thread count expected to arrive/wait on this barrier
+{
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+  uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_barrier);
+  asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(smem_int_ptr),
+               "r"(thread_count));
+#endif
+}
+
+// Barrier wait
+__device__ void wait_barrier(
+    uint64_t* smem_barrier,  // 64 bits user-manged barrier in smem
+    int phase_bit)           // Current phase bit the barrier waiting to flip
+{
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+  uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_barrier);
+  asm volatile(
+      "{\n"
+      ".reg .pred                P1;\n"
+      "LAB_WAIT:\n"
+      "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1;\n"
+      "@P1                       bra DONE;\n"
+      "bra                   LAB_WAIT;\n"
+      "DONE:\n"
+      "}\n" ::"r"(smem_int_ptr),
+      "r"(phase_bit));
+#endif
+}
+
+__device__ bool try_wait_barrier(uint64_t* smem_ptr, int phase_bit) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+  uint32_t wait_complete;
+  uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_ptr);
+  asm volatile(
+      "{\n\t"
+      ".reg .pred P1; \n\t"
+      "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2; \n\t"
+      "selp.b32 %0, 1, 0, P1; \n\t"
+      "}"
+      : "=r"(wait_complete)
+      : "r"(smem_int_ptr), "r"(phase_bit));
+  return static_cast<bool>(wait_complete);
+#endif
+  return false;
+}
+
+// Barrier arrive
+__device__ void arrive_barrier(
+    uint64_t* smem_barrier)  // 64 bits user-manged barrier in smem
+{
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+  uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_barrier);
+  asm volatile(
+      "{\n"
+      ".reg .b64 state; \n"
+      "mbarrier.arrive.shared::cta.b64   state, [%0];\n"
+      "}\n" ::"r"(smem_int_ptr));
+#endif
+}
+
+__device__ void ldgsts_arrive(uint64_t* smem_barrier) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+  uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_barrier);
+  asm volatile("cp.async.mbarrier.arrive.noinc.shared.b64 [%0];"
+               :
+               : "r"(smem_int_ptr));
+#endif
+}
+
+template <int gemm_k, int tile_m, int tile_k, int stage_cnt>
+struct GmemLoaderA {
+  static constexpr int elem_bytes = 2;
+  static constexpr int vec_bytes = 16;
+  static constexpr int vec_elems = vec_bytes / elem_bytes;
+  static constexpr int thread_cnt = 64;
+  static_assert((tile_m * tile_k) % (vec_elems * thread_cnt) == 0);
+  static constexpr int a_inst_cnt_per_iter =
+      (tile_m * tile_k) / (vec_elems * thread_cnt);
+  static_assert(gemm_k % tile_k == 0);
+  static constexpr int k_iter_cnt = gemm_k / tile_k;
+
+  // Extra params to keep the order of k reduction...
+  static constexpr int mma_warp_cnt = 4;
+  static constexpr int per_mma_warp_k = tile_k / mma_warp_cnt;
+  static constexpr int k_each_chunk = gemm_k / mma_warp_cnt;
+
+ private:
+  __device__ int k_project(int tile_k_idx) {
+    return (tile_k_idx / per_mma_warp_k * k_each_chunk) +
+           (tile_k_idx % per_mma_warp_k);
+  }
+
+ public:
+  __device__ GmemLoaderA(bf16_t const* gmem_a_local_, bf16_t* smem_a_,
+                         uint64_t* smem_barrier_)
+      : gmem_a(gmem_a_local_),
+        smem_a(smem_a_),
+        smem_barrier(smem_barrier_),
+        local_tid(threadIdx.x % thread_cnt) {}
+
+  __device__ void prepare() {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+  // swizzle, that's what we want.
+  #pragma unroll
+    for (int i = 0; i < a_inst_cnt_per_iter; i++) {
+      int linear_idx = local_tid * vec_elems + i * thread_cnt * vec_elems;
+      int m_idx = linear_idx / tile_k;
+      int k_idx = linear_idx % tile_k;
+      k_idx = apply_swizzle_343_on_elem_row_col<bf16_t>(m_idx, k_idx);
+      a_smem_offsets[i] = m_idx * tile_k + k_idx;
+    }
+#endif
+  }
+
+  __device__ void issue_mainloop() {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+  #pragma unroll 1
+    for (int loop_idx = 0; loop_idx < k_iter_cnt; loop_idx++) {
+      if (need_wait) {
+        wait_barrier(smem_barrier + 1 + stage_idx * 2, phase_bit);
+      }
+      int next_stage_idx = stage_idx + 1;
+      int next_phase_bit =
+          next_stage_idx == stage_cnt ? phase_bit ^ 1 : phase_bit;
+      next_stage_idx = next_stage_idx == stage_cnt ? 0 : next_stage_idx;
+      if (loop_idx != k_iter_cnt - 1) {
+        need_wait = !try_wait_barrier(smem_barrier + 1 + next_stage_idx * 2,
+                                      next_phase_bit);
+      }
+
+  #pragma unroll
+      for (int i = 0; i < a_inst_cnt_per_iter; i++) {
+        int smem_offset = a_smem_offsets[i];
+        bf16_t* smem_ptr_this_iter =
+            smem_a + stage_idx * tile_m * tile_k + smem_offset;
+        int linear_idx = local_tid * vec_elems + i * thread_cnt * vec_elems;
+        int m_idx = linear_idx / tile_k;
+        int k_idx = linear_idx % tile_k;
+        int gmem_offset = m_idx * gemm_k + k_project(k_idx);
+        bf16_t const* gmem_ptr_this_iter = gmem_a + gmem_offset;
+        ldgsts_128(gmem_ptr_this_iter, smem_ptr_this_iter, true);
+      }
+      ldgsts_arrive(smem_barrier + stage_idx * 2);
+
+      stage_idx = next_stage_idx;
+      phase_bit = next_phase_bit;
+      gmem_a += per_mma_warp_k;
+    }
+#endif
+  }
+
+  bf16_t const* gmem_a;
+  bf16_t* smem_a;
+  uint64_t* smem_barrier;
+  int local_tid;
+  int stage_idx = 0;
+  int phase_bit = 1;
+  bool need_wait = true;
+
+  // per smem_stage, store with swizzle information
+  int a_smem_offsets[a_inst_cnt_per_iter];
+};
+
+template <int gemm_k, int tile_n, int tile_k, int stage_cnt>
+struct GmemLoaderB {
+  static constexpr int elem_bytes = 2;
+  static constexpr int vec_bytes = 16;
+  static constexpr int vec_elems = vec_bytes / elem_bytes;
+  static constexpr int thread_cnt = 64;
+  static_assert((tile_n * tile_k) % (vec_elems * thread_cnt) == 0);
+  static constexpr int b_inst_cnt_per_iter =
+      (tile_n * tile_k) / (vec_elems * thread_cnt);
+  static_assert(gemm_k % tile_k == 0);
+  static constexpr int k_iter_cnt = gemm_k / tile_k;
+
+  // Extra params to keep the order of k reduction...
+  static constexpr int mma_warp_cnt = 4;
+  static constexpr int per_mma_warp_k = tile_k / mma_warp_cnt;
+  static constexpr int k_each_chunk = gemm_k / mma_warp_cnt;
+
+ private:
+  __device__ int k_project(int tile_k_idx) {
+    return (tile_k_idx / per_mma_warp_k * k_each_chunk) +
+           (tile_k_idx % per_mma_warp_k);
+  }
+
+ public:
+  __device__ GmemLoaderB(bf16_t const* gmem_b_local_, bf16_t* smem_b_,
+                         uint64_t* smem_barrier_, int gemm_n_)
+      : gmem_b(gmem_b_local_),
+        smem_b(smem_b_),
+        smem_barrier(smem_barrier_),
+        gemm_n(gemm_n_),
+        local_tid(threadIdx.x % thread_cnt) {}
+
+  __device__ void prepare() {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+  // swizzle, that's what we want.
+  #pragma unroll
+    for (int i = 0; i < b_inst_cnt_per_iter; i++) {
+      int linear_idx = local_tid * vec_elems + i * thread_cnt * vec_elems;
+      int n_idx = linear_idx / tile_k;
+      int k_idx = linear_idx % tile_k;
+      k_idx = apply_swizzle_343_on_elem_row_col<bf16_t>(n_idx, k_idx);
+      b_smem_offsets[i] = n_idx * tile_k + k_idx;
+      preds[i] = n_idx < gemm_n;
+    }
+#endif
+  }
+
+  __device__ void issue_mainloop() {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+    asm volatile("griddepcontrol.wait;");
+  #pragma unroll 1
+    for (int loop_idx = 0; loop_idx < k_iter_cnt; loop_idx++) {
+      if (need_wait) {
+        wait_barrier(smem_barrier + 1 + stage_idx * 2, phase_bit);
+      }
+      int next_stage_idx = stage_idx + 1;
+      int next_phase_bit =
+          next_stage_idx == stage_cnt ? phase_bit ^ 1 : phase_bit;
+      next_stage_idx = next_stage_idx == stage_cnt ? 0 : next_stage_idx;
+      if (loop_idx != k_iter_cnt - 1) {
+        need_wait = !try_wait_barrier(smem_barrier + 1 + next_stage_idx * 2,
+                                      next_phase_bit);
+      }
+  #pragma unroll
+      for (int i = 0; i < b_inst_cnt_per_iter; i++) {
+        int smem_offset = b_smem_offsets[i];
+        bf16_t* smem_ptr_this_iter =
+            smem_b + stage_idx * tile_n * tile_k + smem_offset;
+        int linear_idx = local_tid * vec_elems + i * thread_cnt * vec_elems;
+        int n_idx = linear_idx / tile_k;
+        int k_idx = linear_idx % tile_k;
+        int gmem_offset = n_idx * gemm_k + k_project(k_idx);
+        bf16_t const* gmem_ptr_this_iter = gmem_b + gmem_offset;
+        ldgsts_128(gmem_ptr_this_iter, smem_ptr_this_iter, preds[i]);
+      }
+      ldgsts_arrive(smem_barrier + stage_idx * 2);
+
+      stage_idx = next_stage_idx;
+      phase_bit = next_phase_bit;
+      gmem_b += per_mma_warp_k;
+    }
+#endif
+  }
+
+  bf16_t const* gmem_b;
+  bf16_t* smem_b;
+  uint64_t* smem_barrier;
+  int gemm_n;
+  int local_tid;
+  int stage_idx = 0;
+  int phase_bit = 1;
+  bool need_wait = true;
+
+  // per smem_stage, store with swizzle information
+  int b_smem_offsets[b_inst_cnt_per_iter];
+  uint32_t preds[b_inst_cnt_per_iter];
+};
+
+template <int gemm_m, int gemm_k, int tile_m, int tile_n, int tile_k,
+          int stage_cnt>
+struct MmaComputer {
+  static constexpr int elem_bytes = 2;
+  static constexpr int thread_cnt = 128;
+  static_assert(gemm_k % tile_k == 0);
+  static_assert(tile_k % (thread_cnt / 32) == 0);
+  static constexpr int per_warp_tile_k = tile_k / (thread_cnt / 32);
+  static constexpr int k_iter_cnt = gemm_k / tile_k;
+  static constexpr int k_phase_cnt = per_warp_tile_k / 16;
+  static constexpr int m_iter_cnt = (tile_m + 15) / 16;
+  static constexpr int n_iter_cnt =
+      (tile_n + 7) /
+      8;  // Possible to have non-1 n_iter_cnt for ab_swap m16 case.
+  static_assert(m_iter_cnt == 1);
+  static_assert(n_iter_cnt == 1 || n_iter_cnt == 2);
+
+  __device__ MmaComputer(bf16_t* gmem_c_local_, bf16_t* smem_a_,
+                         bf16_t* smem_b_, uint64_t* smem_barrier_,
+                         int warp_idx_, int gemm_n_)
+      : gmem_c(gmem_c_local_),
+        smem_a(smem_a_),
+        smem_b(smem_b_),
+        smem_barrier(smem_barrier_),
+        warp_idx(warp_idx_ - (thread_cnt / 32)),
+        gemm_n(gemm_n_) {}
+
+ private:
+  __device__ constexpr int internal_b_atom_func(int tid) {
+    if constexpr (tile_n < 8) {
+      return (tid % tile_n) + ((tid % 8) / tile_n * 0) + tid / 8 * 8 * tile_n;
+    } else {
+      return (tid % 8) + ((tid % 32) / 8 * (tile_n * 8));
+    }
+  }
+
+ public:
+  __device__ void prepare() {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+  #pragma unroll
+    for (int i = 0; i < k_phase_cnt; i++) {
+      int linear_idx = (lane_idx % 16) + (lane_idx / 16) * 128 + i * 256;
+      int m_idx = linear_idx % tile_m;
+      int k_idx = linear_idx / tile_m + warp_k_offset_in_tile_k;
+      k_idx = apply_swizzle_343_on_elem_row_col<bf16_t>(m_idx, k_idx);
+      a_smem_offsets[0][i] = m_idx * tile_k + k_idx;
+    }
+  #pragma unroll
+    for (int n_iter_idx = 0; n_iter_idx < n_iter_cnt; n_iter_idx++) {
+  #pragma unroll
+      for (int i = 0; i < k_phase_cnt; i += 2) {  // Special i+=2 for B.
+        int linear_idx =
+            internal_b_atom_func(lane_idx) + i * tile_n * 16 + n_iter_idx * 8;
+        int n_idx = linear_idx % tile_n;
+        int k_idx = linear_idx / tile_n + warp_k_offset_in_tile_k;
+        k_idx = apply_swizzle_343_on_elem_row_col<bf16_t>(n_idx, k_idx);
+        b_smem_offsets[n_iter_idx][i] = n_idx * tile_k + k_idx;
+      }
+    }
+#endif
+  }
+
+  __device__ void issue_mainloop() {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+  #pragma unroll 1
+    for (int loop_idx = 0; loop_idx < k_iter_cnt; loop_idx++) {
+      wait_barrier(smem_barrier + 0 + stage_idx * 2, phase_bit);
+
+  #pragma unroll
+      for (int i = 0; i < k_phase_cnt; i++) {
+        int smem_offset = a_smem_offsets[0][i];
+        bf16_t* smem_ptr_this_iter =
+            smem_a + stage_idx * tile_m * tile_k + smem_offset;
+        ldsm_x4(smem_ptr_this_iter, reinterpret_cast<uint32_t*>(a_reg[0][i]));
+      }
+
+  #pragma unroll
+      for (int n_iter_idx = 0; n_iter_idx < n_iter_cnt; n_iter_idx++) {
+  #pragma unroll
+        for (int i = 0; i < k_phase_cnt; i += 2) {
+          int smem_offset = b_smem_offsets[n_iter_idx][i];
+          bf16_t* smem_ptr_this_iter =
+              smem_b + stage_idx * tile_n * tile_k + smem_offset;
+          ldsm_x4(smem_ptr_this_iter,
+                  reinterpret_cast<uint32_t*>(b_reg[n_iter_idx][i]));
+        }
+      }
+
+  #pragma unroll
+      for (int k_iter_idx = 0; k_iter_idx < k_phase_cnt; k_iter_idx++) {
+  #pragma unroll
+        for (int n_iter_idx = 0; n_iter_idx < n_iter_cnt; n_iter_idx++) {
+          hmma_16_8_16_f32acc_bf16ab(
+              acc_reg[0][n_iter_idx], a_reg[0][k_iter_idx],
+              b_reg[n_iter_idx][k_iter_idx], acc_reg[0][n_iter_idx]);
+        }
+      }
+      ::arrive_barrier(smem_barrier + 1 + stage_idx * 2);
+      stage_idx += 1;
+      phase_bit = stage_idx == stage_cnt ? phase_bit ^ 1 : phase_bit;
+      stage_idx = stage_idx == stage_cnt ? 0 : stage_idx;
+    }
+#endif
+  }
+
+  __device__ void epi() {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+    asm volatile("bar.sync %0, %1;" : : "r"(1), "r"(thread_cnt));
+    // reorganize the acc_reg
+    constexpr int thread_m = 2;
+    constexpr int thread_n = 2 * n_iter_cnt;
+    constexpr int cta_mma_n = n_iter_cnt * 8;
+    float acc_reg_reorg[thread_m][thread_n];
+
+    for (int i = 0; i < thread_m; i++) {
+      for (int j = 0; j < thread_n; j++) {
+        acc_reg_reorg[i][j] = acc_reg[0][j / 2][(j % 2) + (i * 2)];
+      }
+    }
+
+    // 4 x cosize(smem_c_layout)
+    float* smem_c = reinterpret_cast<float*>(smem_a);
+    // coord -> index
+    auto smem_c_index_func = [&](int m_idx, int n_idx) {
+      int group_rows = 32 / cta_mma_n;
+      int group_cnt = 2;
+      return (m_idx % group_rows * cta_mma_n) +
+             (m_idx / group_rows * (32 + group_cnt)) + n_idx;
+    };
+    constexpr int cosize_smem_c = ((tile_m * cta_mma_n) / 32) * (32 + 2);
+
+  // This should be optimized to STS.64 but can not be STS.128 due to the bank
+  // index.
+  #pragma unroll
+    for (int m_idx_thread = 0; m_idx_thread < thread_m; m_idx_thread++) {
+  #pragma unroll
+      for (int n_idx_thread = 0; n_idx_thread < thread_n; n_idx_thread++) {
+        int m_idx = (lane_idx / 4) + m_idx_thread * 8;
+        int n_idx =
+            ((lane_idx % 4) * 2) + (n_idx_thread % 2) + (n_idx_thread / 2) * 8;
+        smem_c[cosize_smem_c * warp_idx + smem_c_index_func(m_idx, n_idx)] =
+            acc_reg_reorg[m_idx_thread][n_idx_thread];
+      }
+    }
+    asm volatile("bar.sync %0, %1;" : : "r"(1), "r"(thread_cnt));
+
+    if (warp_idx == 0) {
+      constexpr int final_acc_reg_cnt = (tile_m * tile_n + 31) / 32;
+      float acc_final[final_acc_reg_cnt]{};
+
+  #pragma unroll
+      for (int reg_idx = 0; reg_idx < final_acc_reg_cnt; reg_idx++) {
+        int linear_idx = reg_idx * 32 + lane_idx;
+        int m_idx = linear_idx % tile_m;
+        int n_idx = linear_idx / tile_m;
+        acc_final[reg_idx] +=
+            smem_c[smem_c_index_func(m_idx, n_idx) + 0 * cosize_smem_c] +
+            smem_c[smem_c_index_func(m_idx, n_idx) + 1 * cosize_smem_c] +
+            smem_c[smem_c_index_func(m_idx, n_idx) + 2 * cosize_smem_c] +
+            smem_c[smem_c_index_func(m_idx, n_idx) + 3 * cosize_smem_c];
+      }
+
+  #pragma unroll
+      for (int reg_idx = 0; reg_idx < final_acc_reg_cnt; reg_idx++) {
+        int linear_idx = reg_idx * 32 + lane_idx;
+        int m_idx = linear_idx % tile_m;
+        int n_idx = linear_idx / tile_m;
+        if (m_idx < tile_m && n_idx < gemm_n) {
+          gmem_c[n_idx * gemm_m + m_idx] = acc_final[reg_idx];
+        }
+      }
+    }
+#endif
+  }
+
+  bf16_t* gmem_c;
+  bf16_t* smem_a;
+  bf16_t* smem_b;
+  uint64_t* smem_barrier;
+  int warp_idx;
+  int gemm_n;
+  int stage_idx = 0;
+  int phase_bit = 0;
+  int lane_idx = threadIdx.x % 32;
+  int warp_k_offset_in_tile_k = warp_idx * per_warp_tile_k;
+
+  int a_smem_offsets[m_iter_cnt][k_phase_cnt];
+  int b_smem_offsets[n_iter_cnt][k_phase_cnt];
+
+  bf16_t a_reg[m_iter_cnt][k_phase_cnt][8];
+  bf16_t b_reg[n_iter_cnt][k_phase_cnt][4];
+  float acc_reg[m_iter_cnt][n_iter_cnt][4]{};
+};
+
+// AB swapped, kernel is k-major, k-major, m-major
+template <int batch_size, int gemm_m, int gemm_k, int tile_m, int tile_n,
+          int tile_k, int stage_cnt>
+__global__ __launch_bounds__(256, 1) void fused_a_gemm_kernel(
+    bf16_t* output, bf16_t const* mat_a, bf16_t const* mat_b, int gemm_n) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+  constexpr int load_thread_cnt = 128;
+  constexpr int compute_thread_cnt = 128;
+  constexpr int thread_cnt = load_thread_cnt + compute_thread_cnt;
+  (void)thread_cnt;
+  static_assert(gemm_m % 16 == 0);
+  static_assert(gemm_k % tile_k == 0);
+  static_assert(gemm_m % tile_m == 0);
+  static_assert(
+      tile_k == 128 || tile_k == 256 || tile_k == 512 ||
+      tile_k == 1024);  // tile_k must be larger than 64 since 4 warp splitK.
+  static_assert(tile_m == 16);
+  constexpr int g2s_vec_bytes = 16;
+  constexpr int a_elem_bytes = 2;
+  constexpr int b_elem_bytes = 2;
+  static_assert((tile_m * a_elem_bytes + tile_n * b_elem_bytes) * tile_k *
+                    stage_cnt <=
+                225 * 1024);
+  static_assert((tile_m * tile_k * a_elem_bytes) %
+                    (load_thread_cnt * g2s_vec_bytes) ==
+                0);
+  static_assert((tile_n * tile_k * b_elem_bytes) %
+                    (load_thread_cnt * g2s_vec_bytes) ==
+                0);
+
+  extern __shared__ char smem[];
+  uint64_t* smem_barrier = reinterpret_cast<uint64_t*>(
+      smem);  // producer,consumer; producer,consumer; ...
+  bf16_t* smem_a = reinterpret_cast<bf16_t*>(smem + (stage_cnt * 8 * 2 + 1024) /
+                                                        1024 * 1024);
+  bf16_t* smem_b = smem_a + tile_m * tile_k * stage_cnt;
+
+  int cta_m_idx = tile_m * blockIdx.x;
+  int cta_n_idx = tile_n * blockIdx.y;
+  bf16_t const* gmem_a_local = mat_a + cta_m_idx * gemm_k;
+  bf16_t const* gmem_b_local = mat_b + cta_n_idx * gemm_k;
+  bf16_t* gmem_c_local = output + cta_n_idx * gemm_m + cta_m_idx;
+
+  int warp_idx = __shfl_sync(0xffffffff, threadIdx.x / 32, 0);
+
+  if (warp_idx == 4) {
+    for (int i = 0; i < stage_cnt; i++) {
+      initialize_barrier(smem_barrier + i * 2 + 0,
+                         load_thread_cnt);  // producer
+      initialize_barrier(smem_barrier + i * 2 + 1,
+                         compute_thread_cnt);  // consumer
+    }
+  }
+  __syncthreads();
+
+  if (warp_idx < 2) {
+    GmemLoaderA<gemm_k, tile_m, tile_k, stage_cnt> a_loader(
+        gmem_a_local, smem_a, smem_barrier);
+    a_loader.prepare();
+    a_loader.issue_mainloop();
+  } else if (warp_idx < 4) {
+    GmemLoaderB<gemm_k, tile_n, tile_k, stage_cnt> b_loader(
+        gmem_b_local, smem_b, smem_barrier, gemm_n);
+    b_loader.prepare();
+    b_loader.issue_mainloop();
+  } else {
+    MmaComputer<gemm_m, gemm_k, tile_m, tile_n, tile_k, stage_cnt> mma_computer(
+        gmem_c_local, smem_a, smem_b, smem_barrier, warp_idx, gemm_n);
+    mma_computer.prepare();
+    mma_computer.issue_mainloop();
+    mma_computer.epi();
+  }
+  asm volatile("griddepcontrol.launch_dependents;");
+#endif
+}
+
+template <typename T, int kHdIn, int kHdOut, int kTileN>
+void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens,
+                      cudaStream_t const stream) {
+  constexpr int gemm_m = kHdOut;  // 2112
+  int const gemm_n = num_tokens;  // 1-16
+  constexpr int gemm_k = kHdIn;   // 7168
+  constexpr int batch_size = 1;
+  std::swap(mat_a, mat_b);
+  constexpr int tile_m = 16;
+  constexpr int tile_n = kTileN;                        // 8 or 16
+  constexpr int tile_k = std::max(256, 1024 / tile_n);  // 256
+  constexpr int max_stage_cnt =
+      1024 * 192 / ((tile_m + tile_n) * tile_k * sizeof(bf16_t));
+  constexpr int k_iter_cnt = gemm_k / tile_k;
+  constexpr int stage_cnt =
+      k_iter_cnt > max_stage_cnt ? max_stage_cnt : k_iter_cnt;
+  int cta_m_cnt = gemm_m / tile_m;
+  int cta_n_cnt = (gemm_n + tile_n - 1) / tile_n;
+  constexpr int barrier_bytes = (stage_cnt * 16 + 1023) / 1024 * 1024;
+  constexpr int smem_bytes =
+      ((tile_m * 2 + tile_n * 2) * tile_k * stage_cnt + barrier_bytes + 1023) /
+      1024 * 1024;
+
+  dim3 grid(cta_m_cnt, cta_n_cnt, 1);
+  dim3 block_size(256);
+  cudaLaunchConfig_t config;
+  config.gridDim = grid;
+  config.blockDim = block_size;
+  config.dynamicSmemBytes = smem_bytes;
+  config.stream = stream;
+  cudaLaunchAttribute attrs[1];
+  attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+  attrs[0].val.programmaticStreamSerializationAllowed = getEnvEnablePDL();
+  config.numAttrs = 1;
+  config.attrs = attrs;
+  if (smem_bytes >= (48 * 1024)) {
+    cudaFuncSetAttribute(fused_a_gemm_kernel<batch_size, gemm_m, gemm_k, tile_m,
+                                             tile_n, tile_k, stage_cnt>,
+                         cudaFuncAttributeMaxDynamicSharedMemorySize,
+                         smem_bytes);
+  }
+  cudaLaunchKernelEx(&config,
+                     fused_a_gemm_kernel<batch_size, gemm_m, gemm_k, tile_m,
+                                         tile_n, tile_k, stage_cnt>,
+                     output, mat_a, mat_b, gemm_n);
+}
+
+template void invokeFusedAGemm<__nv_bfloat16, 7168, 2112, 8>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens,
+    cudaStream_t);
+
+template void invokeFusedAGemm<__nv_bfloat16, 7168, 2112, 16>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens,
+    cudaStream_t);
+
+void dsv3_fused_a_gemm(torch::Tensor& output, torch::Tensor const& mat_a,
+                       torch::Tensor const& mat_b) {
+  TORCH_CHECK(mat_a.dim() == 2 && mat_b.dim() == 2 && output.dim() == 2);
+  int const num_tokens = mat_a.size(0);
+  int const hd_in = mat_a.size(1);
+  int const hd_out = mat_b.size(1);
+
+  constexpr int kHdIn = 7168;
+  constexpr int kHdOut = 2112;
+  TORCH_CHECK(num_tokens >= 1 && num_tokens <= 16,
+              "required 1 <= mat_a.shape[0] <= 16")
+  TORCH_CHECK(hd_in == kHdIn, "required mat_a.shape[1] == 7168")
+  TORCH_CHECK(hd_out == kHdOut, "required mat_b.shape[1] == 2112")
+  TORCH_CHECK(output.size(0) == num_tokens,
+              "required output.shape[0] == mat_a.shape[0]")
+  TORCH_CHECK(output.size(1) == hd_out,
+              "required output.shape[1] == mat_b.shape[1]")
+
+  TORCH_CHECK(mat_a.stride(1) == 1, "mat_a must be a row major tensor");
+  TORCH_CHECK(output.stride(1) == 1, "output must be a row major tensor");
+  TORCH_CHECK(mat_b.stride(0) == 1, "mat_b must be a column major tensor");
+
+  TORCH_CHECK(mat_a.scalar_type() == torch::kBFloat16 &&
+                  mat_b.scalar_type() == torch::kBFloat16,
+              "Only BFloat16 input dtype is supported")
+  TORCH_CHECK(output.scalar_type() == torch::kBFloat16,
+              "Only BFloat16 output dtype is supported")
+
+  TORCH_CHECK(getSMVersion() >= 90, "required CUDA ARCH >= SM_90");
+
+  auto stream = at::cuda::getCurrentCUDAStream(mat_a.get_device());
+  if (num_tokens <= 8) {
+    invokeFusedAGemm<__nv_bfloat16, kHdIn, kHdOut, 8>(
+        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()), num_tokens,
+        stream);
+  } else {
+    invokeFusedAGemm<__nv_bfloat16, kHdIn, kHdOut, 16>(
+        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()), num_tokens,
+        stream);
+  }
+}
@@ -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
@@ -315,7 +315,9 @@ void silu_and_mul_scaled_fp4_experts_quant(
 void per_token_group_quant_fp8(const torch::Tensor& input,
                               torch::Tensor& output_q, torch::Tensor& output_s,
                               int64_t group_size, double eps, double fp8_min,
-                               double fp8_max, bool scale_ue8m0);
+                               double fp8_max, bool scale_ue8m0,
+                               bool dummy_is_scale_transposed,
+                               bool dummy_is_tma_aligned);

 void per_token_group_quant_int8(const torch::Tensor& input,
                                torch::Tensor& output_q,
@@ -408,3 +410,8 @@ void qr_all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out,
                   int64_t quant_level, bool cast_bf2half = false);
 int64_t qr_max_size();
 #endif
+
+#ifndef USE_ROCM
+void dsv3_fused_a_gemm(torch::Tensor& output, torch::Tensor const& mat_a,
+                       torch::Tensor const& mat_b);
+#endif
@@ -97,7 +97,7 @@ __global__ void rms_norm_per_block_quant_kernel(
    scalar_t const* __restrict__ input,   // [..., hidden_size]
    scalar_t const* __restrict__ weight,  // [hidden_size]
    float const* scale_ub, float const var_epsilon, int32_t const hidden_size,
-    scalar_t* __restrict__ residual = nullptr) {
+    scalar_t* __restrict__ residual = nullptr, int64_t outer_scale_stride = 1) {
  float rms;
  // Compute RMS
  // Always able to vectorize due to constraints on hidden_size
@@ -108,7 +108,8 @@ __global__ void rms_norm_per_block_quant_kernel(
  // Always able to vectorize due to constraints on hidden_size and group_size
  vllm::vectorized::compute_dynamic_per_token_scales<
      scalar_t, scalar_out_t, has_residual, is_scale_transposed, group_size>(
-      nullptr, scales, input, weight, rms, scale_ub, hidden_size, residual);
+      nullptr, scales, input, weight, rms, scale_ub, hidden_size, residual,
+      outer_scale_stride);

  // RMS Norm + Quant
  // Always able to vectorize due to constraints on hidden_size
@@ -119,7 +120,8 @@ __global__ void rms_norm_per_block_quant_kernel(
  vllm::vectorized::norm_and_quant<
      scalar_t, scalar_out_t, std::is_same_v<scalar_out_t, int8_t>,
      has_residual, is_scale_transposed, group_size>(
-      out, input, weight, rms, scales, hidden_size, residual);
+      out, input, weight, rms, scales, hidden_size, residual,
+      outer_scale_stride);
 }

 }  // namespace vllm
@@ -225,7 +227,8 @@ void rms_norm_per_block_quant_dispatch(
                                                 : nullptr,
                            var_epsilon, hidden_size,
                            has_residual ? residual->data_ptr<scalar_in_t>()
-                                         : nullptr);
+                                         : nullptr,
+                            scales.stride(1));
                  });
            });
          });
@@ -257,6 +260,11 @@ void rms_norm_per_block_quant(torch::Tensor& out, torch::Tensor const& input,
  TORCH_CHECK(group_size == 128 || group_size == 64,
              "Unsupported group size: ", group_size);

+  if (scales.stride(1) > 1) {
+    TORCH_CHECK(is_scale_transposed,
+                "Outer scale stride must be 1 when scales are not transposed");
+  }
+
  rms_norm_per_block_quant_dispatch(out, input, weight, scales, group_size,
                                    var_epsilon, scale_ub, residual,
                                    is_scale_transposed);
@@ -74,7 +74,7 @@ __device__ void compute_dynamic_per_token_scales(
    scalar_t const* __restrict__ input, scalar_t const* __restrict__ weight,
    float const rms, float const* __restrict__ scale_ub,
    int32_t const hidden_size, scalar_t const* __restrict__ residual = nullptr,
-    int32_t const group_size = 0) {
+    int32_t const group_size = 0, int64_t outer_scale_stride = 1) {
  float block_absmax_val_maybe = 0.0f;
  constexpr scalar_out_t qmax{quant_type_max_v<scalar_out_t>};
  __syncthreads();
@@ -133,7 +133,9 @@ __device__ void compute_dynamic_per_token_scales(
      scale = max(scale / qmax, min_scaling_factor<scalar_out_t>::val());
      // Global output store
      if constexpr (is_scale_transposed) {
-        all_token_scales[(threadIdx.x / threads_per_group) * gridDim.x +
+        int64_t const scale_rows = (gridDim.x + outer_scale_stride - 1) /
+                                   outer_scale_stride * outer_scale_stride;
+        all_token_scales[(threadIdx.x / threads_per_group) * scale_rows +
                         blockIdx.x] = scale;
      } else {
        all_token_scales[blockIdx.x * num_groups +
@@ -180,13 +182,11 @@ __device__ void compute_dynamic_per_token_scales(

 template <typename scalar_t, typename scalar_out_t, bool is_scale_inverted,
          bool has_residual = false, bool is_scale_transposed = false>
-__device__ void norm_and_quant(scalar_out_t* __restrict__ output,
-                               scalar_t const* __restrict__ input,
-                               scalar_t const* __restrict__ weight,
-                               float const rms, float* const scale,
-                               int32_t const hidden_size,
-                               scalar_t* __restrict__ residual = nullptr,
-                               int32_t const group_size = 0) {
+__device__ void norm_and_quant(
+    scalar_out_t* __restrict__ output, scalar_t const* __restrict__ input,
+    scalar_t const* __restrict__ weight, float const rms, float* const scale,
+    int32_t const hidden_size, scalar_t* __restrict__ residual = nullptr,
+    int32_t const group_size = 0, int64_t outer_scale_stride = 1) {
  int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);

  for (auto i = threadIdx.x; i < hidden_size; i += blockDim.x) {
@@ -202,7 +202,9 @@ __device__ void norm_and_quant(scalar_out_t* __restrict__ output,
    int64_t scale_idx = 0;
    if (group_size > 0) {
      if constexpr (is_scale_transposed) {
-        scale_idx = (i / group_size) * gridDim.x + blockIdx.x;
+        int64_t const scale_rows = (gridDim.x + outer_scale_stride - 1) /
+                                   outer_scale_stride * outer_scale_stride;
+        scale_idx = (i / group_size) * scale_rows + blockIdx.x;
      } else {
        scale_idx = blockIdx.x * (hidden_size / group_size) + i / group_size;
      }
@@ -286,8 +288,8 @@ __device__ void compute_dynamic_per_token_scales(
    float* __restrict__ token_scale, float* __restrict__ all_token_scales,
    scalar_t const* __restrict__ input, scalar_t const* __restrict__ weight,
    float const rms, float const* __restrict__ scale_ub,
-    int32_t const hidden_size,
-    scalar_t const* __restrict__ residual = nullptr) {
+    int32_t const hidden_size, scalar_t const* __restrict__ residual = nullptr,
+    int64_t outer_scale_stride = 1) {
  constexpr scalar_out_t qmax{quant_type_max_v<scalar_out_t>};

  const int VEC_SIZE = 4;
@@ -382,7 +384,9 @@ __device__ void compute_dynamic_per_token_scales(
      scale = max(scale / qmax, min_scaling_factor<scalar_out_t>::val());
      // Global output store
      if constexpr (is_scale_transposed) {
-        all_token_scales[(threadIdx.x / threads_per_group) * gridDim.x +
+        int64_t const scale_rows = (gridDim.x + outer_scale_stride - 1) /
+                                   outer_scale_stride * outer_scale_stride;
+        all_token_scales[(threadIdx.x / threads_per_group) * scale_rows +
                         blockIdx.x] = scale;
      } else {
        all_token_scales[blockIdx.x * num_groups +
@@ -463,7 +467,8 @@ __device__ void norm_and_quant(scalar_out_t* __restrict__ output,
                               scalar_t const* __restrict__ weight,
                               float const rms, float* const scale,
                               int32_t const hidden_size,
-                               scalar_t* __restrict__ residual = nullptr) {
+                               scalar_t* __restrict__ residual = nullptr,
+                               int64_t outer_scale_stride = 1) {
  int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);

  // Vectorized input/output/weight/residual to better utilize memory bandwidth.
@@ -516,7 +521,9 @@ __device__ void norm_and_quant(scalar_out_t* __restrict__ output,
      int64_t const num_groups = hidden_size / group_size;
      int64_t scale_idx = 0;
      if constexpr (is_scale_transposed) {
-        scale_idx = (i * VEC_SIZE / group_size) * gridDim.x + blockIdx.x;
+        int64_t const scale_rows = (gridDim.x + outer_scale_stride - 1) /
+                                   outer_scale_stride * outer_scale_stride;
+        scale_idx = (i * VEC_SIZE / group_size) * scale_rows + blockIdx.x;
      } else {
        scale_idx = blockIdx.x * num_groups + i * VEC_SIZE / group_size;
      }
@@ -379,7 +379,9 @@ void per_token_group_quant_8bit_packed(const torch::Tensor& input,
 void per_token_group_quant_fp8(const torch::Tensor& input,
                               torch::Tensor& output_q, torch::Tensor& output_s,
                               int64_t group_size, double eps, double fp8_min,
-                               double fp8_max, bool scale_ue8m0) {
+                               double fp8_max, bool scale_ue8m0,
+                               bool dummy_is_scale_transposed = false,
+                               bool dummy_is_tma_aligned = false) {
  per_token_group_quant_8bit(input, output_q, output_s, group_size, eps,
                             fp8_min, fp8_max, scale_ue8m0);
 }
@@ -239,6 +239,11 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {

  // Quantization ops
 #ifndef USE_ROCM
+  // DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens).
+  ops.def(
+      "dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
+  ops.impl("dsv3_fused_a_gemm", torch::kCUDA, &dsv3_fused_a_gemm);
+
  // Quantized GEMM for AWQ.
  ops.def(
      "awq_gemm(Tensor _in_feats, Tensor _kernel, Tensor _scaling_factors, "
@@ -643,11 +648,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {

 #ifndef USE_ROCM
  // Compute per-token-group FP8 quantized tensor and scaling factor.
+  // The dummy arguments are here so we can correctly fuse with RMSNorm.
  ops.def(
      "per_token_group_fp8_quant(Tensor input, Tensor! output_q, Tensor! "
      "output_s, "
      "int group_size, float eps, float fp8_min, float fp8_max, bool "
-      "scale_ue8m0) -> ()");
+      "scale_ue8m0, bool dummy_is_scale_transposed, bool dummy_is_tma_aligned "
+      ") -> ()");
  ops.impl("per_token_group_fp8_quant", torch::kCUDA,
           &per_token_group_quant_fp8);

@@ -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"
@@ -155,3 +155,4 @@ The interface for the model/module may change during vLLM's development. If you
    - `use_v1` parameter in `Platform.get_attn_backend_cls` is deprecated. It has been removed in v0.13.0.
    - `_Backend` in `vllm.attention` is deprecated. It has been removed in v0.13.0. Please use `vllm.v1.attention.backends.registry.register_backend` to add new attention backend to `AttentionBackendEnum` instead.
    - `seed_everything` platform interface is deprecated. It has been removed in v0.16.0. Please use `vllm.utils.torch_utils.set_random_seed` instead.
+    - `prompt` in `Platform.validate_request` is deprecated and will be removed in v0.18.0.
@@ -36,12 +36,12 @@ th:not(:first-child) {
 }
 </style>

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

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

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

 markdown_extensions:
  - attr_list
@@ -113,7 +113,6 @@ markers = [
    "cpu_test: mark test as CPU-only test",
    "split: run this test as part of a split",
    "distributed: run this test only in distributed GPU tests",
-    "skip_v1: do not run this test with v1",
    "optional: optional tests that are automatically skipped, include --optional to run them",
 ]

@@ -10,4 +10,4 @@ torchaudio==2.10.0
 # These must be updated alongside torch
 torchvision==0.25.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
 # FlashInfer should be updated together with the Dockerfile
-flashinfer-python==0.6.3
+flashinfer-python==0.6.4
@@ -1,4 +1,4 @@
-mkdocs
+mkdocs<2.0.0
 mkdocs-api-autonav
 mkdocs-material
 mkdocstrings-python
@@ -7,6 +7,7 @@ mkdocs-awesome-nav
 mkdocs-glightbox
 mkdocs-git-revision-date-localized-plugin
 mkdocs-minify-plugin
+mkdocs-redirects
 regex
 ruff
 pydantic
@@ -1,2 +1,3 @@
 lmcache >= 0.3.9
 nixl >= 0.7.1 # Required for disaggregated prefill
+mooncake-transfer-engine >= 0.3.8
@@ -50,10 +50,9 @@ def test_tp1_fp8_fusions(
    run_e2e_fusion_test,
    monkeypatch,
 ):
-    if use_deepgemm:
-        # TODO(luka/eliza) DeepGEMM uses different quants, matching not supported
+    if use_deepgemm and is_blackwell():
+        # TODO(luka) DeepGEMM uses different quants, matching not supported
        #  - on Blackwell, uses a special quant fp8, currently not supported
-        #  - on Hopper, tma-aligned scales inhibit matching (fix WIP)
        pytest.skip("DeepGEMM & quant matching not currently supported")

    matches = matches_fn(n_layers)
@@ -66,7 +65,6 @@ def test_tp1_fp8_fusions(
    model_kwargs["hf_overrides"] = hf_overrides(n_layers)
    model_kwargs["load_format"] = "dummy"
    model_kwargs["max_model_len"] = 1024
-
    compilation_config = dict(
        use_inductor_graph_partition=inductor_graph_partition,
        custom_ops=custom_ops.split(","),
@@ -7,7 +7,6 @@ from vllm.entrypoints.llm import LLM
 from vllm.sampling_params import SamplingParams


-@pytest.mark.skip_v1
@pytest.mark.parametrize("model", ["distilbert/distilgpt2"])
 def test_computed_prefix_blocks(model: str):
    # This test checks if the engine generates completions both with and
@@ -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)
@@ -134,6 +134,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
@@ -121,3 +121,75 @@ 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})
+
+            # 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.
+            event = await receive_event(ws, timeout=30.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})
+
+            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"""

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

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

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

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

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

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

@@ -0,0 +1,8 @@
+model_name: "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8"
+accuracy_threshold: 0.29
+num_questions: 1319
+num_fewshot: 5
+server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2"
+env:
+  VLLM_USE_FLASHINFER_MOE_FP8: "1"
+  VLLM_FLASHINFER_MOE_BACKEND: "latency"
@@ -0,0 +1,8 @@
+model_name: "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4"
+accuracy_threshold: 0.29
+num_questions: 1319
+num_fewshot: 5
+server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2"
+env:
+  VLLM_USE_FLASHINFER_MOE_FP4: "1"
+  VLLM_FLASHINFER_MOE_BACKEND: "throughput"
@@ -13,3 +13,5 @@ Llama-4-Scout-BF16-fi-cutlass.yaml
 Llama-4-Scout-BF16-triton.yaml
 Mixtral-8x7B-BF16-fi-cutlass.yaml
 Mixtral-8x7B-BF16-triton.yaml
+Nemotron-Nano-30B-Fp8-ModelOpt-fi-trtllm.yaml
+Nemotron-Nano-30B-NvFp4-ModelOpt-fi-cutlass.yaml
@@ -10,7 +10,7 @@ from vllm.utils.math_utils import next_power_of_2
 from vllm.utils.torch_utils import set_random_seed
 from vllm.v1.attention.ops.triton_unified_attention import unified_attention

-NUM_HEADS = [(4, 4), (8, 2)]
+NUM_HEADS = [(4, 4), (8, 2), (5, 1)]
 HEAD_SIZES = [128, 256]
 BLOCK_SIZES = [16]

@@ -20,6 +20,8 @@ QDTYPES = (
    if not current_platform.is_rocm()
    else [None, torch.float8_e4m3fnuz]
 )
+FP8_DTYPE = current_platform.fp8_dtype()
+
 # one value large enough to test overflow in index calculation.
 # one value small enough to test the schema op check
 NUM_BLOCKS = [32768, 2048]
@@ -217,3 +219,127 @@ def test_triton_unified_attn(
        torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol),
        f"{torch.max(torch.abs(output - ref_output))}",
    )
+
+
+@pytest.mark.parametrize(
+    "seq_lens",
+    [
+        [(1, 1328), (5, 18), (129, 463)],
+        [(1, 523), (1, 37), (1, 2011)],
+        [(1, 1)] * 533,
+        [(533, 533)] * 533,
+    ],
+)
+@pytest.mark.parametrize("num_heads", NUM_HEADS)
+@pytest.mark.parametrize("head_size", HEAD_SIZES)
+@pytest.mark.parametrize("block_size", BLOCK_SIZES)
+@pytest.mark.parametrize("sliding_window", [None, 64, 128, 256])
+@pytest.mark.parametrize("soft_cap", [None, 50.0])
+@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
+@pytest.mark.parametrize("seq_threshold_3D", SEQ_THRESHOLD_3D_VALUES)
+@torch.inference_mode()
+def test_triton_unified_attn_fp16_input_fp8_output(
+    seq_lens: list[tuple[int, int]],
+    num_heads: tuple[int, int],
+    head_size: int,
+    sliding_window: int | None,
+    block_size: int,
+    soft_cap: float | None,
+    num_blocks: int,
+    seq_threshold_3D: int,
+) -> None:
+    """Test with fp16 input and fp8 output using output_scale."""
+    torch.set_default_device("cuda")
+
+    set_random_seed(0)
+    num_seqs = len(seq_lens)
+    query_lens = [x[0] for x in seq_lens]
+    kv_lens = [x[1] for x in seq_lens]
+    num_query_heads = num_heads[0]
+    num_kv_heads = num_heads[1]
+    assert num_query_heads % num_kv_heads == 0
+    max_query_len = max(query_lens)
+    max_kv_len = max(kv_lens)
+    window_size = (sliding_window - 1, 0) if sliding_window is not None else (-1, -1)
+    scale = head_size**-0.5
+
+    dtype = torch.float16
+    query = torch.randn(sum(query_lens), num_query_heads, head_size, dtype=dtype)
+    key_cache = torch.randn(
+        num_blocks, block_size, num_kv_heads, head_size, dtype=dtype
+    )
+    value_cache = torch.randn_like(key_cache)
+    cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(
+        dim=0, dtype=torch.int32
+    )
+    kv_lens_tensor = torch.tensor(kv_lens, dtype=torch.int32)
+
+    max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
+    block_tables = torch.randint(
+        0, num_blocks, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32
+    )
+
+    output = torch.empty(sum(query_lens), num_query_heads, head_size, dtype=FP8_DTYPE)
+
+    output_scale = torch.tensor(0.5, dtype=torch.float32)
+
+    num_par_softmax_segments = 16
+    head_size_padded = next_power_of_2(head_size)
+    softmax_segm_output = torch.empty(
+        (seq_threshold_3D, num_query_heads, num_par_softmax_segments, head_size_padded),
+        dtype=torch.float32,
+    )
+    softmax_segm_max = torch.empty(
+        (seq_threshold_3D, num_query_heads, num_par_softmax_segments),
+        dtype=torch.float32,
+    )
+    softmax_segm_expsum = torch.empty(
+        (seq_threshold_3D, num_query_heads, num_par_softmax_segments),
+        dtype=torch.float32,
+    )
+
+    unified_attention(
+        q=query,
+        k=key_cache,
+        v=value_cache,
+        out=output,
+        cu_seqlens_q=cu_query_lens,
+        seqused_k=kv_lens_tensor,
+        max_seqlen_q=max_query_len,
+        max_seqlen_k=max_kv_len,
+        softmax_scale=scale,
+        causal=True,
+        window_size=window_size,
+        block_table=block_tables,
+        softcap=soft_cap if soft_cap is not None else 0,
+        q_descale=None,
+        k_descale=None,
+        v_descale=None,
+        output_scale=output_scale,
+        seq_threshold_3D=seq_threshold_3D,
+        num_par_softmax_segments=num_par_softmax_segments,
+        softmax_segm_output=softmax_segm_output,
+        softmax_segm_max=softmax_segm_max,
+        softmax_segm_expsum=softmax_segm_expsum,
+    )
+
+    ref_output = ref_paged_attn(
+        query=query,
+        key_cache=key_cache,
+        value_cache=value_cache,
+        query_lens=query_lens,
+        kv_lens=kv_lens,
+        block_tables=block_tables,
+        scale=scale,
+        sliding_window=sliding_window,
+        soft_cap=soft_cap,
+    )
+
+    output_fp16 = output.to(torch.float32) * output_scale.item()
+    output_fp16 = output_fp16.to(torch.float16)
+
+    atol, rtol = 2e-1, 2e-1
+    (
+        torch.testing.assert_close(output_fp16, ref_output, atol=atol, rtol=rtol),
+        f"{torch.max(torch.abs(output_fp16 - ref_output))}",
+    )
@@ -2,6 +2,8 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project


+import itertools
+
 import pytest
 import torch

@@ -21,7 +23,7 @@ QUANT_DTYPES = [torch.int8, current_platform.fp8_dtype()]
 VEC_HIDDEN_SIZES = [1024, 1025, 1027, 1029]
 # Avoid combinatorial explosion with full Cartesian product
 NUM_TOKENS_HIDDEN_SIZES = [
-    *[(1, i) for i in [1, 64, *VEC_HIDDEN_SIZES, 5120, 5137]],
+    *[(1, i) for i in [1, 64, 128, *VEC_HIDDEN_SIZES, 5120, 5137]],
    *[(2048, i) for i in [1, 64, *VEC_HIDDEN_SIZES, 5137]],
    *[(4096, i) for i in [1, 64, 5137]],
 ]
@@ -29,6 +31,7 @@ NUM_TOKENS_HIDDEN_SIZES = [
 ADD_RESIDUAL = [False, True]
 SCALE_UBS = [True, False]
 GROUP_SIZES = [None, [1, 64], [1, 128]]
+TMA_ALIGNMENTS = [0, 4]
 SEEDS = [0]
 CUDA_DEVICES = [f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)]

@@ -110,12 +113,21 @@ def ops_dynamic_per_token_or_block_quant(
    residual: torch.Tensor | None,
    scale_ub: torch.Tensor | None,
    group_size: list[int] | None,
+    tma_alignment: int,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
    if residual is not None:
        residual = residual.clone()
    if group_size is not None:
        out, scales = ops.rms_norm_per_block_quant(
-            x, weight, EPS, quant_dtype, group_size, scale_ub, residual, True
+            x,
+            weight,
+            EPS,
+            quant_dtype,
+            group_size,
+            scale_ub,
+            residual,
+            True,
+            tma_alignment,
        )
        scales = scales.contiguous()
    else:
@@ -132,9 +144,10 @@ def ops_impl(
    residual: torch.Tensor | None,
    scale_ub: torch.Tensor | None,
    group_size: list[int] | None,
+    tma_alignment: int,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
    return ops_dynamic_per_token_or_block_quant(
-        weight, x, quant_dtype, residual, scale_ub, group_size
+        weight, x, quant_dtype, residual, scale_ub, group_size, tma_alignment
    )


@@ -143,7 +156,10 @@ def ops_impl(
@pytest.mark.parametrize("has_scale_ub", SCALE_UBS)
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("quant_dtype", QUANT_DTYPES)
-@pytest.mark.parametrize("group_size", GROUP_SIZES)
+@pytest.mark.parametrize(
+    "group_size, tma_alignment",
+    [(None, 0), *itertools.product(GROUP_SIZES, TMA_ALIGNMENTS)],
+)
@pytest.mark.parametrize("seed", SEEDS)
@pytest.mark.parametrize("device", CUDA_DEVICES)
@torch.inference_mode()
@@ -156,6 +172,7 @@ def test_rms_norm(
    dtype: torch.dtype,
    quant_dtype: torch.dtype,
    group_size: list[int] | None,
+    tma_alignment: int,
    seed: int,
    device: str,
 ) -> None:
@@ -173,6 +190,20 @@ def test_rms_norm(
        # blockwise baseline doesn't support scale_ub
        return

+    if (
+        group_size is None or quant_dtype != current_platform.fp8_dtype()
+    ) and tma_alignment != 0:
+        # TMA alignment is only supported for groupwise fp8 kernels
+        return
+
+    if (
+        group_size is not None
+        and tma_alignment != 0
+        and hidden_size // group_size[1] % tma_alignment == 0
+    ):
+        # Skip tests where TMA alignment doesn't create extra padding to save time
+        return
+
    if has_scale_ub and quant_dtype != current_platform.fp8_dtype():
        # skip
        return
@@ -196,7 +227,7 @@ def test_rms_norm(
        layer, x, quant_dtype, residual, scale_ub, group_size
    )
    ops_out, ops_scales, ops_residual = ops_impl(
-        layer.weight, x, quant_dtype, residual, scale_ub, group_size
+        layer.weight, x, quant_dtype, residual, scale_ub, group_size, tma_alignment
    )

    assert ref_out.dtype == quant_dtype
@@ -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])
@@ -398,80 +398,3 @@ def test_convert_moe_weights_to_flashinfer_trtllm_block_layout(

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

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

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

-        if renormalize:
-            torch.testing.assert_close(
-                baseline_topk_weights, test_topk_weights, atol=2e-2, rtol=0
-            )
+        torch.testing.assert_close(
+            baseline_topk_weights, test_topk_weights, atol=2e-2, rtol=0
+        )
        torch.testing.assert_close(baseline_topk_ids, test_topk_ids, atol=0, rtol=0)
@@ -9,6 +9,7 @@ from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
    UnquantizedMoeBackend,
    select_unquantized_moe_backend,
 )
+from vllm.platforms import current_platform


@pytest.mark.parametrize(
@@ -65,6 +66,9 @@ def test_select_default_backend_by_platform(
    "vllm.model_executor.layers.fused_moe.oracle.unquantized.is_supported_config_trtllm_bf16",
    return_value=(True, None),
 )
+@pytest.mark.skipif(
+    not current_platform.is_cuda(), reason="Only supported on NVIDIA platforms."
+)
 def test_select_cuda_flashinfer_trtllm_backend(
    mock_has_flashinfer, mock_is_supported_trtllm, monkeypatch
 ):
@@ -101,6 +105,9 @@ def test_select_cuda_flashinfer_trtllm_backend(
    "vllm.model_executor.layers.fused_moe.oracle.unquantized.is_supported_config_trtllm_bf16",
    return_value=(False, None),
 )
+@pytest.mark.skipif(
+    not current_platform.is_cuda(), reason="Only supported on NVIDIA platforms."
+)
 def test_select_cuda_flashinfer_cutlass_backend(
    mock_has_flashinfer, mock_is_supported_trtllm, monkeypatch
 ):
@@ -155,9 +155,9 @@ def test_rocm_wvsplitkrc_kernel(xnorm, n, k, m, dtype, seed, bias_mode):
    out = ops.wvSplitKrc(B, A.view(-1, A.size(-1)), cu_count, BIAS)

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


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

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


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

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


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

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


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

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


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

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

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


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

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

@@ -169,13 +169,6 @@ VLM_TEST_SETTINGS = {
        auto_cls=AutoModelForImageTextToText,
        vllm_output_post_proc=model_utils.qwen2_vllm_to_hf_output,
        patch_hf_runner=model_utils.qwen3_vl_patch_hf_runner,
-        vllm_runner_kwargs={
-            "attention_config": {
-                "backend": "ROCM_AITER_FA",
-            },
-        }
-        if current_platform.is_rocm()
-        else None,
        image_size_factors=[(0.25,), (0.25, 0.25, 0.25), (0.25, 0.2, 0.15)],
        marks=[
            pytest.mark.core_model,
@@ -7,8 +7,10 @@ import mimetypes
 import os
 from tempfile import NamedTemporaryFile, TemporaryDirectory

+import aiohttp
 import numpy as np
 import pytest
+import requests
 import torch
 from PIL import Image, ImageChops

@@ -318,3 +320,58 @@ async def test_allowed_media_domains(video_url: str, num_frames: int):

    with pytest.raises(ValueError):
        _, _ = await connector.fetch_video_async(disallowed_url)
+
+
+@pytest.mark.asyncio
+async def test_ssrf_bypass_backslash_in_url(local_asset_server):
+    """Verify that backslash-@ URL parsing confusion cannot bypass the
+    allowed_media_domains check (GHSA-v359-jj2v-j536).
+
+    urllib3.parse_url() and aiohttp/yarl disagree on how to parse a
+    backslash before ``@``.  urllib3 treats ``\\`` as part of the path
+    (encoding it as ``%5C``), while yarl treats it as a userinfo
+    separator, changing the effective host.  The fix normalises the URL
+    through urllib3 *before* handing it to aiohttp so both layers agree.
+    """
+    port = local_asset_server.port
+    asset = TEST_IMAGE_ASSETS[0]
+
+    # Craft the bypass payload: urllib3 sees host=127.0.0.1, but an
+    # un-patched aiohttp would see host=example.com.
+    bypass_url = f"http://127.0.0.1:{port}\\@example.com/{asset}"
+
+    connector = MediaConnector(
+        allowed_media_domains=["127.0.0.1"],
+    )
+
+    # After the fix the request is made to 127.0.0.1 (the local asset
+    # server) using the normalised URL.  The normalised path will be
+    # /%5C@example.com/<asset> which won't match any file the server
+    # knows about, so we expect an HTTP error — but crucially NOT a
+    # successful fetch from example.com.
+    with pytest.raises(requests.exceptions.HTTPError):
+        connector.fetch_image(bypass_url)
+
+    with pytest.raises(aiohttp.ClientResponseError):
+        await connector.fetch_image_async(bypass_url)
+
+
+@pytest.mark.asyncio
+async def test_ssrf_bypass_backslash_disallowed_domain():
+    """The reverse direction: even when the *attacker-controlled* host
+    appears in the urllib3-parsed hostname position the allowlist must
+    still block it.
+    """
+    # urllib3.parse_url sees host=example.com which is NOT in the
+    # allowlist, so this must be rejected before any request is made.
+    bypass_url = "https://example.com\\@safe.example.org/image.png"
+
+    connector = MediaConnector(
+        allowed_media_domains=["safe.example.org"],
+    )
+
+    with pytest.raises(ValueError, match="allowed domains"):
+        connector.fetch_image(bypass_url)
+
+    with pytest.raises(ValueError, match="allowed domains"):
+        await connector.fetch_image_async(bypass_url)
@@ -816,3 +816,26 @@ def test_compressed_tensors_moe_ignore_with_model(vllm_runner):
        # Verify the model can generate output
        output = llm.generate_greedy("Hello, my name is", max_tokens=4)
        assert output
+
+
+def test_w4a16_moe_torch_compile(vllm_runner):
+    """Regression test: MoE quant_config must be initialized inside the
+    moe_forward custom op, not just in forward_native which is compiled by
+    Dynamo (attribute mutations are not replayed at runtime).
+
+    Without the fix in _moe_forward/_moe_forward_shared, this hits:
+        AssertionError: Hidden size mismatch 2048 != 1024
+    because use_int4_w4a16 is False (moe_quant_config stays None).
+    """
+    model_path = "nm-testing/tinysmokeqwen3moe-W4A16-first-only-CTstable"
+
+    with vllm_runner(
+        model_path,
+        enforce_eager=False,
+        max_model_len=256,
+        compilation_config={
+            "cudagraph_mode": "NONE",
+        },
+    ) as llm:
+        output = llm.generate_greedy("Hi", max_tokens=1)
+        assert output
@@ -167,6 +167,23 @@ class TestBaseThinkingReasoningParserMethods:
            is False
        )

+    def test_count_reasoning_tokens(self, test_tokenizer):
+        """Count tokens between start/end markers."""
+        parser = TestThinkingReasoningParser(test_tokenizer)
+        start = parser.start_token_id
+        end = parser.end_token_id
+        token_ids = [0, start, 11, 12, end, 99]
+        assert parser.count_reasoning_tokens(token_ids) == 2
+
+    def test_count_reasoning_tokens_nested(self, test_tokenizer):
+        """Ensure nested thinking spans count all inner tokens safely."""
+        parser = TestThinkingReasoningParser(test_tokenizer)
+        s = parser.start_token_id
+        e = parser.end_token_id
+        token_ids = [s, 1, s, 2, e, 3, e]
+        # Tokens 1,2,3 are inside reasoning (depth>0) => 3 tokens
+        assert parser.count_reasoning_tokens(token_ids) == 3
+
    def test_extract_content_ids(self, test_tokenizer):
        """Test the extract_content_ids method."""
        parser = TestThinkingReasoningParser(test_tokenizer)
@@ -6,6 +6,7 @@ import pytest
 from vllm.assets.image import ImageAsset
 from vllm.assets.video import VideoAsset
 from vllm.config import CacheConfig, ModelConfig, VllmConfig
+from vllm.multimodal.parse import parse_mm_uuids
 from vllm.renderers.hf import HfRenderer
 from vllm.tokenizers.registry import tokenizer_args_from_config

@@ -41,14 +42,25 @@ def test_multi_modal_uuids_length_mismatch_raises():

    mm_data = {"image": [cherry_pil_image, stop_pil_image]}

+    # Mismatch: 2 items but only 0 uuids provided
+    mm_uuids = {"image": []}  # type: ignore[var-annotated]
+
+    mm_processor = renderer.get_mm_processor()
+    mm_data_items = mm_processor.info.parse_mm_data(mm_data)
+    mm_uuid_items = parse_mm_uuids(mm_uuids)
+
+    with pytest.raises(ValueError, match="must have same length as"):
+        renderer._process_mm_uuids(mm_data, mm_data_items, mm_uuid_items, "req-1a")
+
    # Mismatch: 2 items but only 1 uuid provided
    mm_uuids = {"image": ["hash_cherry"]}

    mm_processor = renderer.get_mm_processor()
-    mm_items = mm_processor.info.parse_mm_data(mm_data)
+    mm_data_items = mm_processor.info.parse_mm_data(mm_data)
+    mm_uuid_items = parse_mm_uuids(mm_uuids)

    with pytest.raises(ValueError, match="must have same length as"):
-        renderer._process_mm_uuids(mm_data, mm_items, mm_uuids, "req-1")
+        renderer._process_mm_uuids(mm_data, mm_data_items, mm_uuid_items, "req-1b")


 def test_multi_modal_uuids_missing_modality_raises():
@@ -63,10 +75,11 @@ def test_multi_modal_uuids_missing_modality_raises():
    mm_uuids = {"image": ["hash_cherry"]}

    mm_processor = renderer.get_mm_processor()
-    mm_items = mm_processor.info.parse_mm_data(mm_data)
+    mm_data_items = mm_processor.info.parse_mm_data(mm_data)
+    mm_uuid_items = parse_mm_uuids(mm_uuids)

    with pytest.raises(ValueError, match="is empty but .* is missing"):
-        renderer._process_mm_uuids(mm_data, mm_items, mm_uuids, "req-2")
+        renderer._process_mm_uuids(mm_data, mm_data_items, mm_uuid_items, "req-2")


@pytest.mark.parametrize(
@@ -78,7 +91,7 @@ def test_multi_modal_uuids_missing_modality_raises():
    ],
 )
 def test_multi_modal_uuids_accepts_none_and_passes_through(
-    monkeypatch, mm_cache_gb: float, enable_prefix_caching: bool
+    mm_cache_gb: float, enable_prefix_caching: bool
 ):
    renderer = _build_renderer(
        mm_cache_gb=mm_cache_gb,
@@ -94,9 +107,11 @@ def test_multi_modal_uuids_accepts_none_and_passes_through(
    mm_uuids = {"image": [None, "hash_stop"], "video": None}

    mm_processor = renderer.get_mm_processor()
-    mm_items = mm_processor.info.parse_mm_data(mm_data)
+    mm_data_items = mm_processor.info.parse_mm_data(mm_data)
+    mm_uuid_items = parse_mm_uuids(mm_uuids)
+
    processed_mm_uuids = renderer._process_mm_uuids(
-        mm_data, mm_items, mm_uuids, "req-3"
+        mm_data, mm_data_items, mm_uuid_items, "req-3"
    )

    assert processed_mm_uuids == mm_uuids
@@ -111,7 +126,7 @@ def test_multi_modal_uuids_accepts_none_and_passes_through(
    ],
 )
 def test_multi_modal_uuids_accepts_empty(
-    monkeypatch, mm_cache_gb: float, enable_prefix_caching: bool
+    mm_cache_gb: float, enable_prefix_caching: bool
 ):
    renderer = _build_renderer(
        mm_cache_gb=mm_cache_gb,
@@ -120,19 +135,21 @@ def test_multi_modal_uuids_accepts_empty(

    # While None means cached multi-modal input requiring UUIDs
    # an empty list means no multi-modal input
-    mm_data = {"image": [], "video": []}  # type: ignore[var-annotated]
-    mm_uuids = {"image": [], "video": None}  # type: ignore[var-annotated]
+    mm_data = {"image": [], "video": [], "audio": None}  # type: ignore[var-annotated]
+    mm_uuids = {"image": [], "video": None, "audio": []}  # type: ignore[var-annotated]

    mm_processor = renderer.get_mm_processor()
-    mm_items = mm_processor.info.parse_mm_data(mm_data)
+    mm_data_items = mm_processor.info.parse_mm_data(mm_data)
+    mm_uuid_items = parse_mm_uuids(mm_uuids)
+
    processed_mm_uuids = renderer._process_mm_uuids(
-        mm_data, mm_items, mm_uuids, "req-4"
+        mm_data, mm_data_items, mm_uuid_items, "req-4"
    )

    assert processed_mm_uuids == mm_uuids


-def test_multi_modal_uuids_ignored_when_caching_disabled(monkeypatch):
+def test_multi_modal_uuids_ignored_when_caching_disabled():
    # When both processor cache is 0 and prefix caching disabled, the
    # processor builds overrides from request id instead of using user UUIDs.
    renderer = _build_renderer(mm_cache_gb=0.0, enable_prefix_caching=False)
@@ -145,9 +162,11 @@ def test_multi_modal_uuids_ignored_when_caching_disabled(monkeypatch):
    mm_uuids = {"image": ["hash_cherry", "hash_stop"], "video": ["hash_video"]}

    mm_processor = renderer.get_mm_processor()
-    mm_items = mm_processor.info.parse_mm_data(mm_data)
+    mm_data_items = mm_processor.info.parse_mm_data(mm_data)
+    mm_uuid_items = parse_mm_uuids(mm_uuids)
+
    processed_mm_uuids = renderer._process_mm_uuids(
-        mm_data, mm_items, mm_uuids, request_id
+        mm_data, mm_data_items, mm_uuid_items, request_id
    )

    # Expect request-id-based overrides are passed through
@@ -9,6 +9,26 @@ import pytest
 from transformers import AutoModelForSeq2SeqLM

 from vllm.assets.audio import AudioAsset
+from vllm.platforms import current_platform
+
+# Extra engine kwargs needed for numerically deterministic beam search.
+# On ROCm, floating-point reductions in attention and GEMM kernels are
+# non-associative and sensitive to batch geometry, so we:
+#   async_scheduling=False      – deterministic batch composition
+#   enforce_eager=True          – no CUDA-graph padding changing effective size
+#   enable_prefix_caching=False – avoid prefix-sharing side effects
+#   max_num_seqs=1              – fixed batch size across runs
+# On other platforms these are not needed and the dict is empty.
+EXTRA_ENGINE_KWARGS: dict = (
+    dict(
+        async_scheduling=False,
+        enforce_eager=True,
+        enable_prefix_caching=False,
+        max_num_seqs=1,
+    )
+    if current_platform.is_rocm()
+    else dict(async_scheduling=False, max_num_seqs=1)
+)

 # FIXME(zhuohan): The test can not pass if we:
 #   1. Increase max_tokens to 256.
@@ -25,6 +45,7 @@ MODELS = ["TinyLlama/TinyLlama-1.1B-Chat-v1.0"]
@pytest.mark.parametrize("max_tokens", MAX_TOKENS)
@pytest.mark.parametrize("beam_width", BEAM_WIDTHS)
 def test_beam_search_single_input(
+    monkeypatch,
    hf_runner,
    vllm_runner,
    example_prompts,
@@ -33,13 +54,16 @@ def test_beam_search_single_input(
    max_tokens: int,
    beam_width: int,
 ) -> None:
+    if current_platform.is_rocm():
+        monkeypatch.setenv("VLLM_ROCM_USE_SKINNY_GEMM", "0")
+
    example_prompts = example_prompts[:1]
    with hf_runner(model, dtype=dtype) as hf_model:
        hf_outputs = hf_model.generate_beam_search(
            example_prompts, beam_width, max_tokens
        )

-    with vllm_runner(model, dtype=dtype) as vllm_model:
+    with vllm_runner(model, dtype=dtype, **EXTRA_ENGINE_KWARGS) as vllm_model:
        vllm_outputs = vllm_model.generate_beam_search(
            example_prompts, beam_width, max_tokens
        )
@@ -66,6 +90,7 @@ def test_beam_search_single_input(
@pytest.mark.parametrize("max_tokens", MAX_TOKENS)
@pytest.mark.parametrize("beam_width", BEAM_WIDTHS)
 def test_beam_search_with_concurrency_limit(
+    monkeypatch,
    hf_runner,
    vllm_runner,
    example_prompts,
@@ -74,21 +99,29 @@ def test_beam_search_with_concurrency_limit(
    max_tokens: int,
    beam_width: int,
 ) -> None:
+    if current_platform.is_rocm():
+        monkeypatch.setenv("VLLM_ROCM_USE_SKINNY_GEMM", "0")
+
    # example_prompts[1]&[3]&[7] fails due to unknown reason even without
    # concurrency limit. skip them for now.
    example_prompts = example_prompts[:8]
    concurrency_limit = 2
    assert len(example_prompts) > concurrency_limit
-    with vllm_runner(model, dtype=dtype) as vllm_model:
+    with vllm_runner(model, dtype=dtype, **EXTRA_ENGINE_KWARGS) as vllm_model:
        outputs_with_limit = vllm_model.generate_beam_search(
-            example_prompts, beam_width, max_tokens, concurrency_limit=concurrency_limit
+            example_prompts,
+            beam_width,
+            max_tokens,
+            concurrency_limit=concurrency_limit,
        )
        outputs_without_limit = []

        for i in range(0, len(example_prompts), concurrency_limit):
            outputs_without_limit.extend(
                vllm_model.generate_beam_search(
-                    example_prompts[i : i + concurrency_limit], beam_width, max_tokens
+                    example_prompts[i : i + concurrency_limit],
+                    beam_width,
+                    max_tokens,
                )
            )

@@ -118,6 +151,7 @@ def test_beam_search_with_concurrency_limit(
@pytest.mark.parametrize("max_tokens", MAX_TOKENS)
@pytest.mark.parametrize("beam_width", MM_BEAM_WIDTHS)
 def test_beam_search_passes_multimodal_data(
+    monkeypatch,
    hf_runner,
    vllm_runner,
    dtype: str,
@@ -125,6 +159,9 @@ def test_beam_search_passes_multimodal_data(
    beam_width: int,
 ) -> None:
    """Ensure that beam search passes multimodal data through correctly."""
+    if current_platform.is_rocm():
+        monkeypatch.setenv("VLLM_ROCM_USE_SKINNY_GEMM", "0")
+
    # NOTE - this test is primarily to check that mm data is passed to beams
    # correctly. As such, we just need to check one extra modality to make
    # sure things pass through properly.
@@ -145,7 +182,7 @@ def test_beam_search_passes_multimodal_data(
            audios=audios,
        )

-    with vllm_runner(model, dtype=dtype) as vllm_model:
+    with vllm_runner(model, dtype=dtype, **EXTRA_ENGINE_KWARGS) as vllm_model:
        vllm_outputs = vllm_model.generate_beam_search(
            prompts,
            beam_width=beam_width,
@@ -190,7 +190,7 @@ class RemoteOpenAIServer:
            model_loader.download_model(model_config)

        self._start_server(model, vllm_serve_args, env_dict)
-        max_wait_seconds = max_wait_seconds or 240
+        max_wait_seconds = max_wait_seconds or 360
        self._wait_for_server(url=self.url_for("health"), timeout=max_wait_seconds)

    def __enter__(self):
@@ -19,8 +19,13 @@ from tests.v1.attention.utils import (
 )
 from vllm import _custom_ops as ops
 from vllm.config.vllm import set_current_vllm_config
-from vllm.model_executor.layers.attention.mla_attention import QueryLenSupport
+from vllm.model_executor.layers.attention.mla_attention import (
+    QueryLenSupport,
+    _DecodeConcatQuantFP8,
+)
 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
+from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
+from vllm.platforms import current_platform
 from vllm.utils.math_utils import cdiv
 from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
 from vllm.v1.attention.backend import CommonAttentionMetadata
@@ -50,6 +55,7 @@ if not flash_attn_supports_mla():
 if not is_flashmla_dense_supported()[0]:
    BACKENDS_TO_TEST.remove(AttentionBackendEnum.FLASHMLA)

+
 SPEC_DECODE_BACKENDS = []
 for backend in BACKENDS_TO_TEST:
    builder_cls, _ = try_get_attention_backend(backend)
@@ -144,9 +150,8 @@ def create_and_prepopulate_kv_cache(
        common_attn_metadata: Common attention metadata
        randomize_blocks: Whether to randomly permute blocks
                          or use sequential order
-        kv_cache_dtype: Optional kv cache dtype string. When set to
-                        "fp8_ds_mla" the cache is populated using the
-                        fp8 DeepSeek MLA layout via concat_and_cache_mla.
+        kv_cache_dtype: Optional kv cache dtype string. For fp8 cache dtype,
+                        the cache is populated via concat_and_cache_mla.
        scale: Scaling factor forwarded to concat_and_cache_mla when the
               fp8 cache layout is requested.

@@ -163,18 +168,21 @@ def create_and_prepopulate_kv_cache(
    block_table = common_attn_metadata.block_table_tensor
    slot_mapping = common_attn_metadata.slot_mapping

+    fp8_attention = kv_cache_dtype and kv_cache_dtype.startswith("fp8")
    use_fp8_ds_mla = kv_cache_dtype == "fp8_ds_mla"

-    if use_fp8_ds_mla:
-        if not kv_c_contexts:
-            raise ValueError(
-                "kv_c_contexts cannot be empty when using fp8_ds_mla cache dtype"
-            )
-        kv_lora_rank = kv_c_contexts[0].shape[-1]
-        rope_dim = k_pe_contexts[0].shape[-1]
-        entry_size = kv_lora_rank + 4 * 4 + 2 * rope_dim
+    if fp8_attention:
+        if use_fp8_ds_mla:
+            kv_lora_rank = kv_c_contexts[0].shape[-1]
+            rope_dim = k_pe_contexts[0].shape[-1]
+            # 4 * 4: 4 float32 scale values for 128-element tiles
+            # 2 * rope_dim: 16-bit RoPE values
+            kv_entry_size = kv_lora_rank + 4 * 4 + 2 * rope_dim
+        else:
+            kv_entry_size = head_size
+
        kv_cache = torch.zeros(
-            num_blocks, block_size, entry_size, dtype=torch.uint8, device=device
+            num_blocks, block_size, kv_entry_size, dtype=torch.uint8, device=device
        )
        scale_tensor = (
            scale
@@ -201,14 +209,14 @@ def create_and_prepopulate_kv_cache(

        start = start_block_idx * block_size

-        if use_fp8_ds_mla:
+        if fp8_attention:
            slots = torch.arange(context_len, device=device, dtype=torch.long) + start
            ops.concat_and_cache_mla(
                kv_c_context,
                k_pe_context.squeeze(1),
                kv_cache,
                slots,
-                kv_cache_dtype="fp8_ds_mla",
+                kv_cache_dtype=kv_cache_dtype,
                scale=scale_tensor,
            )
        else:
@@ -329,8 +337,9 @@ class MockSparseMLAAttentionLayer:
        output: torch.Tensor,
    ) -> torch.Tensor:
        """Forward for sparse MLA - uses forward_mqa for all tokens."""
-        # Write to KV cache
        kv_cache_dtype = getattr(self.impl, "kv_cache_dtype", "auto")
+
+        # Write to KV cache
        if kv_cache.numel() > 0:
            ops.concat_and_cache_mla(
                kv_c,
@@ -426,6 +435,12 @@ class MockMLAAttentionLayer(AttentionLayerBase):
        self._k_scale_float = 1.0
        self._v_scale_float = 1.0

+        self._decode_concat_quant_fp8_op = _DecodeConcatQuantFP8(
+            static=True,
+            group_shape=GroupShape.PER_TENSOR,
+            compile_native=True,
+        )
+
    def get_attn_backend(self):
        raise NotImplementedError

@@ -443,16 +458,21 @@ class MockMLAAttentionLayer(AttentionLayerBase):
    ) -> torch.Tensor:
        """Replicates MLAAttention.forward_impl logic for testing."""
        # Write to KV cache
+        kv_cache_dtype = getattr(self.impl, "kv_cache_dtype", "auto")
+        fp8_attention = kv_cache_dtype.startswith("fp8")
        if kv_cache.numel() > 0:
            ops.concat_and_cache_mla(
                kv_c,
                k_pe.squeeze(1),
                kv_cache,
                attn_metadata.slot_mapping.flatten(),
-                kv_cache_dtype="auto",
+                kv_cache_dtype=kv_cache_dtype,
                scale=self._k_scale,
            )

+        if fp8_attention and kv_cache_dtype != "fp8_ds_mla":
+            kv_cache = kv_cache.view(current_platform.fp8_dtype())
+
        # Determine decode vs prefill split
        num_decode_tokens = attn_metadata.num_decode_tokens or 0
        has_decode = (attn_metadata.num_decodes or 0) > 0
@@ -491,8 +511,14 @@ class MockMLAAttentionLayer(AttentionLayerBase):
            # Convert from (N, B, L) to (B, N, L)
            mqa_ql_nope = mqa_ql_nope.transpose(0, 1)

-            # Pass as tuple to forward_mqa
-            mqa_q = (mqa_ql_nope, mqa_q_pe)
+            if fp8_attention and self.impl.supports_quant_query_input:
+                assert mqa_ql_nope.shape[0] == mqa_q_pe.shape[0]
+                assert mqa_ql_nope.shape[1] == mqa_q_pe.shape[1]
+                mqa_q = self._decode_concat_quant_fp8_op(
+                    mqa_ql_nope, mqa_q_pe, self._q_scale
+                )
+            else:
+                mqa_q = (mqa_ql_nope, mqa_q_pe)

            attn_out, _ = self.impl.forward_mqa(mqa_q, kv_cache, attn_metadata, self)

@@ -526,6 +552,7 @@ def run_attention_backend(
    qk_rope_head_dim: int,
    v_head_dim: int,
    mock_kv_b_proj,
+    kv_cache_dtype: str = "auto",
 ) -> torch.Tensor:
    """Run attention computation using the specified backend's AttentionImpl."""

@@ -550,7 +577,7 @@ def run_attention_backend(
            num_kv_heads=num_kv_heads,
            alibi_slopes=None,
            sliding_window=None,
-            kv_cache_dtype="auto",
+            kv_cache_dtype=kv_cache_dtype,
            logits_soft_cap=None,
            attn_type="decoder",
            kv_sharing_target_layer_name=None,
@@ -630,12 +657,14 @@ def run_attention_backend(
 )
@pytest.mark.parametrize("model", ["deepseek-ai/DeepSeek-R1"])
@pytest.mark.parametrize("tensor_parallel_size", [1, 4, 8, 16])
+@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8", "fp8_e4m3"])
 def test_backend_correctness(
    default_vllm_config,
    dist_init,
    batch_spec_name: str,
    model: str,
    tensor_parallel_size: int,
+    kv_cache_dtype: str,
 ):
    """
    Test that all backends produce similar outputs to a reference implementation
@@ -658,9 +687,18 @@ def test_backend_correctness(
    head counts.
    """

+    # Filter backends to those that support the requested kv_cache_dtype
+    backends_to_test = [
+        b
+        for b in BACKENDS_TO_TEST
+        if kv_cache_dtype in b.get_class().supported_kv_cache_dtypes
+    ]
+    if not backends_to_test:
+        pytest.skip(f"No backends support kv_cache_dtype={kv_cache_dtype}")
+
    batch_spec = BATCH_SPECS[batch_spec_name]
    is_spec_decode_test = batch_spec_name.startswith("spec_decode")
-    unique_block_sizes = sorted(set(BACKEND_BLOCK_SIZES.values()))
+    unique_block_sizes = sorted(set(BACKEND_BLOCK_SIZES[b] for b in backends_to_test))
    default_block_size = unique_block_sizes[0]
    required_blocks = sum(
        (seq_len + default_block_size - 1) // default_block_size
@@ -694,6 +732,7 @@ def test_backend_correctness(
        block_size=default_block_size,
        hf_config_override=hf_config_override,
    )
+    vllm_config.cache_config.cache_dtype = kv_cache_dtype

    # For spec decode tests, add a speculative_config to set the reorder_batch_threshold
    if is_spec_decode_test:
@@ -751,7 +790,7 @@ def test_backend_correctness(

    kv_b_proj_weight = torch.cat([W_UK, W_UV], dim=-1)

-    for i, backend in enumerate(BACKENDS_TO_TEST):
+    for i, backend in enumerate(backends_to_test):
        all_sdpa_outputs.append([])

    for i in range(batch_size):
@@ -785,7 +824,7 @@ def test_backend_correctness(
        # pipeline (MHA-style). This ensures the reference implementation
        # matches each backend's actual decode/prefill pipeline path.
        is_decode = []
-        for backend_idx, backend in enumerate(BACKENDS_TO_TEST):
+        for backend_idx, backend in enumerate(backends_to_test):
            builder_cls, _ = try_get_attention_backend(backend)
            if is_spec_decode_test:
                query_len_support = getattr(
@@ -885,7 +924,7 @@ def test_backend_correctness(
        sdpa_out_i_prefill = sdpa_out_i_prefill.transpose(1, 2).squeeze(0)
        sdpa_out_i_prefill = sdpa_out_i_prefill.flatten(start_dim=-2)

-        for backend_idx, backend in enumerate(BACKENDS_TO_TEST):
+        for backend_idx, backend in enumerate(backends_to_test):
            if is_decode[backend_idx]:
                all_sdpa_outputs[backend_idx].append(sdpa_out_i_decode)
            else:
@@ -905,7 +944,7 @@ def test_backend_correctness(
    kv_c_vllm = torch.cat(all_kv_c_vllm, dim=0)
    k_pe_vllm = torch.cat(all_k_pe_vllm, dim=0)
    sdpa_outputs = {}
-    for backend_idx, backend in enumerate(BACKENDS_TO_TEST):
+    for backend_idx, backend in enumerate(backends_to_test):
        sdpa_outputs[backend] = torch.cat(all_sdpa_outputs[backend_idx], dim=0)

    # Create mock kv_b_proj using the same weights as reference implementation
@@ -973,12 +1012,13 @@ def test_backend_correctness(
            num_blocks=num_blocks_for_size,
            common_attn_metadata=common_attn_metadata,
            randomize_blocks=True,
+            kv_cache_dtype=kv_cache_dtype,
        )
        kv_cache_per_block_size[block_size] = kv_cache

    # 4. Run vLLM backends and compare
    failures = []
-    for backend_idx, backend_name in enumerate(BACKENDS_TO_TEST):
+    for backend_idx, backend_name in enumerate(backends_to_test):
        # Skip backends that don't support spec decode for spec decode tests
        if is_spec_decode_test and backend_name not in SPEC_DECODE_BACKENDS:
            continue
@@ -997,7 +1037,7 @@ def test_backend_correctness(
            head_size=vllm_config.model_config.get_head_size(),
            dtype=vllm_config.model_config.dtype,
            sliding_window=vllm_config.model_config.get_sliding_window(),
-            cache_dtype_str=vllm_config.cache_config.cache_dtype,
+            cache_dtype_str=kv_cache_dtype,
        )

        backend_output = run_attention_backend(
@@ -1016,6 +1056,7 @@ def test_backend_correctness(
            qk_rope_head_dim,
            v_head_dim,
            mock_kv_b_proj,
+            kv_cache_dtype=kv_cache_dtype,
        )

        # Use backend_idx to get the correct SDPA output for this backend
@@ -17,6 +17,7 @@ from vllm.config import (
    ModelConfig,
    SchedulerConfig,
    VllmConfig,
+    set_current_vllm_config,
 )
 from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import (
    MoRIIOAgentMetadata,
@@ -433,10 +434,11 @@ def test_register_kv_caches(mock_parallel_groups):
            }
        )

-        connector = MoRIIOConnector(vllm_config, KVConnectorRole.WORKER)
-        connector.connector_worker = FakeMorIIOConnectorWorker(
-            vllm_config, connector.engine_id, hand_shake_latency=0
-        )
+        with set_current_vllm_config(vllm_config):
+            connector = MoRIIOConnector(vllm_config, KVConnectorRole.WORKER)
+            connector.connector_worker = FakeMorIIOConnectorWorker(
+                vllm_config, connector.engine_id, hand_shake_latency=0
+            )

        from mori.io import (
            MemoryDesc,
@@ -523,7 +525,8 @@ def test_moriio_handshake_returns_metadata(mock_parallel_groups):
                "handshake_port": handshake_port,
            }
        )
-        connector = MoRIIOConnector(vllm_config, KVConnectorRole.WORKER)
+        with set_current_vllm_config(vllm_config):
+            connector = MoRIIOConnector(vllm_config, KVConnectorRole.WORKER)

        # Execute register_kv_caches
        connector.register_kv_caches(kv_caches)
@@ -8,7 +8,7 @@ MODEL = "meta-llama/Llama-3.2-1B-Instruct"


 def test_unsupported_configs():
-    with pytest.raises(NotImplementedError):
+    with pytest.raises(ValueError):
        AsyncEngineArgs(
            model=MODEL,
            speculative_config={
@@ -450,15 +450,30 @@ def rms_norm_per_block_quant(
    scale_ub: torch.Tensor | None = None,
    residual: torch.Tensor | None = None,
    is_scale_transposed: bool = False,
+    tma_alignment: int = 0,
 ) -> tuple[torch.Tensor, torch.Tensor]:
    assert len(group_size) == 2
    output = torch.empty_like(input, dtype=quant_dtype)
    if is_scale_transposed:
-        scales = torch.empty(
-            (input.shape[-1] // group_size[1], input.numel() // input.shape[-1]),
-            device=input.device,
-            dtype=torch.float32,
-        ).transpose(0, 1)
+        if tma_alignment == 0:
+            scales = torch.empty(
+                (input.shape[-1] // group_size[1], input.numel() // input.shape[-1]),
+                device=input.device,
+                dtype=torch.float32,
+            ).transpose(0, 1)
+        else:
+            m = input.shape[-2]
+            sf_k = input.shape[-1] // group_size[1]
+            tma_aligned_m = (m + tma_alignment - 1) // tma_alignment * tma_alignment
+            shape = input.shape[:-2] + (m, sf_k)
+            stride = (
+                (1, tma_aligned_m)
+                if input.dim() == 2
+                else (tma_aligned_m * sf_k, 1, tma_aligned_m)
+            )
+            scales = torch.empty_strided(
+                shape, stride, device=input.device, dtype=torch.float32
+            )
    else:
        scales = torch.empty(
            (input.numel() // input.shape[-1], input.shape[-1] // group_size[1]),
@@ -466,6 +481,10 @@ def rms_norm_per_block_quant(
            dtype=torch.float32,
        )

+    assert tma_alignment in [0, 4], "Expected TMA alignment 0 or 4, but got " + str(
+        tma_alignment
+    )
+
    torch.ops._C.rms_norm_per_block_quant(
        output,
        input,
@@ -2770,6 +2789,24 @@ def sm100_cutlass_mla_get_workspace_size(
    )


+def dsv3_fused_a_gemm(
+    output: torch.Tensor,
+    mat_a: torch.Tensor,
+    mat_b: torch.Tensor,
+) -> None:
+    """DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens).
+
+    Computes output = mat_a @ mat_b.T where:
+      mat_a: [num_tokens, 7168] row-major bf16 (hidden states)
+      mat_b: [7168, 2112] column-major bf16 (weight transposed)
+      output: [num_tokens, 2112] row-major bf16
+
+    Optimized for the DeepSeek V2/V3 QKV A-projection at small batch sizes.
+    Requires SM 9.0+ (Hopper).
+    """
+    torch.ops._C.dsv3_fused_a_gemm(output, mat_a, mat_b)
+
+
 if hasattr(torch.ops._C, "weight_packed_linear"):

    @register_fake("_C::weight_packed_linear")
@@ -249,7 +249,7 @@ class CompilerManager:
        if graph_index == 0:
            # before compiling the first graph, record the start time
            global compilation_start_time
-            compilation_start_time = time.time()
+            compilation_start_time = time.perf_counter()

        compilation_counter.num_backend_compilations += 1

@@ -261,8 +261,7 @@ class CompilerManager:
            if graph_index == num_graphs - 1:
                # after loading the last graph for this shape, record the time.
                # there can be multiple graphs due to piecewise compilation.
-                now = time.time()
-                elapsed = now - compilation_start_time
+                elapsed = time.perf_counter() - compilation_start_time
                compilation_config.compilation_time += elapsed
                logger.info_once(
                    "Directly load the compiled graph(s) for compile range %s "
@@ -285,7 +284,7 @@ class CompilerManager:
        with self.compile_context(compile_range):
            # There is a compilation time optimization here.
            #
-            # If the (input metdata, graph, compiler config) are the same, then
+            # If the (input metadata, graph, compiler config) are the same, then
            # we want to avoid compiling the same artifact again. If we didn't
            # do this optimization, the backend compilation (InductorAdaptor or
            # InductorStandaloneAdaptor)
@@ -362,8 +361,7 @@ class CompilerManager:

        # after compiling the last graph, record the end time
        if graph_index == num_graphs - 1:
-            now = time.time()
-            elapsed = now - compilation_start_time
+            elapsed = time.perf_counter() - compilation_start_time
            compilation_config.compilation_time += elapsed
            logger.info_once(
                "Compiling a graph for compile range %s takes %.2f s",
@@ -974,7 +972,7 @@ class VllmBackend:
        compilation_counter.num_graphs_seen += 1
        from .monitor import torch_compile_start_time

-        dynamo_time = time.time() - torch_compile_start_time
+        dynamo_time = time.perf_counter() - torch_compile_start_time
        logger.info_once(
            "Dynamo bytecode transform time: %.2f s", dynamo_time, scope="local"
        )
@@ -407,10 +407,10 @@ def _support_torch_compile(
        if envs.VLLM_USE_AOT_COMPILE:
            """
            When using torch.compile in AOT mode, we store the cache artifacts
-            under VLLM_CACHE_ROOT/torch_aot_compile/{hash}/rank_i_j. The {hash}
-            contains all of the factors except for the source files being
-            traced through, because we don't actually know which source files
-            to check at this point (before dynamo runs).
+            under VLLM_CACHE_ROOT/torch_compile_cache/torch_aot_compile/{hash}
+            The {hash} contains all of the factors except for the source files
+            being traced through, because we don't actually know which source
+            files to check at this point (before dynamo runs).
            On loading we will actually look at the source files being traced
            through. If any source file have changed (compared with the
            serialized backend artifacts), then we need to generate a new AOT
@@ -424,6 +424,7 @@ def _support_torch_compile(
            hash_key = hashlib.sha256(str(factors).encode()).hexdigest()
            cache_dir = os.path.join(
                envs.VLLM_CACHE_ROOT,
+                "torch_compile_cache",
                "torch_aot_compile",
                hash_key,
            )
@@ -14,7 +14,7 @@ torch_compile_start_time: float = 0.0

 def start_monitoring_torch_compile(vllm_config: VllmConfig) -> None:
    global torch_compile_start_time
-    torch_compile_start_time = time.time()
+    torch_compile_start_time = time.perf_counter()

    compilation_config: CompilationConfig = vllm_config.compilation_config
    path = vllm_config.compile_debug_dump_path()
@@ -30,10 +30,11 @@ def start_monitoring_torch_compile(vllm_config: VllmConfig) -> None:

 def end_monitoring_torch_compile(vllm_config: VllmConfig) -> None:
    compilation_config: CompilationConfig = vllm_config.compilation_config
+    total_compile_time: float = time.perf_counter() - torch_compile_start_time
    if compilation_config.mode == CompilationMode.VLLM_COMPILE:
        logger.info_once(
            "torch.compile takes %.2f s in total",
-            compilation_config.compilation_time,
+            total_compile_time,
            scope="local",
        )
        global context_manager
@@ -292,6 +292,7 @@ class MatcherQuantFP8(MatcherCustomOp):
        has_col_major_scales: bool = False,
        is_e8m0: bool = False,
        match_rocm_aiter: bool = False,
+        is_tma_aligned: bool = False,
    ) -> None:
        if enabled is None:
            enabled = QuantFP8.enabled()
@@ -301,6 +302,7 @@ class MatcherQuantFP8(MatcherCustomOp):
        self.has_col_major_scales = has_col_major_scales
        self.is_e8m0 = is_e8m0
        self.match_rocm_aiter = match_rocm_aiter
+        self.is_tma_aligned = is_tma_aligned

        if match_rocm_aiter:
            assert not quant_key.scale.group_shape.is_per_tensor(), (
@@ -336,6 +338,7 @@ class MatcherQuantFP8(MatcherCustomOp):
            quant_key.scale.group_shape,
            column_major_scales=has_col_major_scales,
            use_ue8m0=is_e8m0,
+            tma_aligned_scales=self.is_tma_aligned,
            compile_native=False,
        )

@@ -367,8 +370,11 @@ class MatcherQuantFP8(MatcherCustomOp):
        )

        if self.quant_key.scale.group_shape.is_per_group():
-            assert scale is None
-            scale = self.make_scale(input, transposed=self.has_col_major_scales)
+            # for tma_aligned, the scale must be passed to forward_custom
+            # tma_aligned fusion then matches by custom op arguments
+            if not self.is_tma_aligned:
+                assert scale is None
+                scale = self.make_scale(input, transposed=self.has_col_major_scales)

            finfo = torch.finfo(self.quant_key.dtype)
            fp8_min = finfo.min
@@ -384,6 +390,8 @@ class MatcherQuantFP8(MatcherCustomOp):
                fp8_min=fp8_min,
                fp8_max=fp8_max,
                scale_ue8m0=self.is_e8m0,
+                dummy_is_scale_transposed=self.has_col_major_scales,
+                dummy_is_tma_aligned=self.is_tma_aligned,
            )
            return result, scale

@@ -121,6 +121,7 @@ class RMSNormQuantPattern:
        key: FusedRMSQuantKey,
        has_col_major_scales: bool = False,
        is_e8m0: bool = False,
+        is_tma_aligned: bool = False,
    ) -> None:
        self.epsilon = epsilon
        self.quant_dtype = key.quant.dtype
@@ -136,7 +137,10 @@ class RMSNormQuantPattern:
            else MatcherFusedAddRMSNorm(epsilon)
        )
        self.quant_matcher = MatcherQuantFP8(
-            key.quant, has_col_major_scales=has_col_major_scales, is_e8m0=is_e8m0
+            key.quant,
+            has_col_major_scales=has_col_major_scales,
+            is_e8m0=is_e8m0,
+            is_tma_aligned=is_tma_aligned,
        )


@@ -262,8 +266,9 @@ class FusedAddRMSNormGroupQuantPattern(RMSNormQuantPattern):
        quant_dtype: torch.dtype,
        group_shape: GroupShape,
        symmetric: bool = True,
-        has_col_major_scales: bool = False,
        is_e8m0: bool = False,
+        has_col_major_scales: bool = True,
+        is_tma_aligned: bool = True,
    ) -> None:
        scale = ScaleDesc(torch.float32, False, group_shape)
        key = FusedRMSQuantKey(
@@ -271,29 +276,63 @@ class FusedAddRMSNormGroupQuantPattern(RMSNormQuantPattern):
            quant=QuantKey(dtype=quant_dtype, scale=scale, symmetric=symmetric),
        )
        self.group_shape = group_shape
-        self.has_col_major_scales = has_col_major_scales
        self.is_e8m0 = is_e8m0
+        self.has_col_major_scales = has_col_major_scales
+        self.is_tma_aligned = is_tma_aligned
        super().__init__(
-            epsilon, key, has_col_major_scales=has_col_major_scales, is_e8m0=is_e8m0
+            epsilon,
+            key,
+            has_col_major_scales=has_col_major_scales,
+            is_e8m0=is_e8m0,
+            is_tma_aligned=is_tma_aligned,
        )

    def register(self, pm_pass: PatternMatcherPass) -> None:
        def pattern(
-            input: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor
+            input: torch.Tensor,
+            weight: torch.Tensor,
+            residual: torch.Tensor,
+            scale: torch.Tensor,
        ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
            result_rms, residual = self.rmsnorm_matcher(input, weight, residual)
-            result, scale = self.quant_matcher(result_rms)
+            result = torch.empty(
+                result_rms.shape,
+                device=result_rms.device,
+                dtype=self.quant_matcher.quant_key.dtype,
+            )
+            assert scale is not None
+            finfo = torch.finfo(self.quant_matcher.quant_key.dtype)
+            fp8_min = finfo.min
+            fp8_max = finfo.max
+
+            _, result, scale = auto_functionalized(
+                self.quant_matcher.QUANT_OP,
+                input=result_rms,
+                output_q=result,
+                output_s=scale,
+                group_size=self.quant_matcher.quant_key.scale.group_shape[1],
+                eps=1e-10,
+                fp8_min=fp8_min,
+                fp8_max=fp8_max,
+                scale_ue8m0=self.quant_matcher.is_e8m0,
+                dummy_is_scale_transposed=self.has_col_major_scales,
+                dummy_is_tma_aligned=self.is_tma_aligned,
+            )
+
            return result, residual, scale

        def replacement(
-            input: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor
+            input: torch.Tensor,
+            weight: torch.Tensor,
+            residual: torch.Tensor,
+            scale: torch.Tensor,
        ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
            # In case we're matching native rms-norm, conversions might be
            # optimized out. We convert here just to be safe.
            input = input.to(dtype=self.model_dtype)

            result = torch.empty_like(input, dtype=self.quant_dtype)
-            scale = self.quant_matcher.make_scale(input, self.has_col_major_scales)
+
            at = auto_functionalized(
                self.FUSED_OP,
                result=result,
@@ -310,10 +349,12 @@ class FusedAddRMSNormGroupQuantPattern(RMSNormQuantPattern):
            # result, residual, scale
            return at[1], at[3], at[2]

+        scale = self.quant_matcher.empty_f32(1, 1)
+
        pm.register_replacement(
            pattern,
            replacement,
-            self.rmsnorm_matcher.inputs(),
+            self.rmsnorm_matcher.inputs() + [scale],
            pm.fwd_only,
            pm_pass,
        )
@@ -326,8 +367,9 @@ class RMSNormGroupQuantPattern(RMSNormQuantPattern):
        quant_dtype: torch.dtype,
        group_shape: GroupShape,
        symmetric: bool = True,
-        has_col_major_scales: bool = False,
        is_e8m0: bool = False,
+        has_col_major_scales: bool = True,
+        is_tma_aligned: bool = True,
    ) -> None:
        scale = ScaleDesc(torch.float32, False, group_shape)
        key = FusedRMSQuantKey(
@@ -335,29 +377,55 @@ class RMSNormGroupQuantPattern(RMSNormQuantPattern):
            quant=QuantKey(dtype=quant_dtype, scale=scale, symmetric=symmetric),
        )
        self.group_shape = group_shape
+        self.has_col_major_scales = has_col_major_scales
+        self.is_tma_aligned = is_tma_aligned
        super().__init__(
-            epsilon, key, has_col_major_scales=has_col_major_scales, is_e8m0=is_e8m0
+            epsilon,
+            key,
+            has_col_major_scales=self.has_col_major_scales,
+            is_e8m0=is_e8m0,
+            is_tma_aligned=is_tma_aligned,
        )

    def register(self, pm_pass: PatternMatcherPass) -> None:
        def pattern(
-            input: torch.Tensor, weight: torch.Tensor
+            input: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor
        ) -> tuple[torch.Tensor, torch.Tensor]:
            result_rms = self.rmsnorm_matcher(input, weight)
-            result, scale = self.quant_matcher(result_rms)
+            result = torch.empty(
+                result_rms.shape,
+                device=result_rms.device,
+                dtype=self.quant_matcher.quant_key.dtype,
+            )
+            assert scale is not None
+            finfo = torch.finfo(self.quant_matcher.quant_key.dtype)
+            fp8_min = finfo.min
+            fp8_max = finfo.max
+
+            _, result, scale = auto_functionalized(
+                self.quant_matcher.QUANT_OP,
+                input=result_rms,
+                output_q=result,
+                output_s=scale,
+                group_size=self.quant_matcher.quant_key.scale.group_shape[1],
+                eps=1e-10,
+                fp8_min=fp8_min,
+                fp8_max=fp8_max,
+                scale_ue8m0=self.quant_matcher.is_e8m0,
+                dummy_is_scale_transposed=self.has_col_major_scales,
+                dummy_is_tma_aligned=self.is_tma_aligned,
+            )
+
            return result, scale

        def replacement(
-            input: torch.Tensor, weight: torch.Tensor
+            input: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor
        ) -> tuple[torch.Tensor, torch.Tensor]:
            # In case we're matching native rms-norm, conversions might be
            # optimized out. We convert here just to be safe.
            input = input.to(dtype=self.model_dtype)

            result = torch.empty_like(input, dtype=self.quant_dtype)
-            scale = self.quant_matcher.make_scale(
-                input, transposed=self.quant_matcher.has_col_major_scales
-            )
            at = auto_functionalized(
                self.FUSED_OP,
                result=result,
@@ -368,16 +436,18 @@ class RMSNormGroupQuantPattern(RMSNormQuantPattern):
                scale_ub=None,
                residual=None,
                group_size=self.group_shape[1],
-                is_scale_transposed=self.quant_matcher.has_col_major_scales,
+                is_scale_transposed=self.has_col_major_scales,
            )

            # result, scale
            return at[1], at[2]

+        scale = self.quant_matcher.empty_f32(1, 1)
+
        pm.register_replacement(
            pattern,
            replacement,
-            self.rmsnorm_matcher.inputs(),
+            self.rmsnorm_matcher.inputs() + [scale],
            pm.fwd_only,
            pm_pass,
        )
@@ -532,23 +602,26 @@ class RMSNormQuantFusionPass(VllmPatternMatcherPass):
                for group_shape in [GroupShape(1, 128), GroupShape(1, 64)]:
                    for has_col_major_scales in [True, False]:
                        for is_e8m0 in [True, False]:
-                            # Fuse fused_add_rms_norm + fp8 group quant
-                            FusedAddRMSNormGroupQuantPattern(
-                                epsilon,
-                                FP8_DTYPE,
-                                group_shape=group_shape,
-                                has_col_major_scales=has_col_major_scales,
-                                is_e8m0=is_e8m0,
-                            ).register(self.patterns)
+                            for is_tma_aligned in [False, True]:
+                                # Fuse fused_add_rms_norm + fp8 group quant
+                                FusedAddRMSNormGroupQuantPattern(
+                                    epsilon,
+                                    FP8_DTYPE,
+                                    group_shape=group_shape,
+                                    is_e8m0=is_e8m0,
+                                    has_col_major_scales=has_col_major_scales,
+                                    is_tma_aligned=is_tma_aligned,
+                                ).register(self.patterns)

-                            # Fuse rms_norm + fp8 group quant
-                            RMSNormGroupQuantPattern(
-                                epsilon,
-                                FP8_DTYPE,
-                                group_shape=group_shape,
-                                has_col_major_scales=has_col_major_scales,
-                                is_e8m0=is_e8m0,
-                            ).register(self.patterns)
+                                # Fuse rms_norm + fp8 group quant
+                                RMSNormGroupQuantPattern(
+                                    epsilon,
+                                    FP8_DTYPE,
+                                    group_shape=group_shape,
+                                    is_e8m0=is_e8m0,
+                                    has_col_major_scales=has_col_major_scales,
+                                    is_tma_aligned=is_tma_aligned,
+                                ).register(self.patterns)

        self.dump_patterns(config, self.patterns)

@@ -5,6 +5,7 @@ import dataclasses
 import io
 import json
 import pickle
+import time
 from collections.abc import Callable
 from pickle import Pickler
 from typing import Any
@@ -164,7 +165,16 @@ class PiecewiseBackend:
        if self.is_last_graph and not self.to_be_compiled_ranges:
            # no specific sizes to compile
            # save the hash of the inductor graph for the next run
+            time_before_saving = time.perf_counter()
            self.vllm_backend.compiler_manager.save_to_file()
+            elapsed = time.perf_counter() - time_before_saving
+            if elapsed > 1:
+                logger.info_once(
+                    "Saved compiler manager cache in %.2f seconds.",
+                    elapsed,
+                    scope="local",
+                )
+
            end_monitoring_torch_compile(self.vllm_config)
            # Call the completion callback (e.g., to save AOT compiled function)
            if self.on_compilation_complete is not None:
@@ -19,6 +19,7 @@ else:

 logger = init_logger(__name__)

+BlockSize = Literal[1, 8, 16, 32, 64, 128, 256]
 CacheDType = Literal[
    "auto",
    "bfloat16",
@@ -38,11 +39,13 @@ KVOffloadingBackend = Literal["native", "lmcache"]
 class CacheConfig:
    """Configuration for the KV cache."""

-    block_size: SkipValidation[int] = None  # type: ignore[assignment]
-    """Size of a contiguous cache block in number of tokens.
+    block_size: SkipValidation[BlockSize] = None  # type: ignore[assignment]
+    """Size of a contiguous cache block in number of tokens. On CUDA devices,
+    only block sizes up to 32 are supported.

-    This is None until `Platform.check_and_update_config()` sets it based on
-    the current platform. Always an int by the time the engine starts."""
+    This config has no static default. If left unspecified by the user, it will
+    be set in `Platform.check_and_update_config()` based on the current
+    platform."""
    gpu_memory_utilization: float = Field(default=0.9, gt=0, le=1)
    """The fraction of GPU memory to be used for the model executor, which can
    range from 0 to 1. For example, a value of 0.5 would imply 50% GPU memory
@@ -305,6 +305,13 @@ class SpeculativeConfig:
        # can not be detected, it will be considered as the "draft_model" by
        # default.

+        # infer method from user args
+        if self.method is None:
+            if self.model in ("ngram", "[ngram]"):
+                self.method = "ngram"
+            else:
+                self.method = "draft_model"
+
        if self.method in get_args(MTPModelTypes) and self.method != "mtp":
            logger.warning(
                "method `%s` is deprecated and replaced with mtp.", self.method
@@ -334,13 +341,6 @@ class SpeculativeConfig:
                    "num_speculative_tokens was provided but without speculative model."
                )

-        # Automatically configure the method for ngram when "model" is used
-        # instead of "method"
-        if self.method is None and (
-            self.model is not None and self.model in ("ngram", "[ngram]")
-        ):
-            self.method = "ngram"
-
        if self.method in ("ngram", "[ngram]"):
            # Unified to "ngram" internally
            self.method = "ngram"
@@ -505,6 +505,13 @@ class SpeculativeConfig:
                        )

                if self.speculative_token_tree is None:
+                    if self.num_speculative_tokens is None:
+                        raise ValueError(
+                            "A speculative model was provided, but neither "
+                            "`speculative_token_tree` nor `num_speculative_tokens` "
+                            "was provided"
+                        )
+
                    # Generate chain of tokens.
                    self.speculative_token_tree = str(
                        [(i + 1) * (0,) for i in range(self.num_speculative_tokens)]
@@ -7,6 +7,7 @@ import enum
 import hashlib
 import inspect
 import json
+import os
 import pathlib
 import textwrap
 from collections.abc import Callable, Mapping, Sequence, Set
@@ -21,6 +22,7 @@ from pydantic.fields import Field as PydanticField
 from pydantic.fields import FieldInfo
 from typing_extensions import dataclass_transform, runtime_checkable

+import vllm.envs as envs
 from vllm.logger import init_logger

 logger = init_logger(__name__)
@@ -380,3 +382,66 @@ def handle_deprecated(

    for new_name in new_names:
        setattr(config, new_name, old_val)
+
+
+def get_from_deprecated_env_if_set(
+    env_name: str,
+    removal_version: str,
+    field_name: str | None = None,
+) -> str | None:
+    """
+    Get value from deprecated environment variable with warning.
+
+    Args:
+        env_name: Name of the deprecated environment variable
+        removal_version: Version when it will be removed
+        field_name: Name of the field to suggest as alternative
+
+    Returns:
+        The environment variable value if set, None otherwise
+    """
+    if envs.is_set(env_name):
+        value = os.environ.get(env_name)
+        alt_msg = f" Please use {field_name} instead." if field_name else ""
+        logger.warning_once(
+            "Using %s environment variable is deprecated and will be removed in %s.%s",
+            env_name,
+            removal_version,
+            alt_msg,
+        )
+        return value
+    return None
+
+
+def set_from_deprecated_env_if_set(
+    config: ConfigT,
+    env_name: str,
+    removal_version: str,
+    field_name: str,
+    to_bool: bool = False,
+    to_int: bool = False,
+) -> None:
+    """
+    Set object field from deprecated environment variable with warning.
+
+    Args:
+        config: Config object to set the field on
+        env_name: Name of the deprecated environment variable
+        removal_version: Version when the env var will be removed
+        field_name: Name of the field to set
+        to_bool: Whether to convert the environment variable value to boolean
+        to_int: Whether to convert the environment variable value to integer
+    Returns:
+        None
+    """
+    if to_bool and to_int:
+        raise ValueError("Cannot convert to both boolean and integer.")
+
+    env_value = get_from_deprecated_env_if_set(env_name, removal_version, field_name)
+    if env_value is not None:
+        field_value: str | bool | int = env_value
+        if to_bool:
+            field_value = env_value.lower() in ("1", "true")
+        elif to_int:
+            field_value = int(env_value)
+        setattr(config, field_name, field_value)
@@ -95,11 +95,16 @@ def enable_norm_fusion(cfg: "VllmConfig") -> bool:


 def enable_act_fusion(cfg: "VllmConfig") -> bool:
-    """Enable if either SiLU+Mul or quant FP8 custom op is active;
-    otherwise Inductor handles fusion."""
-    return cfg.compilation_config.is_custom_op_enabled(
-        "silu_and_mul"
-    ) or cfg.compilation_config.is_custom_op_enabled("quant_fp8")
+    """
+    Enable if either SiLU+Mul or quant FP8 custom op is active;
+    otherwise Inductor handles fusion.
+    Also enable for FP4 models as FP4 quant is always custom so Inductor cannot fuse it.
+    """
+    return (
+        cfg.compilation_config.is_custom_op_enabled("silu_and_mul")
+        or cfg.compilation_config.is_custom_op_enabled("quant_fp8")
+        or (cfg.model_config is not None and cfg.model_config.is_nvfp4_quantized())
+    )


 def enable_allreduce_rms_fusion(cfg: "VllmConfig") -> bool:
@@ -513,8 +513,8 @@ class MessageQueue:
        assert self._is_local_reader, "Only readers can acquire read"
        start_time = time.monotonic()
        n_warning = 1
-        while True:
-            with self.buffer.get_metadata(self.current_idx) as metadata_buffer:
+        with self.buffer.get_metadata(self.current_idx) as metadata_buffer:
+            while True:
                # Memory fence ensures we see the latest writes from the writer.
                # Without this, we may read stale flags from our CPU cache
                # and spin indefinitely even though writer has updated them.
@@ -59,6 +59,7 @@ from vllm.config import (
    get_attr_docs,
 )
 from vllm.config.cache import (
+    BlockSize,
    CacheDType,
    KVOffloadingBackend,
    MambaCacheMode,
@@ -430,7 +431,7 @@ class EngineArgs:
    max_parallel_loading_workers: int | None = (
        ParallelConfig.max_parallel_loading_workers
    )
-    block_size: int = None  # type: ignore[assignment]
+    block_size: BlockSize = CacheConfig.block_size
    enable_prefix_caching: bool | None = None
    prefix_caching_hash_algo: PrefixCachingHashAlgo = (
        CacheConfig.prefix_caching_hash_algo
@@ -40,6 +40,7 @@ class AnthropicContentBlock(BaseModel):
    source: dict[str, Any] | None = None
    # For tool use/result
    id: str | None = None
+    tool_use_id: str | None = None
    name: str | None = None
    input: dict[str, Any] | None = None
    content: str | list[dict[str, Any]] | None = None
@@ -139,7 +139,7 @@ class AnthropicServingMessages(OpenAIServingChat):
                            openai_messages.append(
                                {
                                    "role": "tool",
-                                    "tool_call_id": block.id or "",
+                                    "tool_call_id": block.tool_use_id or "",
                                    "content": str(block.content)
                                    if block.content
                                    else "",
@@ -10,7 +10,7 @@ import cloudpickle
 import torch.nn as nn
 from pydantic import ValidationError
 from tqdm.auto import tqdm
-from typing_extensions import TypeVar
+from typing_extensions import TypeVar, overload

 from vllm.beam_search import (
    BeamSearchInstance,
@@ -94,6 +94,11 @@ if TYPE_CHECKING:

 logger = init_logger(__name__)

+_O = TypeVar(
+    "_O",
+    bound=RequestOutput | PoolingRequestOutput,
+    default=RequestOutput | PoolingRequestOutput,
+)
 _P = TypeVar("_P", bound=SamplingParams | PoolingParams | None)
 _R = TypeVar("_R", default=Any)

@@ -447,17 +452,16 @@ class LLM:
        if sampling_params is None:
            sampling_params = self.get_default_sampling_params()

-        outputs = self._run_completion(
+        return self._run_completion(
            prompts=prompts,
            params=sampling_params,
+            output_type=RequestOutput,
            use_tqdm=use_tqdm,
            lora_request=lora_request,
            tokenization_kwargs=tokenization_kwargs,
            priority=priority,
        )

-        return self.engine_class.validate_outputs(outputs, RequestOutput)
-
    def enqueue(
        self,
        prompts: PromptType | Sequence[PromptType],
@@ -519,75 +523,74 @@ class LLM:
            ),
            params=seq_params,
            lora_requests=seq_lora_requests,
-            tokenization_kwargs=tokenization_kwargs,
            priorities=seq_priority,
        )

        return request_ids

+    @overload
    def wait_for_completion(
        self,
+        *,
        use_tqdm: bool | Callable[..., tqdm] = True,
-    ) -> list[RequestOutput]:
+    ) -> list[RequestOutput | PoolingRequestOutput]: ...
+
+    @overload
+    def wait_for_completion(
+        self,
+        output_type: type[_O] | tuple[type[_O], ...],
+        *,
+        use_tqdm: bool | Callable[..., tqdm] = True,
+    ) -> list[_O]: ...
+
+    def wait_for_completion(
+        self,
+        output_type: type[Any] | tuple[type[Any], ...] | None = None,
+        *,
+        use_tqdm: bool | Callable[..., tqdm] = True,
+    ) -> list[Any]:
        """Wait for all enqueued requests to complete and return results.

        This method processes all requests currently in the engine queue
        and returns their outputs. Use after enqueue() to get results.

        Args:
+            output_type: The expected output type, defaults to RequestOutput.
            use_tqdm: If True, shows a tqdm progress bar.

        Returns:
-            A list of RequestOutput objects for all completed requests.
+            A list of output objects for all completed requests.
        """
-        outputs = self._run_engine(use_tqdm=use_tqdm)
-        return self.engine_class.validate_outputs(outputs, RequestOutput)
+        if output_type is None:
+            output_type = (RequestOutput, PoolingRequestOutput)

-    def _resolve_lora_reqs(
-        self,
-        prompts: Sequence[ProcessorInputs],
-        lora_request: Sequence[LoRARequest | None] | LoRARequest | None,
-    ):
-        lora_config = self.llm_engine.vllm_config.lora_config
-        seq_lora_requests = self._lora_request_to_seq(lora_request, len(prompts))
+        return self._run_engine(output_type, use_tqdm=use_tqdm)

-        if (
-            lora_config is None
-            or not self.model_config.is_multimodal_model
-            or (lora_config and lora_config.default_mm_loras is None)
-        ):
-            return seq_lora_requests
-
-        return [
-            self._resolve_single_prompt_mm_lora(
-                prompt,
-                lora_req,
-                lora_config.default_mm_loras,
-            )
-            for prompt, lora_req in zip(prompts, seq_lora_requests)
-        ]
-
-    def _resolve_single_prompt_mm_lora(
+    def _resolve_mm_lora(
        self,
        prompt: ProcessorInputs,
        lora_request: LoRARequest | None,
-        default_mm_loras: dict[str, str] | None,
-    ):
-        if not default_mm_loras or prompt["type"] != "multimodal":
+    ) -> LoRARequest | None:
+        if prompt["type"] != "multimodal":
+            return lora_request
+
+        lora_config = self.llm_engine.vllm_config.lora_config
+        default_mm_loras = None if lora_config is None else lora_config.default_mm_loras
+        if not default_mm_loras:
            return lora_request

        prompt_modalities = prompt["mm_placeholders"].keys()
        intersection = set(prompt_modalities).intersection(default_mm_loras.keys())
        if not intersection:
            return lora_request
+
        if len(intersection) > 1:
            # TODO: Would be nice to be able to have multiple loras per prompt
            logger.warning(
-                "Multiple modality specific loras were registered and would be"
-                " used by a single prompt consuming several modalities; "
-                " currently we only support one lora per request; as such,"
-                " lora(s) registered with modalities: %s"
-                " will be skipped",
+                "Multiple modality specific loras were registered and would be "
+                "used by a single prompt consuming several modalities; "
+                "currently we only support one lora per request; as such, "
+                "lora(s) registered with modalities: %s will be skipped",
                intersection,
            )
            return lora_request
@@ -765,13 +768,13 @@ class LLM:

                # only runs for one step
                # we don't need to use tqdm here
-                raw_output = self._render_and_run_requests(
+                output = self._render_and_run_requests(
                    prompts=(beam.get_prompt() for beam in all_beams),
                    params=self._params_to_seq(sampling_params, len(all_beams)),
+                    output_type=RequestOutput,
                    lora_requests=[beam.lora_request for beam in all_beams],
                    use_tqdm=False,
                )
-                output = self.engine_class.validate_outputs(raw_output, RequestOutput)

                for (start, end), instance in zip(
                    instance_start_and_end, instances_batch
@@ -1008,9 +1011,10 @@ class LLM:
        if sampling_params is None:
            sampling_params = self.get_default_sampling_params()

-        outputs = self._run_chat(
+        return self._run_chat(
            messages=messages,
            params=sampling_params,
+            output_type=RequestOutput,
            use_tqdm=use_tqdm,
            lora_request=lora_request,
            chat_template=chat_template,
@@ -1023,8 +1027,6 @@ class LLM:
            mm_processor_kwargs=mm_processor_kwargs,
        )

-        return self.engine_class.validate_outputs(outputs, RequestOutput)
-
    def encode(
        self,
        prompts: PromptType | Sequence[PromptType] | DataPrompt,
@@ -1156,19 +1158,16 @@ class LLM:
        outputs = self._run_completion(
            prompts=prompts_seq,
            params=params_seq,
+            output_type=PoolingRequestOutput,
            use_tqdm=use_tqdm,
            lora_request=lora_request,
            tokenization_kwargs=tokenization_kwargs,
        )

-        model_outputs = self.engine_class.validate_outputs(
-            outputs, PoolingRequestOutput
-        )
-
        if use_io_processor:
            # get the post-processed model outputs
            assert self.io_processor is not None
-            processed_outputs = self.io_processor.post_process(model_outputs)
+            processed_outputs = self.io_processor.post_process(outputs)

            return [
                PoolingRequestOutput[Any](
@@ -1181,8 +1180,8 @@ class LLM:
                    finished=True,
                )
            ]
-        else:
-            return model_outputs
+
+        return outputs

    def embed(
        self,
@@ -1374,8 +1373,7 @@ class LLM:
            embed_2=encoded_output_2,
        )

-        items = self.engine_class.validate_outputs(scores, PoolingRequestOutput)
-        return [ScoringRequestOutput.from_base(item) for item in items]
+        return [ScoringRequestOutput.from_base(item) for item in scores]

    def _late_interaction_score(
        self,
@@ -1414,7 +1412,7 @@ class LLM:
                )
            text_2.append(text)

-        encoded_output: list[PoolingRequestOutput] = self.encode(
+        encoded_output = self.encode(
            text_1 + text_2,
            use_tqdm=use_tqdm,
            lora_request=lora_request,
@@ -1423,8 +1421,8 @@ class LLM:
            tokenization_kwargs=tokenization_kwargs,
        )

-        encoded_output_1: list[PoolingRequestOutput] = encoded_output[0 : len(text_1)]
-        encoded_output_2: list[PoolingRequestOutput] = encoded_output[len(text_1) :]
+        encoded_output_1 = encoded_output[0 : len(text_1)]
+        encoded_output_2 = encoded_output[len(text_1) :]

        if len(encoded_output_1) == 1:
            encoded_output_1 = encoded_output_1 * len(encoded_output_2)
@@ -1455,8 +1453,7 @@ class LLM:
                )
            )

-        items = self.engine_class.validate_outputs(scores, PoolingRequestOutput)
-        return [ScoringRequestOutput.from_base(item) for item in items]
+        return [ScoringRequestOutput.from_base(item) for item in scores]

    def _cross_encoding_score(
        self,
@@ -1512,13 +1509,12 @@ class LLM:
        outputs = self._run_completion(
            prompts=prompts,
            params=pooling_params_list,
+            output_type=PoolingRequestOutput,
            use_tqdm=use_tqdm,
            lora_request=lora_request,
        )

-        items = self.engine_class.validate_outputs(outputs, PoolingRequestOutput)
-
-        return [ScoringRequestOutput.from_base(item) for item in items]
+        return [ScoringRequestOutput.from_base(item) for item in outputs]

    def score(
        self,
@@ -1780,6 +1776,7 @@ class LLM:
        params: SamplingParams
        | PoolingParams
        | Sequence[SamplingParams | PoolingParams],
+        output_type: type[_O],
        *,
        use_tqdm: bool | Callable[..., tqdm] = True,
        lora_request: Sequence[LoRARequest] | LoRARequest | None = None,
@@ -1811,9 +1808,9 @@ class LLM:
                )
            ),
            params=seq_params,
+            output_type=output_type,
            use_tqdm=use_tqdm,
            lora_requests=seq_lora_requests,
-            tokenization_kwargs=tokenization_kwargs,
            priorities=seq_priority,
        )

@@ -1824,6 +1821,7 @@ class LLM:
        params: SamplingParams
        | PoolingParams
        | Sequence[SamplingParams | PoolingParams],
+        output_type: type[_O],
        *,
        use_tqdm: bool | Callable[..., tqdm] = True,
        lora_request: Sequence[LoRARequest] | LoRARequest | None = None,
@@ -1870,18 +1868,18 @@ class LLM:
                )
            ),
            params=seq_params,
+            output_type=output_type,
            lora_requests=seq_lora_requests,
            use_tqdm=use_tqdm,
-            tokenization_kwargs=tokenization_kwargs,
        )

    def _render_and_run_requests(
        self,
        prompts: Iterable[ProcessorInputs],
        params: Sequence[SamplingParams | PoolingParams],
+        output_type: type[_O],
        *,
        lora_requests: Sequence[LoRARequest | None] | None = None,
-        tokenization_kwargs: dict[str, Any] | None = None,
        priorities: Sequence[int] | None = None,
        use_tqdm: bool | Callable[..., tqdm] = True,
    ):
@@ -1899,11 +1897,10 @@ class LLM:
            prompts=prompts,
            params=params,
            lora_requests=lora_requests,
-            tokenization_kwargs=tokenization_kwargs,
            priorities=priorities,
        )

-        return self._run_engine(use_tqdm=use_tqdm)
+        return self._run_engine(output_type, use_tqdm=use_tqdm)

    def _render_and_add_requests(
        self,
@@ -1911,7 +1908,6 @@ class LLM:
        params: Sequence[SamplingParams | PoolingParams],
        *,
        lora_requests: Sequence[LoRARequest | None] | None = None,
-        tokenization_kwargs: dict[str, Any] | None = None,
        priorities: Sequence[int] | None = None,
    ) -> list[str]:
        added_request_ids: list[str] = []
@@ -1921,8 +1917,10 @@ class LLM:
                request_id = self._add_request(
                    prompt,
                    params[i],
-                    lora_request=None if lora_requests is None else lora_requests[i],
-                    tokenization_kwargs=tokenization_kwargs,
+                    lora_request=self._resolve_mm_lora(
+                        prompt,
+                        None if lora_requests is None else lora_requests[i],
+                    ),
                    priority=0 if priorities is None else priorities[i],
                )
                added_request_ids.append(request_id)
@@ -1938,7 +1936,6 @@ class LLM:
        prompt: ProcessorInputs,
        params: SamplingParams | PoolingParams,
        lora_request: LoRARequest | None = None,
-        tokenization_kwargs: dict[str, Any] | None = None,
        priority: int = 0,
    ) -> str:
        if isinstance(params, SamplingParams):
@@ -1947,35 +1944,20 @@ class LLM:

        request_id = str(next(self.request_counter))

-        if params.truncate_prompt_tokens is not None:
-            params_type = type(params).__name__
-            warnings.warn(
-                f"The `truncate_prompt_tokens` parameter in `{params_type}` "
-                "is deprecated and will be removed in v0.16. "
-                "Please pass it via `tokenization_kwargs` instead.",
-                DeprecationWarning,
-                stacklevel=2,
-            )
-
-            tokenization_kwargs = merge_kwargs(
-                tokenization_kwargs,
-                dict(truncate_prompt_tokens=params.truncate_prompt_tokens),
-            )
-
        return self.llm_engine.add_request(
            request_id,
            prompt,
            params,
            lora_request=lora_request,
-            tokenization_kwargs=tokenization_kwargs,
            priority=priority,
        )

    def _run_engine(
        self,
+        output_type: type[_O] | tuple[type[_O], ...],
        *,
        use_tqdm: bool | Callable[..., tqdm] = True,
-    ) -> list[RequestOutput | PoolingRequestOutput]:
+    ) -> list[_O]:
        # Initialize tqdm.
        if use_tqdm:
            num_requests = self.llm_engine.get_num_unfinished_requests()
@@ -1988,14 +1970,15 @@ class LLM:
            )

        # Run the engine.
-        outputs: list[RequestOutput | PoolingRequestOutput] = []
+        outputs: list[_O] = []
        total_in_toks = 0
        total_out_toks = 0
        while self.llm_engine.has_unfinished_requests():
            step_outputs = self.llm_engine.step()
            for output in step_outputs:
+                assert isinstance(output, output_type)
                if output.finished:
-                    outputs.append(output)
+                    outputs.append(output)  # type: ignore[arg-type]
                    if use_tqdm:
                        if isinstance(output, RequestOutput):
                            # Calculate tokens only for RequestOutput
@@ -674,3 +674,52 @@ class ChatCompletionRequest(OpenAIBaseModel):
                "Parameter 'cache_salt' must be a non-empty string if provided."
            )
        return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_system_message_content_type(cls, data):
+        """Warn if system messages contain non-text content.
+
+        According to OpenAI API spec, system messages can only be of type
+        'text'. We log a warning instead of rejecting to avoid breaking
+        users who intentionally send multimodal system messages.
+        See: https://platform.openai.com/docs/api-reference/chat/create#chat_create-messages-system_message
+        """
+        if not isinstance(data, dict):
+            return data
+        messages = data.get("messages", [])
+        for msg in messages:
+            # Check if this is a system message
+            if isinstance(msg, dict) and msg.get("role") == "system":
+                content = msg.get("content")
+
+                # If content is a list (multimodal format)
+                if isinstance(content, list):
+                    for part in content:
+                        if isinstance(part, dict):
+                            part_type = part.get("type")
+                            # Infer type when 'type' field is not explicit
+                            if part_type is None:
+                                if "image_url" in part or "image_pil" in part:
+                                    part_type = "image_url"
+                                elif "image_embeds" in part:
+                                    part_type = "image_embeds"
+                                elif "audio_url" in part:
+                                    part_type = "audio_url"
+                                elif "input_audio" in part:
+                                    part_type = "input_audio"
+                                elif "audio_embeds" in part:
+                                    part_type = "audio_embeds"
+                                elif "video_url" in part:
+                                    part_type = "video_url"
+
+                            # Warn about non-text content in system messages
+                            if part_type and part_type != "text":
+                                logger.warning_once(
+                                    "System messages should only contain text "
+                                    "content according to the OpenAI API spec. "
+                                    "Found content type: '%s'.",
+                                    part_type,
+                                )
+
+        return data
@@ -280,7 +280,6 @@ class ParsableContext(ConversationContext):
        self.num_prompt_tokens = 0
        self.num_output_tokens = 0
        self.num_cached_tokens = 0
-        # TODO: num_reasoning_tokens is not implemented yet.
        self.num_reasoning_tokens = 0
        # not implemented yet for ParsableContext
        self.all_turn_metrics: list[TurnMetrics] = []
@@ -308,12 +307,15 @@ class ParsableContext(ConversationContext):

        self.input_messages: list[ResponseRawMessageAndToken] = []
        self.output_messages: list[ResponseRawMessageAndToken] = []
+        self._accumulated_token_ids: list[int] = []

    def append_output(self, output: RequestOutput) -> None:
        self.num_prompt_tokens = len(output.prompt_token_ids or [])
        self.num_cached_tokens = output.num_cached_tokens or 0
        self.num_output_tokens += len(output.outputs[0].token_ids or [])
        self.parser.process(output.outputs[0])
+        output_token_ids = output.outputs[0].token_ids or []
+        self._accumulated_token_ids.extend(output_token_ids)

        # only store if enable_response_messages is True, save memory
        if self.request.enable_response_messages:
@@ -2,36 +2,22 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

 import asyncio
-import json
 import time
 import uuid
 from collections import deque
 from collections.abc import AsyncGenerator, AsyncIterator, Callable, Sequence
 from contextlib import AsyncExitStack
 from copy import copy
-from dataclasses import dataclass, replace
+from dataclasses import replace
 from http import HTTPStatus
 from typing import Final

 import jinja2
 from fastapi import Request
 from openai.types.responses import (
-    ResponseCodeInterpreterCallCodeDeltaEvent,
-    ResponseCodeInterpreterCallCodeDoneEvent,
-    ResponseCodeInterpreterCallCompletedEvent,
-    ResponseCodeInterpreterCallInProgressEvent,
-    ResponseCodeInterpreterCallInterpretingEvent,
-    ResponseCodeInterpreterToolCallParam,
    ResponseContentPartAddedEvent,
    ResponseContentPartDoneEvent,
-    ResponseFunctionCallArgumentsDeltaEvent,
-    ResponseFunctionCallArgumentsDoneEvent,
    ResponseFunctionToolCall,
-    ResponseFunctionWebSearch,
-    ResponseMcpCallArgumentsDeltaEvent,
-    ResponseMcpCallArgumentsDoneEvent,
-    ResponseMcpCallCompletedEvent,
-    ResponseMcpCallInProgressEvent,
    ResponseOutputItem,
    ResponseOutputItemAddedEvent,
    ResponseOutputItemDoneEvent,
@@ -43,13 +29,8 @@ from openai.types.responses import (
    ResponseStatus,
    ResponseTextDeltaEvent,
    ResponseTextDoneEvent,
-    ResponseWebSearchCallCompletedEvent,
-    ResponseWebSearchCallInProgressEvent,
-    ResponseWebSearchCallSearchingEvent,
-    response_function_web_search,
    response_text_delta_event,
 )
-from openai.types.responses.response_output_item import McpCall
 from openai.types.responses.response_output_text import Logprob, LogprobTopLogprob
 from openai.types.responses.response_reasoning_item import (
    Content as ResponseReasoningTextContent,
@@ -102,13 +83,17 @@ from vllm.entrypoints.openai.responses.protocol import (
    ResponseCreatedEvent,
    ResponseInProgressEvent,
    ResponseInputOutputMessage,
-    ResponseReasoningPartAddedEvent,
-    ResponseReasoningPartDoneEvent,
    ResponsesRequest,
    ResponsesResponse,
    ResponseUsage,
    StreamingResponsesResponse,
 )
+from vllm.entrypoints.openai.responses.streaming_events import (
+    HarmonyStreamingState,
+    emit_content_delta_events,
+    emit_previous_item_done_events,
+    emit_tool_action_events,
+)
 from vllm.entrypoints.openai.responses.utils import (
    construct_input_messages,
    construct_tool_dicts,
@@ -129,23 +114,6 @@ from vllm.utils import random_uuid
 logger = init_logger(__name__)


-@dataclass
-class HarmonyStreamingState:
-    """Mutable state for harmony streaming event processing."""
-
-    current_content_index: int = -1
-    current_output_index: int = 0
-    current_item_id: str = ""
-    sent_output_item_added: bool = False
-    is_first_function_call_delta: bool = False
-
-    def reset_for_new_item(self) -> None:
-        """Reset state when expecting a new output item."""
-        self.current_output_index += 1
-        self.sent_output_item_added = False
-        self.is_first_function_call_delta = False
-
-
 def _extract_allowed_tools_from_mcp_requests(
    tools: list[Tool],
 ) -> dict[str, list[str] | None]:
@@ -759,6 +727,19 @@ class OpenAIServingResponses(OpenAIServing):
        num_generated_tokens = context.num_output_tokens
        num_cached_tokens = context.num_cached_tokens
        num_reasoning_tokens = context.num_reasoning_tokens
+        # For text-based reasoning parsers (e.g., <think>...</think>),
+        # HarmonyContext already counts reasoning tokens via channels.
+        # For Simple/Parsable contexts, derive reasoning_tokens from
+        # accumulated output token IDs using the parser if not already set.
+        if (
+            num_reasoning_tokens == 0
+            and self.parser is not None
+            and self.parser.reasoning_parser_cls is not None
+            and isinstance(context, (SimpleContext, ParsableContext))
+        ):
+            reasoning_parser = self.parser.reasoning_parser_cls(tokenizer)
+            accumulated = getattr(context, "_accumulated_token_ids", []) or []
+            num_reasoning_tokens = reasoning_parser.count_reasoning_tokens(accumulated)

        usage = ResponseUsage(
            input_tokens=num_prompt_tokens,
@@ -804,26 +785,6 @@ class OpenAIServingResponses(OpenAIServing):
                    self.response_store[response.id] = response
        return response

-    def _is_mcp_tool_by_namespace(self, recipient: str | None) -> bool:
-        """
-        Determine if a tool call is an MCP tool based on recipient prefix.
-
-        - Tools starting with "functions." are function calls
-        - Everything else is an MCP tool
-        """
-        if recipient is None:
-            return False
-
-        # Function calls have "functions." prefix
-        # Everything else is an MCP tool
-        return not recipient.startswith("functions.")
-
-    _TOOL_NAME_TO_MCP_SERVER_LABEL: Final[dict[str, str]] = {
-        "python": "code_interpreter",
-        "container": "container",
-        "browser": "web_search_preview",
-    }
-
    def _topk_logprobs(
        self,
        logprobs: dict[int, SampleLogprob],
@@ -1592,816 +1553,6 @@ class OpenAIServingResponses(OpenAIServing):
                    )
                )

-    def _emit_function_call_done_events(
-        self,
-        previous_item,
-        state: HarmonyStreamingState,
-    ) -> list[StreamingResponsesResponse]:
-        """Emit events when a function call completes."""
-        function_name = previous_item.recipient[len("functions.") :]
-        events = []
-        events.append(
-            ResponseFunctionCallArgumentsDoneEvent(
-                type="response.function_call_arguments.done",
-                arguments=previous_item.content[0].text,
-                name=function_name,
-                item_id=state.current_item_id,
-                output_index=state.current_output_index,
-                sequence_number=-1,
-            )
-        )
-        function_call_item = ResponseFunctionToolCall(
-            type="function_call",
-            arguments=previous_item.content[0].text,
-            name=function_name,
-            item_id=state.current_item_id,
-            output_index=state.current_output_index,
-            sequence_number=-1,
-            call_id=f"fc_{random_uuid()}",
-            status="completed",
-        )
-        events.append(
-            ResponseOutputItemDoneEvent(
-                type="response.output_item.done",
-                sequence_number=-1,
-                output_index=state.current_output_index,
-                item=function_call_item,
-            )
-        )
-        return events
-
-    def _emit_mcp_call_done_events(
-        self,
-        previous_item,
-        state: HarmonyStreamingState,
-    ) -> list[StreamingResponsesResponse]:
-        """Emit events when an MCP tool call completes."""
-        server_label = self._TOOL_NAME_TO_MCP_SERVER_LABEL.get(
-            previous_item.recipient, previous_item.recipient
-        )
-        events = []
-        events.append(
-            ResponseMcpCallArgumentsDoneEvent(
-                type="response.mcp_call_arguments.done",
-                arguments=previous_item.content[0].text,
-                name=previous_item.recipient,
-                item_id=state.current_item_id,
-                output_index=state.current_output_index,
-                sequence_number=-1,
-            )
-        )
-        events.append(
-            ResponseMcpCallCompletedEvent(
-                type="response.mcp_call.completed",
-                sequence_number=-1,
-                output_index=state.current_output_index,
-                item_id=state.current_item_id,
-            )
-        )
-        events.append(
-            ResponseOutputItemDoneEvent(
-                type="response.output_item.done",
-                sequence_number=-1,
-                output_index=state.current_output_index,
-                item=McpCall(
-                    type="mcp_call",
-                    arguments=previous_item.content[0].text,
-                    name=previous_item.recipient,
-                    id=state.current_item_id,
-                    server_label=server_label,
-                    status="completed",
-                ),
-            )
-        )
-        return events
-
-    def _emit_reasoning_done_events(
-        self,
-        previous_item,
-        state: HarmonyStreamingState,
-    ) -> list[StreamingResponsesResponse]:
-        """Emit events when a reasoning (analysis) item completes."""
-        content = ResponseReasoningTextContent(
-            text=previous_item.content[0].text,
-            type="reasoning_text",
-        )
-        reasoning_item = ResponseReasoningItem(
-            type="reasoning",
-            content=[content],
-            status="completed",
-            id=state.current_item_id,
-            summary=[],
-        )
-        events = []
-        events.append(
-            ResponseReasoningTextDoneEvent(
-                type="response.reasoning_text.done",
-                item_id=state.current_item_id,
-                sequence_number=-1,
-                output_index=state.current_output_index,
-                content_index=state.current_content_index,
-                text=previous_item.content[0].text,
-            )
-        )
-        events.append(
-            ResponseReasoningPartDoneEvent(
-                type="response.reasoning_part.done",
-                sequence_number=-1,
-                item_id=state.current_item_id,
-                output_index=state.current_output_index,
-                content_index=state.current_content_index,
-                part=content,
-            )
-        )
-        events.append(
-            ResponseOutputItemDoneEvent(
-                type="response.output_item.done",
-                sequence_number=-1,
-                output_index=state.current_output_index,
-                item=reasoning_item,
-            )
-        )
-        return events
-
-    def _emit_text_output_done_events(
-        self,
-        previous_item,
-        state: HarmonyStreamingState,
-    ) -> list[StreamingResponsesResponse]:
-        """Emit events when a final text output item completes."""
-        text_content = ResponseOutputText(
-            type="output_text",
-            text=previous_item.content[0].text,
-            annotations=[],
-        )
-        events = []
-        events.append(
-            ResponseTextDoneEvent(
-                type="response.output_text.done",
-                sequence_number=-1,
-                output_index=state.current_output_index,
-                content_index=state.current_content_index,
-                text=previous_item.content[0].text,
-                logprobs=[],
-                item_id=state.current_item_id,
-            )
-        )
-        events.append(
-            ResponseContentPartDoneEvent(
-                type="response.content_part.done",
-                sequence_number=-1,
-                item_id=state.current_item_id,
-                output_index=state.current_output_index,
-                content_index=state.current_content_index,
-                part=text_content,
-            )
-        )
-        events.append(
-            ResponseOutputItemDoneEvent(
-                type="response.output_item.done",
-                sequence_number=-1,
-                output_index=state.current_output_index,
-                item=ResponseOutputMessage(
-                    id=state.current_item_id,
-                    type="message",
-                    role="assistant",
-                    content=[text_content],
-                    status="completed",
-                ),
-            )
-        )
-        return events
-
-    def _emit_previous_item_done_events(
-        self,
-        previous_item,
-        state: HarmonyStreamingState,
-    ) -> list[StreamingResponsesResponse]:
-        """Emit done events for the previous item when expecting a new start."""
-        if previous_item.recipient is not None:
-            # Deal with tool call
-            if previous_item.recipient.startswith("functions."):
-                return self._emit_function_call_done_events(previous_item, state)
-            elif (
-                self._is_mcp_tool_by_namespace(previous_item.recipient)
-                and state.current_item_id is not None
-                and state.current_item_id.startswith("mcp_")
-            ):
-                return self._emit_mcp_call_done_events(previous_item, state)
-        elif previous_item.channel == "analysis":
-            return self._emit_reasoning_done_events(previous_item, state)
-        elif previous_item.channel == "final":
-            return self._emit_text_output_done_events(previous_item, state)
-        return []
-
-    def _emit_final_channel_delta_events(
-        self,
-        ctx: StreamingHarmonyContext,
-        state: HarmonyStreamingState,
-    ) -> list[StreamingResponsesResponse]:
-        """Emit events for final channel text delta streaming."""
-        events = []
-        if not state.sent_output_item_added:
-            state.sent_output_item_added = True
-            state.current_item_id = f"msg_{random_uuid()}"
-            events.append(
-                ResponseOutputItemAddedEvent(
-                    type="response.output_item.added",
-                    sequence_number=-1,
-                    output_index=state.current_output_index,
-                    item=ResponseOutputMessage(
-                        id=state.current_item_id,
-                        type="message",
-                        role="assistant",
-                        content=[],
-                        status="in_progress",
-                    ),
-                )
-            )
-            state.current_content_index += 1
-            events.append(
-                ResponseContentPartAddedEvent(
-                    type="response.content_part.added",
-                    sequence_number=-1,
-                    output_index=state.current_output_index,
-                    item_id=state.current_item_id,
-                    content_index=state.current_content_index,
-                    part=ResponseOutputText(
-                        type="output_text",
-                        text="",
-                        annotations=[],
-                        logprobs=[],
-                    ),
-                )
-            )
-        events.append(
-            ResponseTextDeltaEvent(
-                type="response.output_text.delta",
-                sequence_number=-1,
-                content_index=state.current_content_index,
-                output_index=state.current_output_index,
-                item_id=state.current_item_id,
-                delta=ctx.last_content_delta,
-                # TODO, use logprobs from ctx.last_request_output
-                logprobs=[],
-            )
-        )
-        return events
-
-    def _emit_analysis_channel_delta_events(
-        self,
-        ctx: StreamingHarmonyContext,
-        state: HarmonyStreamingState,
-    ) -> list[StreamingResponsesResponse]:
-        """Emit events for analysis channel reasoning delta streaming."""
-        events = []
-        if not state.sent_output_item_added:
-            state.sent_output_item_added = True
-            state.current_item_id = f"msg_{random_uuid()}"
-            events.append(
-                ResponseOutputItemAddedEvent(
-                    type="response.output_item.added",
-                    sequence_number=-1,
-                    output_index=state.current_output_index,
-                    item=ResponseReasoningItem(
-                        type="reasoning",
-                        id=state.current_item_id,
-                        summary=[],
-                        status="in_progress",
-                    ),
-                )
-            )
-            state.current_content_index += 1
-            events.append(
-                ResponseReasoningPartAddedEvent(
-                    type="response.reasoning_part.added",
-                    sequence_number=-1,
-                    output_index=state.current_output_index,
-                    item_id=state.current_item_id,
-                    content_index=state.current_content_index,
-                    part=ResponseReasoningTextContent(
-                        text="",
-                        type="reasoning_text",
-                    ),
-                )
-            )
-        events.append(
-            ResponseReasoningTextDeltaEvent(
-                type="response.reasoning_text.delta",
-                item_id=state.current_item_id,
-                output_index=state.current_output_index,
-                content_index=state.current_content_index,
-                delta=ctx.last_content_delta,
-                sequence_number=-1,
-            )
-        )
-        return events
-
-    def _emit_mcp_tool_delta_events(
-        self,
-        ctx: StreamingHarmonyContext,
-        state: HarmonyStreamingState,
-        recipient: str,
-    ) -> list[StreamingResponsesResponse]:
-        """Emit events for MCP tool delta streaming."""
-        server_label = self._TOOL_NAME_TO_MCP_SERVER_LABEL.get(recipient, recipient)
-        events = []
-        if not state.sent_output_item_added:
-            state.sent_output_item_added = True
-            state.current_item_id = f"mcp_{random_uuid()}"
-            events.append(
-                ResponseOutputItemAddedEvent(
-                    type="response.output_item.added",
-                    sequence_number=-1,
-                    output_index=state.current_output_index,
-                    item=McpCall(
-                        type="mcp_call",
-                        id=state.current_item_id,
-                        name=recipient,
-                        arguments="",
-                        server_label=server_label,
-                        status="in_progress",
-                    ),
-                )
-            )
-            events.append(
-                ResponseMcpCallInProgressEvent(
-                    type="response.mcp_call.in_progress",
-                    sequence_number=-1,
-                    output_index=state.current_output_index,
-                    item_id=state.current_item_id,
-                )
-            )
-        events.append(
-            ResponseMcpCallArgumentsDeltaEvent(
-                type="response.mcp_call_arguments.delta",
-                sequence_number=-1,
-                output_index=state.current_output_index,
-                item_id=state.current_item_id,
-                delta=ctx.last_content_delta,
-            )
-        )
-        return events
-
-    def _emit_code_interpreter_delta_events(
-        self,
-        ctx: StreamingHarmonyContext,
-        state: HarmonyStreamingState,
-    ) -> list[StreamingResponsesResponse]:
-        """Emit events for code interpreter delta streaming."""
-        events = []
-        if not state.sent_output_item_added:
-            state.sent_output_item_added = True
-            state.current_item_id = f"tool_{random_uuid()}"
-            events.append(
-                ResponseOutputItemAddedEvent(
-                    type="response.output_item.added",
-                    sequence_number=-1,
-                    output_index=state.current_output_index,
-                    item=ResponseCodeInterpreterToolCallParam(
-                        type="code_interpreter_call",
-                        id=state.current_item_id,
-                        code=None,
-                        container_id="auto",
-                        outputs=None,
-                        status="in_progress",
-                    ),
-                )
-            )
-            events.append(
-                ResponseCodeInterpreterCallInProgressEvent(
-                    type="response.code_interpreter_call.in_progress",
-                    sequence_number=-1,
-                    output_index=state.current_output_index,
-                    item_id=state.current_item_id,
-                )
-            )
-        events.append(
-            ResponseCodeInterpreterCallCodeDeltaEvent(
-                type="response.code_interpreter_call_code.delta",
-                sequence_number=-1,
-                output_index=state.current_output_index,
-                item_id=state.current_item_id,
-                delta=ctx.last_content_delta,
-            )
-        )
-        return events
-
-    def _emit_mcp_prefix_delta_events(
-        self,
-        ctx: StreamingHarmonyContext,
-        state: HarmonyStreamingState,
-    ) -> list[StreamingResponsesResponse]:
-        """Emit events for MCP prefix (mcp.*) delta streaming."""
-        events = []
-        if not state.sent_output_item_added:
-            state.sent_output_item_added = True
-            state.current_item_id = f"mcp_{random_uuid()}"
-            mcp_name = ctx.parser.current_recipient[len("mcp.") :]
-
-            events.append(
-                ResponseOutputItemAddedEvent(
-                    type="response.output_item.added",
-                    sequence_number=-1,
-                    output_index=state.current_output_index,
-                    item=McpCall(
-                        type="mcp_call",
-                        id=state.current_item_id,
-                        name=mcp_name,
-                        arguments="",
-                        server_label=mcp_name,
-                        status="in_progress",
-                    ),
-                )
-            )
-            events.append(
-                ResponseMcpCallInProgressEvent(
-                    type="response.mcp_call.in_progress",
-                    sequence_number=-1,
-                    output_index=state.current_output_index,
-                    item_id=state.current_item_id,
-                )
-            )
-
-        events.append(
-            ResponseMcpCallArgumentsDeltaEvent(
-                type="response.mcp_call_arguments.delta",
-                sequence_number=-1,
-                output_index=state.current_output_index,
-                item_id=state.current_item_id,
-                delta=ctx.last_content_delta,
-            )
-        )
-        return events
-
-    def _emit_content_delta_events(
-        self,
-        ctx: StreamingHarmonyContext,
-        state: HarmonyStreamingState,
-    ) -> list[StreamingResponsesResponse]:
-        """Emit events for content delta streaming based on channel type."""
-        if not ctx.last_content_delta:
-            return []
-
-        if (
-            ctx.parser.current_channel == "final"
-            and ctx.parser.current_recipient is None
-        ):
-            return self._emit_final_channel_delta_events(ctx, state)
-        elif (
-            ctx.parser.current_channel == "analysis"
-            and ctx.parser.current_recipient is None
-        ):
-            return self._emit_analysis_channel_delta_events(ctx, state)
-        # built-in tools will be triggered on the analysis channel
-        # However, occasionally built-in tools will
-        # still be output to commentary.
-        elif (
-            ctx.parser.current_channel == "commentary"
-            or ctx.parser.current_channel == "analysis"
-        ) and ctx.parser.current_recipient is not None:
-            recipient = ctx.parser.current_recipient
-            # Check for function calls first - they have their own event handling
-            if recipient.startswith("functions."):
-                return self._emit_function_call_delta_events(ctx, state)
-            is_mcp_tool = self._is_mcp_tool_by_namespace(recipient)
-            if is_mcp_tool:
-                return self._emit_mcp_tool_delta_events(ctx, state, recipient)
-            else:
-                return self._emit_code_interpreter_delta_events(ctx, state)
-        elif (
-            (
-                ctx.parser.current_channel == "commentary"
-                or ctx.parser.current_channel == "analysis"
-            )
-            and ctx.parser.current_recipient is not None
-            and ctx.parser.current_recipient.startswith("mcp.")
-        ):
-            return self._emit_mcp_prefix_delta_events(ctx, state)
-
-        return []
-
-    def _emit_browser_tool_events(
-        self,
-        previous_item,
-        state: HarmonyStreamingState,
-    ) -> list[StreamingResponsesResponse]:
-        """Emit events for browser tool calls (web search)."""
-        function_name = previous_item.recipient[len("browser.") :]
-        parsed_args = json.loads(previous_item.content[0].text)
-        action = None
-
-        if function_name == "search":
-            action = response_function_web_search.ActionSearch(
-                type="search",
-                query=parsed_args["query"],
-            )
-        elif function_name == "open":
-            action = response_function_web_search.ActionOpenPage(
-                type="open_page",
-                # TODO: translate to url
-                url=f"cursor:{parsed_args.get('cursor', '')}",
-            )
-        elif function_name == "find":
-            action = response_function_web_search.ActionFind(
-                type="find",
-                pattern=parsed_args["pattern"],
-                # TODO: translate to url
-                url=f"cursor:{parsed_args.get('cursor', '')}",
-            )
-        else:
-            raise ValueError(f"Unknown function name: {function_name}")
-
-        state.current_item_id = f"tool_{random_uuid()}"
-        events = []
-        events.append(
-            ResponseOutputItemAddedEvent(
-                type="response.output_item.added",
-                sequence_number=-1,
-                output_index=state.current_output_index,
-                item=response_function_web_search.ResponseFunctionWebSearch(
-                    # TODO: generate a unique id for web search call
-                    type="web_search_call",
-                    id=state.current_item_id,
-                    action=action,
-                    status="in_progress",
-                ),
-            )
-        )
-        events.append(
-            ResponseWebSearchCallInProgressEvent(
-                type="response.web_search_call.in_progress",
-                sequence_number=-1,
-                output_index=state.current_output_index,
-                item_id=state.current_item_id,
-            )
-        )
-        events.append(
-            ResponseWebSearchCallSearchingEvent(
-                type="response.web_search_call.searching",
-                sequence_number=-1,
-                output_index=state.current_output_index,
-                item_id=state.current_item_id,
-            )
-        )
-        # enqueue
-        events.append(
-            ResponseWebSearchCallCompletedEvent(
-                type="response.web_search_call.completed",
-                sequence_number=-1,
-                output_index=state.current_output_index,
-                item_id=state.current_item_id,
-            )
-        )
-        events.append(
-            ResponseOutputItemDoneEvent(
-                type="response.output_item.done",
-                sequence_number=-1,
-                output_index=state.current_output_index,
-                item=ResponseFunctionWebSearch(
-                    type="web_search_call",
-                    id=state.current_item_id,
-                    action=action,
-                    status="completed",
-                ),
-            )
-        )
-        return events
-
-    def _emit_mcp_tool_completion_events(
-        self,
-        previous_item,
-        state: HarmonyStreamingState,
-    ) -> list[StreamingResponsesResponse]:
-        """Emit events when an MCP tool completes during assistant action turn."""
-        recipient = previous_item.recipient
-        server_label = self._TOOL_NAME_TO_MCP_SERVER_LABEL.get(recipient, recipient)
-        events = []
-        events.append(
-            ResponseMcpCallArgumentsDoneEvent(
-                type="response.mcp_call_arguments.done",
-                sequence_number=-1,
-                output_index=state.current_output_index,
-                item_id=state.current_item_id,
-                arguments=previous_item.content[0].text,
-                name=recipient,
-            )
-        )
-        events.append(
-            ResponseMcpCallCompletedEvent(
-                type="response.mcp_call.completed",
-                sequence_number=-1,
-                output_index=state.current_output_index,
-                item_id=state.current_item_id,
-            )
-        )
-        events.append(
-            ResponseOutputItemDoneEvent(
-                type="response.output_item.done",
-                sequence_number=-1,
-                output_index=state.current_output_index,
-                item=McpCall(
-                    type="mcp_call",
-                    id=state.current_item_id,
-                    name=recipient,
-                    arguments=previous_item.content[0].text,
-                    server_label=server_label,
-                    status="completed",
-                ),
-            )
-        )
-        return events
-
-    def _emit_code_interpreter_completion_events(
-        self,
-        previous_item,
-        state: HarmonyStreamingState,
-    ) -> list[StreamingResponsesResponse]:
-        """Emit events when code interpreter completes."""
-        events = []
-        events.append(
-            ResponseCodeInterpreterCallCodeDoneEvent(
-                type="response.code_interpreter_call_code.done",
-                sequence_number=-1,
-                output_index=state.current_output_index,
-                item_id=state.current_item_id,
-                code=previous_item.content[0].text,
-            )
-        )
-        events.append(
-            ResponseCodeInterpreterCallInterpretingEvent(
-                type="response.code_interpreter_call.interpreting",
-                sequence_number=-1,
-                output_index=state.current_output_index,
-                item_id=state.current_item_id,
-            )
-        )
-        events.append(
-            ResponseCodeInterpreterCallCompletedEvent(
-                type="response.code_interpreter_call.completed",
-                sequence_number=-1,
-                output_index=state.current_output_index,
-                item_id=state.current_item_id,
-            )
-        )
-        events.append(
-            ResponseOutputItemDoneEvent(
-                type="response.output_item.done",
-                sequence_number=-1,
-                output_index=state.current_output_index,
-                item=ResponseCodeInterpreterToolCallParam(
-                    type="code_interpreter_call",
-                    id=state.current_item_id,
-                    code=previous_item.content[0].text,
-                    container_id="auto",
-                    outputs=[],
-                    status="completed",
-                ),
-            )
-        )
-        return events
-
-    def _emit_mcp_prefix_completion_events(
-        self,
-        previous_item,
-        state: HarmonyStreamingState,
-    ) -> list[StreamingResponsesResponse]:
-        """Emit events when an MCP prefix tool (mcp.*) completes."""
-        mcp_name = previous_item.recipient[len("mcp.") :]
-        events = []
-        events.append(
-            ResponseMcpCallArgumentsDoneEvent(
-                type="response.mcp_call_arguments.done",
-                sequence_number=-1,
-                output_index=state.current_output_index,
-                item_id=state.current_item_id,
-                arguments=previous_item.content[0].text,
-                name=mcp_name,
-            )
-        )
-        events.append(
-            ResponseMcpCallCompletedEvent(
-                type="response.mcp_call.completed",
-                sequence_number=-1,
-                output_index=state.current_output_index,
-                item_id=state.current_item_id,
-            )
-        )
-        events.append(
-            ResponseOutputItemDoneEvent(
-                type="response.output_item.done",
-                sequence_number=-1,
-                output_index=state.current_output_index,
-                item=McpCall(
-                    type="mcp_call",
-                    id=state.current_item_id,
-                    name=mcp_name,
-                    arguments=previous_item.content[0].text,
-                    server_label=mcp_name,
-                    status="completed",
-                ),
-            )
-        )
-        return events
-
-    def _emit_tool_action_events(
-        self,
-        ctx: StreamingHarmonyContext,
-        state: HarmonyStreamingState,
-    ) -> list[StreamingResponsesResponse]:
-        """Emit events for tool action turn."""
-        if not ctx.is_assistant_action_turn() or len(ctx.parser.messages) == 0:
-            return []
-
-        events = []
-        previous_item = ctx.parser.messages[-1]
-
-        # Handle browser tool
-        if (
-            self.tool_server is not None
-            and self.tool_server.has_tool("browser")
-            and previous_item.recipient is not None
-            and previous_item.recipient.startswith("browser.")
-        ):
-            events.extend(self._emit_browser_tool_events(previous_item, state))
-
-        # Handle tool completion
-        if (
-            self.tool_server is not None
-            and previous_item.recipient is not None
-            and state.current_item_id is not None
-            and state.sent_output_item_added
-        ):
-            recipient = previous_item.recipient
-            # Handle MCP prefix tool completion first
-            if recipient.startswith("mcp."):
-                events.extend(
-                    self._emit_mcp_prefix_completion_events(previous_item, state)
-                )
-            else:
-                # Handle other MCP tool and code interpreter completion
-                is_mcp_tool = self._is_mcp_tool_by_namespace(
-                    recipient
-                ) and state.current_item_id.startswith("mcp_")
-                if is_mcp_tool:
-                    events.extend(
-                        self._emit_mcp_tool_completion_events(previous_item, state)
-                    )
-                else:
-                    events.extend(
-                        self._emit_code_interpreter_completion_events(
-                            previous_item, state
-                        )
-                    )
-
-        return events
-
-    def _emit_function_call_delta_events(
-        self,
-        ctx: StreamingHarmonyContext,
-        state: HarmonyStreamingState,
-    ) -> list[StreamingResponsesResponse]:
-        """Emit events for developer function calls on commentary channel."""
-        if not (
-            ctx.parser.current_channel == "commentary"
-            and ctx.parser.current_recipient
-            and ctx.parser.current_recipient.startswith("functions.")
-        ):
-            return []
-
-        events = []
-        if state.is_first_function_call_delta is False:
-            state.is_first_function_call_delta = True
-            fc_name = ctx.parser.current_recipient[len("functions.") :]
-            state.current_item_id = f"fc_{random_uuid()}"
-            tool_call_item = ResponseFunctionToolCall(
-                name=fc_name,
-                type="function_call",
-                id=state.current_item_id,
-                call_id=f"call_{random_uuid()}",
-                arguments="",
-                status="in_progress",
-            )
-            events.append(
-                ResponseOutputItemAddedEvent(
-                    type="response.output_item.added",
-                    sequence_number=-1,
-                    output_index=state.current_output_index,
-                    item=tool_call_item,
-                )
-            )
-        # Always emit the delta (including on first call)
-        events.append(
-            ResponseFunctionCallArgumentsDeltaEvent(
-                item_id=state.current_item_id,
-                delta=ctx.last_content_delta,
-                output_index=state.current_output_index,
-                sequence_number=-1,
-                type="response.function_call_arguments.delta",
-            )
-        )
-        return events
-
    async def _process_harmony_streaming_events(
        self,
        request: ResponsesRequest,
@@ -2427,18 +1578,16 @@ class OpenAIServingResponses(OpenAIServing):
            if ctx.is_expecting_start():
                if len(ctx.parser.messages) > 0:
                    previous_item = ctx.parser.messages[-1]
-                    for event in self._emit_previous_item_done_events(
-                        previous_item, state
-                    ):
+                    for event in emit_previous_item_done_events(previous_item, state):
                        yield _increment_sequence_number_and_return(event)
                state.reset_for_new_item()

            # Stream the output of a harmony message
-            for event in self._emit_content_delta_events(ctx, state):
+            for event in emit_content_delta_events(ctx, state):
                yield _increment_sequence_number_and_return(event)

            # Stream tool call outputs
-            for event in self._emit_tool_action_events(ctx, state):
+            for event in emit_tool_action_events(ctx, state, self.tool_server):
                yield _increment_sequence_number_and_return(event)

    async def responses_stream_generator(
@@ -0,0 +1,897 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Streaming SSE event builders for the Responses API.
+
+Pure functions that translate streaming state + delta data into
+OpenAI Response API SSE events. Used by the streaming event
+processors in serving.py.
+"""
+
+import json
+from dataclasses import dataclass
+from typing import Final
+
+from openai.types.responses import (
+    ResponseCodeInterpreterCallCodeDeltaEvent,
+    ResponseCodeInterpreterCallCodeDoneEvent,
+    ResponseCodeInterpreterCallCompletedEvent,
+    ResponseCodeInterpreterCallInProgressEvent,
+    ResponseCodeInterpreterCallInterpretingEvent,
+    ResponseCodeInterpreterToolCallParam,
+    ResponseContentPartAddedEvent,
+    ResponseContentPartDoneEvent,
+    ResponseFunctionCallArgumentsDeltaEvent,
+    ResponseFunctionCallArgumentsDoneEvent,
+    ResponseFunctionToolCall,
+    ResponseFunctionWebSearch,
+    ResponseMcpCallArgumentsDeltaEvent,
+    ResponseMcpCallArgumentsDoneEvent,
+    ResponseMcpCallCompletedEvent,
+    ResponseMcpCallInProgressEvent,
+    ResponseOutputItemAddedEvent,
+    ResponseOutputItemDoneEvent,
+    ResponseOutputMessage,
+    ResponseOutputText,
+    ResponseReasoningItem,
+    ResponseReasoningTextDeltaEvent,
+    ResponseReasoningTextDoneEvent,
+    ResponseTextDeltaEvent,
+    ResponseTextDoneEvent,
+    ResponseWebSearchCallCompletedEvent,
+    ResponseWebSearchCallInProgressEvent,
+    ResponseWebSearchCallSearchingEvent,
+    response_function_web_search,
+)
+from openai.types.responses.response_output_item import McpCall
+from openai.types.responses.response_reasoning_item import (
+    Content as ResponseReasoningTextContent,
+)
+
+from vllm.entrypoints.mcp.tool_server import ToolServer
+from vllm.entrypoints.openai.responses.context import StreamingHarmonyContext
+from vllm.entrypoints.openai.responses.protocol import (
+    ResponseReasoningPartAddedEvent,
+    ResponseReasoningPartDoneEvent,
+    StreamingResponsesResponse,
+)
+from vllm.utils import random_uuid
+
+TOOL_NAME_TO_MCP_SERVER_LABEL: Final[dict[str, str]] = {
+    "python": "code_interpreter",
+    "container": "container",
+    "browser": "web_search_preview",
+}
+
+
+@dataclass
+class HarmonyStreamingState:
+    """Mutable state for harmony streaming event processing."""
+
+    current_content_index: int = -1
+    current_output_index: int = 0
+    current_item_id: str = ""
+    sent_output_item_added: bool = False
+    is_first_function_call_delta: bool = False
+
+    def reset_for_new_item(self) -> None:
+        """Reset state when expecting a new output item."""
+        self.current_output_index += 1
+        self.sent_output_item_added = False
+        self.is_first_function_call_delta = False
+
+
+def is_mcp_tool_by_namespace(recipient: str | None) -> bool:
+    """
+    Determine if a tool call is an MCP tool based on recipient prefix.
+
+    - Tools starting with "functions." are function calls
+    - Everything else is an MCP tool
+    """
+    if recipient is None:
+        return False
+
+    # Function calls have "functions." prefix
+    # Everything else is an MCP tool
+    return not recipient.startswith("functions.")
+
+
+def emit_function_call_done_events(
+    previous_item,
+    state: HarmonyStreamingState,
+) -> list[StreamingResponsesResponse]:
+    """Emit events when a function call completes."""
+    function_name = previous_item.recipient[len("functions.") :]
+    events: list[StreamingResponsesResponse] = []
+    events.append(
+        ResponseFunctionCallArgumentsDoneEvent(
+            type="response.function_call_arguments.done",
+            arguments=previous_item.content[0].text,
+            name=function_name,
+            item_id=state.current_item_id,
+            output_index=state.current_output_index,
+            sequence_number=-1,
+        )
+    )
+    function_call_item = ResponseFunctionToolCall(
+        type="function_call",
+        arguments=previous_item.content[0].text,
+        name=function_name,
+        item_id=state.current_item_id,
+        output_index=state.current_output_index,
+        sequence_number=-1,
+        call_id=f"fc_{random_uuid()}",
+        status="completed",
+    )
+    events.append(
+        ResponseOutputItemDoneEvent(
+            type="response.output_item.done",
+            sequence_number=-1,
+            output_index=state.current_output_index,
+            item=function_call_item,
+        )
+    )
+    return events
+
+
+def emit_mcp_call_done_events(
+    previous_item,
+    state: HarmonyStreamingState,
+) -> list[StreamingResponsesResponse]:
+    """Emit events when an MCP tool call completes."""
+    server_label = TOOL_NAME_TO_MCP_SERVER_LABEL.get(
+        previous_item.recipient, previous_item.recipient
+    )
+    events: list[StreamingResponsesResponse] = []
+    events.append(
+        ResponseMcpCallArgumentsDoneEvent(
+            type="response.mcp_call_arguments.done",
+            arguments=previous_item.content[0].text,
+            name=previous_item.recipient,
+            item_id=state.current_item_id,
+            output_index=state.current_output_index,
+            sequence_number=-1,
+        )
+    )
+    events.append(
+        ResponseMcpCallCompletedEvent(
+            type="response.mcp_call.completed",
+            sequence_number=-1,
+            output_index=state.current_output_index,
+            item_id=state.current_item_id,
+        )
+    )
+    events.append(
+        ResponseOutputItemDoneEvent(
+            type="response.output_item.done",
+            sequence_number=-1,
+            output_index=state.current_output_index,
+            item=McpCall(
+                type="mcp_call",
+                arguments=previous_item.content[0].text,
+                name=previous_item.recipient,
+                id=state.current_item_id,
+                server_label=server_label,
+                status="completed",
+            ),
+        )
+    )
+    return events
+
+
+def emit_reasoning_done_events(
+    previous_item,
+    state: HarmonyStreamingState,
+) -> list[StreamingResponsesResponse]:
+    """Emit events when a reasoning (analysis) item completes."""
+    content = ResponseReasoningTextContent(
+        text=previous_item.content[0].text,
+        type="reasoning_text",
+    )
+    reasoning_item = ResponseReasoningItem(
+        type="reasoning",
+        content=[content],
+        status="completed",
+        id=state.current_item_id,
+        summary=[],
+    )
+    events: list[StreamingResponsesResponse] = []
+    events.append(
+        ResponseReasoningTextDoneEvent(
+            type="response.reasoning_text.done",
+            item_id=state.current_item_id,
+            sequence_number=-1,
+            output_index=state.current_output_index,
+            content_index=state.current_content_index,
+            text=previous_item.content[0].text,
+        )
+    )
+    events.append(
+        ResponseReasoningPartDoneEvent(
+            type="response.reasoning_part.done",
+            sequence_number=-1,
+            item_id=state.current_item_id,
+            output_index=state.current_output_index,
+            content_index=state.current_content_index,
+            part=content,
+        )
+    )
+    events.append(
+        ResponseOutputItemDoneEvent(
+            type="response.output_item.done",
+            sequence_number=-1,
+            output_index=state.current_output_index,
+            item=reasoning_item,
+        )
+    )
+    return events
+
+
+def emit_text_output_done_events(
+    previous_item,
+    state: HarmonyStreamingState,
+) -> list[StreamingResponsesResponse]:
+    """Emit events when a final text output item completes."""
+    text_content = ResponseOutputText(
+        type="output_text",
+        text=previous_item.content[0].text,
+        annotations=[],
+    )
+    events: list[StreamingResponsesResponse] = []
+    events.append(
+        ResponseTextDoneEvent(
+            type="response.output_text.done",
+            sequence_number=-1,
+            output_index=state.current_output_index,
+            content_index=state.current_content_index,
+            text=previous_item.content[0].text,
+            logprobs=[],
+            item_id=state.current_item_id,
+        )
+    )
+    events.append(
+        ResponseContentPartDoneEvent(
+            type="response.content_part.done",
+            sequence_number=-1,
+            item_id=state.current_item_id,
+            output_index=state.current_output_index,
+            content_index=state.current_content_index,
+            part=text_content,
+        )
+    )
+    events.append(
+        ResponseOutputItemDoneEvent(
+            type="response.output_item.done",
+            sequence_number=-1,
+            output_index=state.current_output_index,
+            item=ResponseOutputMessage(
+                id=state.current_item_id,
+                type="message",
+                role="assistant",
+                content=[text_content],
+                status="completed",
+            ),
+        )
+    )
+    return events
+
+
+def emit_previous_item_done_events(
+    previous_item,
+    state: HarmonyStreamingState,
+) -> list[StreamingResponsesResponse]:
+    """Emit done events for the previous item when expecting a new start."""
+    if previous_item.recipient is not None:
+        # Deal with tool call
+        if previous_item.recipient.startswith("functions."):
+            return emit_function_call_done_events(previous_item, state)
+        elif (
+            is_mcp_tool_by_namespace(previous_item.recipient)
+            and state.current_item_id is not None
+            and state.current_item_id.startswith("mcp_")
+        ):
+            return emit_mcp_call_done_events(previous_item, state)
+    elif previous_item.channel == "analysis":
+        return emit_reasoning_done_events(previous_item, state)
+    elif previous_item.channel == "final":
+        return emit_text_output_done_events(previous_item, state)
+    return []
+
+
+def emit_final_channel_delta_events(
+    ctx: StreamingHarmonyContext,
+    state: HarmonyStreamingState,
+) -> list[StreamingResponsesResponse]:
+    """Emit events for final channel text delta streaming."""
+    events: list[StreamingResponsesResponse] = []
+    if not state.sent_output_item_added:
+        state.sent_output_item_added = True
+        state.current_item_id = f"msg_{random_uuid()}"
+        events.append(
+            ResponseOutputItemAddedEvent(
+                type="response.output_item.added",
+                sequence_number=-1,
+                output_index=state.current_output_index,
+                item=ResponseOutputMessage(
+                    id=state.current_item_id,
+                    type="message",
+                    role="assistant",
+                    content=[],
+                    status="in_progress",
+                ),
+            )
+        )
+        state.current_content_index += 1
+        events.append(
+            ResponseContentPartAddedEvent(
+                type="response.content_part.added",
+                sequence_number=-1,
+                output_index=state.current_output_index,
+                item_id=state.current_item_id,
+                content_index=state.current_content_index,
+                part=ResponseOutputText(
+                    type="output_text",
+                    text="",
+                    annotations=[],
+                    logprobs=[],
+                ),
+            )
+        )
+    events.append(
+        ResponseTextDeltaEvent(
+            type="response.output_text.delta",
+            sequence_number=-1,
+            content_index=state.current_content_index,
+            output_index=state.current_output_index,
+            item_id=state.current_item_id,
+            delta=ctx.last_content_delta,
+            # TODO, use logprobs from ctx.last_request_output
+            logprobs=[],
+        )
+    )
+    return events
+
+
+def emit_analysis_channel_delta_events(
+    ctx: StreamingHarmonyContext,
+    state: HarmonyStreamingState,
+) -> list[StreamingResponsesResponse]:
+    """Emit events for analysis channel reasoning delta streaming."""
+    events: list[StreamingResponsesResponse] = []
+    if not state.sent_output_item_added:
+        state.sent_output_item_added = True
+        state.current_item_id = f"msg_{random_uuid()}"
+        events.append(
+            ResponseOutputItemAddedEvent(
+                type="response.output_item.added",
+                sequence_number=-1,
+                output_index=state.current_output_index,
+                item=ResponseReasoningItem(
+                    type="reasoning",
+                    id=state.current_item_id,
+                    summary=[],
+                    status="in_progress",
+                ),
+            )
+        )
+        state.current_content_index += 1
+        events.append(
+            ResponseReasoningPartAddedEvent(
+                type="response.reasoning_part.added",
+                sequence_number=-1,
+                output_index=state.current_output_index,
+                item_id=state.current_item_id,
+                content_index=state.current_content_index,
+                part=ResponseReasoningTextContent(
+                    text="",
+                    type="reasoning_text",
+                ),
+            )
+        )
+    events.append(
+        ResponseReasoningTextDeltaEvent(
+            type="response.reasoning_text.delta",
+            item_id=state.current_item_id,
+            output_index=state.current_output_index,
+            content_index=state.current_content_index,
+            delta=ctx.last_content_delta,
+            sequence_number=-1,
+        )
+    )
+    return events
+
+
+def emit_mcp_tool_delta_events(
+    ctx: StreamingHarmonyContext,
+    state: HarmonyStreamingState,
+    recipient: str,
+) -> list[StreamingResponsesResponse]:
+    """Emit events for MCP tool delta streaming."""
+    server_label = TOOL_NAME_TO_MCP_SERVER_LABEL.get(recipient, recipient)
+    events: list[StreamingResponsesResponse] = []
+    if not state.sent_output_item_added:
+        state.sent_output_item_added = True
+        state.current_item_id = f"mcp_{random_uuid()}"
+        events.append(
+            ResponseOutputItemAddedEvent(
+                type="response.output_item.added",
+                sequence_number=-1,
+                output_index=state.current_output_index,
+                item=McpCall(
+                    type="mcp_call",
+                    id=state.current_item_id,
+                    name=recipient,
+                    arguments="",
+                    server_label=server_label,
+                    status="in_progress",
+                ),
+            )
+        )
+        events.append(
+            ResponseMcpCallInProgressEvent(
+                type="response.mcp_call.in_progress",
+                sequence_number=-1,
+                output_index=state.current_output_index,
+                item_id=state.current_item_id,
+            )
+        )
+    events.append(
+        ResponseMcpCallArgumentsDeltaEvent(
+            type="response.mcp_call_arguments.delta",
+            sequence_number=-1,
+            output_index=state.current_output_index,
+            item_id=state.current_item_id,
+            delta=ctx.last_content_delta,
+        )
+    )
+    return events
+
+
+def emit_code_interpreter_delta_events(
+    ctx: StreamingHarmonyContext,
+    state: HarmonyStreamingState,
+) -> list[StreamingResponsesResponse]:
+    """Emit events for code interpreter delta streaming."""
+    events: list[StreamingResponsesResponse] = []
+    if not state.sent_output_item_added:
+        state.sent_output_item_added = True
+        state.current_item_id = f"tool_{random_uuid()}"
+        events.append(
+            ResponseOutputItemAddedEvent(
+                type="response.output_item.added",
+                sequence_number=-1,
+                output_index=state.current_output_index,
+                item=ResponseCodeInterpreterToolCallParam(
+                    type="code_interpreter_call",
+                    id=state.current_item_id,
+                    code=None,
+                    container_id="auto",
+                    outputs=None,
+                    status="in_progress",
+                ),
+            )
+        )
+        events.append(
+            ResponseCodeInterpreterCallInProgressEvent(
+                type="response.code_interpreter_call.in_progress",
+                sequence_number=-1,
+                output_index=state.current_output_index,
+                item_id=state.current_item_id,
+            )
+        )
+    events.append(
+        ResponseCodeInterpreterCallCodeDeltaEvent(
+            type="response.code_interpreter_call_code.delta",
+            sequence_number=-1,
+            output_index=state.current_output_index,
+            item_id=state.current_item_id,
+            delta=ctx.last_content_delta,
+        )
+    )
+    return events
+
+
+def emit_mcp_prefix_delta_events(
+    ctx: StreamingHarmonyContext,
+    state: HarmonyStreamingState,
+) -> list[StreamingResponsesResponse]:
+    """Emit events for MCP prefix (mcp.*) delta streaming."""
+    events: list[StreamingResponsesResponse] = []
+    if not state.sent_output_item_added:
+        state.sent_output_item_added = True
+        state.current_item_id = f"mcp_{random_uuid()}"
+        mcp_name = ctx.parser.current_recipient[len("mcp.") :]
+
+        events.append(
+            ResponseOutputItemAddedEvent(
+                type="response.output_item.added",
+                sequence_number=-1,
+                output_index=state.current_output_index,
+                item=McpCall(
+                    type="mcp_call",
+                    id=state.current_item_id,
+                    name=mcp_name,
+                    arguments="",
+                    server_label=mcp_name,
+                    status="in_progress",
+                ),
+            )
+        )
+        events.append(
+            ResponseMcpCallInProgressEvent(
+                type="response.mcp_call.in_progress",
+                sequence_number=-1,
+                output_index=state.current_output_index,
+                item_id=state.current_item_id,
+            )
+        )
+
+    events.append(
+        ResponseMcpCallArgumentsDeltaEvent(
+            type="response.mcp_call_arguments.delta",
+            sequence_number=-1,
+            output_index=state.current_output_index,
+            item_id=state.current_item_id,
+            delta=ctx.last_content_delta,
+        )
+    )
+    return events
+
+
+def emit_function_call_delta_events(
+    ctx: StreamingHarmonyContext,
+    state: HarmonyStreamingState,
+) -> list[StreamingResponsesResponse]:
+    """Emit events for developer function calls on commentary channel."""
+    if not (
+        ctx.parser.current_channel == "commentary"
+        and ctx.parser.current_recipient
+        and ctx.parser.current_recipient.startswith("functions.")
+    ):
+        return []
+
+    events: list[StreamingResponsesResponse] = []
+    if state.is_first_function_call_delta is False:
+        state.is_first_function_call_delta = True
+        fc_name = ctx.parser.current_recipient[len("functions.") :]
+        state.current_item_id = f"fc_{random_uuid()}"
+        tool_call_item = ResponseFunctionToolCall(
+            name=fc_name,
+            type="function_call",
+            id=state.current_item_id,
+            call_id=f"call_{random_uuid()}",
+            arguments="",
+            status="in_progress",
+        )
+        events.append(
+            ResponseOutputItemAddedEvent(
+                type="response.output_item.added",
+                sequence_number=-1,
+                output_index=state.current_output_index,
+                item=tool_call_item,
+            )
+        )
+    # Always emit the delta (including on first call)
+    events.append(
+        ResponseFunctionCallArgumentsDeltaEvent(
+            item_id=state.current_item_id,
+            delta=ctx.last_content_delta,
+            output_index=state.current_output_index,
+            sequence_number=-1,
+            type="response.function_call_arguments.delta",
+        )
+    )
+    return events
+
+
+def emit_content_delta_events(
+    ctx: StreamingHarmonyContext,
+    state: HarmonyStreamingState,
+) -> list[StreamingResponsesResponse]:
+    """Emit events for content delta streaming based on channel type."""
+    if not ctx.last_content_delta:
+        return []
+
+    if ctx.parser.current_channel == "final" and ctx.parser.current_recipient is None:
+        return emit_final_channel_delta_events(ctx, state)
+    elif (
+        ctx.parser.current_channel == "analysis"
+        and ctx.parser.current_recipient is None
+    ):
+        return emit_analysis_channel_delta_events(ctx, state)
+    # built-in tools will be triggered on the analysis channel
+    # However, occasionally built-in tools will
+    # still be output to commentary.
+    elif (
+        ctx.parser.current_channel == "commentary"
+        or ctx.parser.current_channel == "analysis"
+    ) and ctx.parser.current_recipient is not None:
+        recipient = ctx.parser.current_recipient
+        # Check for function calls first - they have their own event handling
+        if recipient.startswith("functions."):
+            return emit_function_call_delta_events(ctx, state)
+        if is_mcp_tool_by_namespace(recipient):
+            return emit_mcp_tool_delta_events(ctx, state, recipient)
+        else:
+            return emit_code_interpreter_delta_events(ctx, state)
+    elif (
+        (
+            ctx.parser.current_channel == "commentary"
+            or ctx.parser.current_channel == "analysis"
+        )
+        and ctx.parser.current_recipient is not None
+        and ctx.parser.current_recipient.startswith("mcp.")
+    ):
+        return emit_mcp_prefix_delta_events(ctx, state)
+
+    return []
+
+
+def emit_browser_tool_events(
+    previous_item,
+    state: HarmonyStreamingState,
+) -> list[StreamingResponsesResponse]:
+    """Emit events for browser tool calls (web search)."""
+    function_name = previous_item.recipient[len("browser.") :]
+    parsed_args = json.loads(previous_item.content[0].text)
+    action = None
+
+    if function_name == "search":
+        action = response_function_web_search.ActionSearch(
+            type="search",
+            query=parsed_args["query"],
+        )
+    elif function_name == "open":
+        action = response_function_web_search.ActionOpenPage(
+            type="open_page",
+            # TODO: translate to url
+            url=f"cursor:{parsed_args.get('cursor', '')}",
+        )
+    elif function_name == "find":
+        action = response_function_web_search.ActionFind(
+            type="find",
+            pattern=parsed_args["pattern"],
+            # TODO: translate to url
+            url=f"cursor:{parsed_args.get('cursor', '')}",
+        )
+    else:
+        raise ValueError(f"Unknown function name: {function_name}")
+
+    state.current_item_id = f"tool_{random_uuid()}"
+    events: list[StreamingResponsesResponse] = []
+    events.append(
+        ResponseOutputItemAddedEvent(
+            type="response.output_item.added",
+            sequence_number=-1,
+            output_index=state.current_output_index,
+            item=response_function_web_search.ResponseFunctionWebSearch(
+                # TODO: generate a unique id for web search call
+                type="web_search_call",
+                id=state.current_item_id,
+                action=action,
+                status="in_progress",
+            ),
+        )
+    )
+    events.append(
+        ResponseWebSearchCallInProgressEvent(
+            type="response.web_search_call.in_progress",
+            sequence_number=-1,
+            output_index=state.current_output_index,
+            item_id=state.current_item_id,
+        )
+    )
+    events.append(
+        ResponseWebSearchCallSearchingEvent(
+            type="response.web_search_call.searching",
+            sequence_number=-1,
+            output_index=state.current_output_index,
+            item_id=state.current_item_id,
+        )
+    )
+    # enqueue
+    events.append(
+        ResponseWebSearchCallCompletedEvent(
+            type="response.web_search_call.completed",
+            sequence_number=-1,
+            output_index=state.current_output_index,
+            item_id=state.current_item_id,
+        )
+    )
+    events.append(
+        ResponseOutputItemDoneEvent(
+            type="response.output_item.done",
+            sequence_number=-1,
+            output_index=state.current_output_index,
+            item=ResponseFunctionWebSearch(
+                type="web_search_call",
+                id=state.current_item_id,
+                action=action,
+                status="completed",
+            ),
+        )
+    )
+    return events
+
+
+def emit_mcp_tool_completion_events(
+    previous_item,
+    state: HarmonyStreamingState,
+) -> list[StreamingResponsesResponse]:
+    """Emit events when an MCP tool completes during assistant action turn."""
+    recipient = previous_item.recipient
+    server_label = TOOL_NAME_TO_MCP_SERVER_LABEL.get(recipient, recipient)
+    events: list[StreamingResponsesResponse] = []
+    events.append(
+        ResponseMcpCallArgumentsDoneEvent(
+            type="response.mcp_call_arguments.done",
+            sequence_number=-1,
+            output_index=state.current_output_index,
+            item_id=state.current_item_id,
+            arguments=previous_item.content[0].text,
+            name=recipient,
+        )
+    )
+    events.append(
+        ResponseMcpCallCompletedEvent(
+            type="response.mcp_call.completed",
+            sequence_number=-1,
+            output_index=state.current_output_index,
+            item_id=state.current_item_id,
+        )
+    )
+    events.append(
+        ResponseOutputItemDoneEvent(
+            type="response.output_item.done",
+            sequence_number=-1,
+            output_index=state.current_output_index,
+            item=McpCall(
+                type="mcp_call",
+                id=state.current_item_id,
+                name=recipient,
+                arguments=previous_item.content[0].text,
+                server_label=server_label,
+                status="completed",
+            ),
+        )
+    )
+    return events
+
+
+def emit_code_interpreter_completion_events(
+    previous_item,
+    state: HarmonyStreamingState,
+) -> list[StreamingResponsesResponse]:
+    """Emit events when code interpreter completes."""
+    events: list[StreamingResponsesResponse] = []
+    events.append(
+        ResponseCodeInterpreterCallCodeDoneEvent(
+            type="response.code_interpreter_call_code.done",
+            sequence_number=-1,
+            output_index=state.current_output_index,
+            item_id=state.current_item_id,
+            code=previous_item.content[0].text,
+        )
+    )
+    events.append(
+        ResponseCodeInterpreterCallInterpretingEvent(
+            type="response.code_interpreter_call.interpreting",
+            sequence_number=-1,
+            output_index=state.current_output_index,
+            item_id=state.current_item_id,
+        )
+    )
+    events.append(
+        ResponseCodeInterpreterCallCompletedEvent(
+            type="response.code_interpreter_call.completed",
+            sequence_number=-1,
+            output_index=state.current_output_index,
+            item_id=state.current_item_id,
+        )
+    )
+    events.append(
+        ResponseOutputItemDoneEvent(
+            type="response.output_item.done",
+            sequence_number=-1,
+            output_index=state.current_output_index,
+            item=ResponseCodeInterpreterToolCallParam(
+                type="code_interpreter_call",
+                id=state.current_item_id,
+                code=previous_item.content[0].text,
+                container_id="auto",
+                outputs=[],
+                status="completed",
+            ),
+        )
+    )
+    return events
+
+
+def emit_mcp_prefix_completion_events(
+    previous_item,
+    state: HarmonyStreamingState,
+) -> list[StreamingResponsesResponse]:
+    """Emit events when an MCP prefix tool (mcp.*) completes."""
+    mcp_name = previous_item.recipient[len("mcp.") :]
+    events: list[StreamingResponsesResponse] = []
+    events.append(
+        ResponseMcpCallArgumentsDoneEvent(
+            type="response.mcp_call_arguments.done",
+            sequence_number=-1,
+            output_index=state.current_output_index,
+            item_id=state.current_item_id,
+            arguments=previous_item.content[0].text,
+            name=mcp_name,
+        )
+    )
+    events.append(
+        ResponseMcpCallCompletedEvent(
+            type="response.mcp_call.completed",
+            sequence_number=-1,
+            output_index=state.current_output_index,
+            item_id=state.current_item_id,
+        )
+    )
+    events.append(
+        ResponseOutputItemDoneEvent(
+            type="response.output_item.done",
+            sequence_number=-1,
+            output_index=state.current_output_index,
+            item=McpCall(
+                type="mcp_call",
+                id=state.current_item_id,
+                name=mcp_name,
+                arguments=previous_item.content[0].text,
+                server_label=mcp_name,
+                status="completed",
+            ),
+        )
+    )
+    return events
+
+
+def emit_tool_action_events(
+    ctx: StreamingHarmonyContext,
+    state: HarmonyStreamingState,
+    tool_server: ToolServer | None,
+) -> list[StreamingResponsesResponse]:
+    """Emit events for tool action turn."""
+    if not ctx.is_assistant_action_turn() or len(ctx.parser.messages) == 0:
+        return []
+
+    events: list[StreamingResponsesResponse] = []
+    previous_item = ctx.parser.messages[-1]
+
+    # Handle browser tool
+    if (
+        tool_server is not None
+        and tool_server.has_tool("browser")
+        and previous_item.recipient is not None
+        and previous_item.recipient.startswith("browser.")
+    ):
+        events.extend(emit_browser_tool_events(previous_item, state))
+
+    # Handle tool completion
+    if (
+        tool_server is not None
+        and previous_item.recipient is not None
+        and state.current_item_id is not None
+        and state.sent_output_item_added
+    ):
+        recipient = previous_item.recipient
+        # Handle MCP prefix tool completion first
+        if recipient.startswith("mcp."):
+            events.extend(emit_mcp_prefix_completion_events(previous_item, state))
+        else:
+            # Handle other MCP tool and code interpreter completion
+            is_mcp_tool = is_mcp_tool_by_namespace(
+                recipient
+            ) and state.current_item_id.startswith("mcp_")
+            if is_mcp_tool:
+                events.extend(emit_mcp_tool_completion_events(previous_item, state))
+            else:
+                events.extend(
+                    emit_code_interpreter_completion_events(previous_item, state)
+                )
+
+    return events
@@ -91,7 +91,7 @@ class InputPreprocessor:
        self,
        prompt: str | list[int],
        mm_data: MultiModalDataDict,
-        mm_processor_kwargs: Mapping[str, object] | None,
+        mm_processor_kwargs: Mapping[str, object] | None = None,
        tokenization_kwargs: dict[str, Any] | None = None,
        *,
        mm_uuids: MultiModalUUIDDict | None = None,
@@ -103,9 +103,9 @@ class InputPreprocessor:
        return self.renderer._process_multimodal(
            prompt,
            mm_data,
+            mm_uuids=mm_uuids,
            mm_processor_kwargs=mm_processor_kwargs,
            tokenization_kwargs=tokenization_kwargs,
-            mm_uuids=mm_uuids,
        )

    def _process_embeds(
@@ -144,7 +144,7 @@ class InputPreprocessor:
            inputs = self._process_multimodal(
                prompt_token_ids,
                multi_modal_data,
-                parsed_content.get("mm_processor_kwargs") or {},
+                parsed_content.get("mm_processor_kwargs"),
                tokenization_kwargs=tokenization_kwargs,
                mm_uuids=parsed_content.get("multi_modal_uuids"),
            )
@@ -3,6 +3,7 @@

 from typing import Any

+import regex as re
 import torch

 from vllm.logger import init_logger
@@ -53,44 +54,78 @@ def silu_mul_fp8(input: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
    return out.view(output_shape)


+@silu_mul_fp8.register_input_generator  # type: ignore[misc]
+def generate_silu_mul_fp8_inputs() -> dict[str, tuple[Any, ...]]:
+    intermediate_sizes = [2048, 2880, 4096, 8192, 11008, 14336]
+
+    # Use the same num_tokens values as vLLM's default cudagraph capture sizes.
+    # See vllm/config/vllm.py _set_cudagraph_sizes() for the canonical formula.
+    num_tokens_list = [1, 2, 4] + list(range(8, 256, 8)) + list(range(256, 513, 16))
+
+    inputs = {}
+    for num_tokens in num_tokens_list:
+        for intermediate_size in intermediate_sizes:
+            # Input tensor has shape (num_tokens, 2 * intermediate_size)
+            # because silu_mul splits it into two halves
+            input_tensor = torch.randn(
+                num_tokens,
+                2 * intermediate_size,
+                device="cuda",
+                dtype=torch.bfloat16,
+            )
+            scale = torch.tensor([1.0], device="cuda", dtype=torch.float32)
+
+            config_key = f"intermediate_{intermediate_size}_numtokens_{num_tokens}"
+            inputs[config_key] = (input_tensor, scale)
+
+    return inputs
+
+
@silu_mul_fp8.register_config_picker  # type: ignore[misc]
 def pick_silu_mul_fp8_config(
    args: tuple[Any, ...], config_keys: list[str]
 ) -> str | None:
+    """Pick the best pre-tuned config for the given input shape.
+
+    Selection strategy:
+      1. Find the closest intermediate_size among available configs
+         (exact match preferred).
+      2. Among the num_tokens values tuned for that intermediate_size, pick
+         the smallest num_tokens >= the input's num_tokens. If the input is
+         larger than all available num_tokens, fall back to the largest.
+
+    Config keys must be "default" or follow the format
+    "intermediate_{int}_numtokens_{int}".
+    """
    if not config_keys:
        return None

-    input_tensor, scale = args
+    input_tensor, _scale = args
    intermediate_size = input_tensor.shape[-1] // 2
-
-    # TODO(gmagosfm): Rerun autotuning to capture config for
-    # other batch sizes.
-    target_key = f"intermediate_{intermediate_size}_batchsize_256"
-    if target_key in config_keys:
-        return target_key
-
-    intermediate_sizes = []
+    num_tokens = input_tensor.view(-1, input_tensor.shape[-1]).shape[0]
+    configs: dict[int, list[int]] = {}
    for key in config_keys:
-        if key.startswith("intermediate_") and "_batchsize_256" in key:
-            try:
-                size_str = key.split("_")[1]
-                size = int(size_str)
-                intermediate_sizes.append((abs(size - intermediate_size), key))
-            except (ValueError, IndexError):
-                continue
+        if key == "default":
+            continue
+        match = re.fullmatch(r"intermediate_(\d+)_numtokens_(\d+)", key)
+        if not match:
+            raise ValueError(
+                f"Malformed config key '{key}', "
+                f"expected format 'intermediate_{{int}}_numtokens_{{int}}'"
+            )
+        isize_str, ntokens_str = match.groups()
+        configs.setdefault(int(isize_str), []).append(int(ntokens_str))

-    if intermediate_sizes:
-        _, best_key = min(intermediate_sizes)
-        logger.debug(
-            "No exact config for intermediate_size=%d, using closest match: %s",
-            intermediate_size,
-            best_key,
-        )
-        return best_key
-    if "default" in config_keys:
-        return "default"
+    if not configs:
+        return "default" if "default" in config_keys else None

-    return None
+    best_isize = min(configs, key=lambda s: abs(s - intermediate_size))
+    available_ntokens = sorted(configs[best_isize])
+    best_ntokens = next(
+        (n for n in available_ntokens if n >= num_tokens), available_ntokens[-1]
+    )
+
+    return f"intermediate_{best_isize}_numtokens_{best_ntokens}"


 def silu_mul_fp8_baseline(input: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
@@ -10,7 +10,6 @@
 import warnings

 import torch
-from einops import rearrange

 from .chunk_delta_h import chunk_gated_delta_rule_fwd_h
 from .chunk_o import chunk_fwd_o
@@ -119,21 +118,20 @@ def chunk_gated_delta_rule(
    initial_state: torch.Tensor = None,
    output_final_state: bool = False,
    cu_seqlens: torch.LongTensor | None = None,
-    head_first: bool = False,
    use_qk_l2norm_in_kernel: bool = False,
 ):
    r"""
    Args:
        q (torch.Tensor):
-            queries of shape `[B, T, H, K]` if `head_first=False` else `[B, H, T, K]`.
+            Queries of shape `[B, T, H, K]`.
        k (torch.Tensor):
-            keys of shape `[B, T, H, K]` if `head_first=False` else `[B, H, T, K]`.
+            Keys of shape `[B, T, H, K]`.
        v (torch.Tensor):
-            values of shape `[B, T, H, V]` if `head_first=False` else `[B, H, T, V]`.
+            Values of shape `[B, T, H, V]`.
        g (torch.Tensor):
-            (forget) gating tensor (in log space!) of shape `[B, T, H]` if `head_first=False` else `[B, H, T]`.
+            (forget) Gating tensor (in log space!) of shape `[B, T, H]`.
        beta (torch.Tensor):
-            betas of shape `[B, T, H]` if `head_first=False` else `[B, H, T]`.
+            Betas of shape `[B, T, H]`.
        scale (Optional[int]):
            Scale factor for the RetNet attention scores.
            If not provided, it will default to `1 / sqrt(K)`. Default: `None`.
@@ -146,13 +144,9 @@ def chunk_gated_delta_rule(
        cu_seqlens (torch.LongTensor):
            Cumulative sequence lengths of shape `[N+1]` used for variable-length training,
            consistent with the FlashAttention API.
-        head_first (Optional[bool]):
-            Whether the inputs are in the head-first format, which is not supported for variable-length inputs.
-            Default: `False`.
-
    Returns:
        o (torch.Tensor):
-            Outputs of shape `[B, T, H, V]` if `head_first=False` else `[B, H, T, V]`.
+            Outputs of shape `[B, T, H, V]`.
        final_state (torch.Tensor):
            Final state of shape `[N, H, V, K]` if `output_final_state=True` else `None`.

@@ -189,24 +183,11 @@ def chunk_gated_delta_rule(
    assert q.dtype != torch.float32, (
        "ChunkGatedDeltaRuleFunction does not support float32. Please use bfloat16."
    )
-    assert len(beta.shape) == 3, (
-        "beta must be of shape [B, T, H] if head_first=False, or [B, H, T] otherwise."
-    )
-
-    if head_first:
-        raise DeprecationWarning(
-            "head_first is deprecated and will be removed in a future version. "
-            "Please use head_first=False for now instead.",
-            stacklevel=2,
-        )
-        q, k, v, beta, g = map(
-            lambda x: rearrange(x, "b h t ... -> b t h ..."), (q, k, v, beta, g)
-        )
-    if not head_first and q.shape[1] < q.shape[2]:
+    assert len(beta.shape) == 3, "beta must be of shape [B, T, H]."
+    if q.shape[1] < q.shape[2]:
        warnings.warn(
            f"Input tensor shape suggests potential format mismatch: seq_len ({q.shape[1]}) < num_heads ({q.shape[2]}). "
            "This may indicate the inputs were passed in head-first format [B, H, T, ...] "
-            "when head_first=False was specified. "
            "Please verify your input tensor format matches the expected shape [B, T, H, ...].",
            stacklevel=2,
        )
@@ -235,6 +216,4 @@ def chunk_gated_delta_rule(
        cu_seqlens,
        use_qk_l2norm_in_kernel,
    )
-    if head_first:
-        o = rearrange(o, "b t h ... -> b h t ...")
    return o, final_state
@@ -8,11 +8,7 @@ import torch

 import vllm.envs as envs
 from vllm.config import ParallelConfig
-from vllm.distributed import (
-    get_dp_group,
-    get_pcp_group,
-    get_tensor_model_parallel_rank,
-)
+from vllm.distributed import get_dp_group, get_pcp_group, get_tensor_model_parallel_rank
 from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
 from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
@@ -126,20 +122,31 @@ class RoutingMethodType(IntEnum):


 def get_routing_method_type(
-    scoring_func: str, top_k: int, renormalize: bool
+    scoring_func: str,
+    top_k: int,
+    renormalize: bool,
+    num_expert_group: int | None,
+    has_e_score_bias: bool,
 ) -> RoutingMethodType:
+    if has_e_score_bias:
+        if (num_expert_group or 0) > 0 and scoring_func == "sigmoid":
+            return RoutingMethodType.DeepSeekV3
+        else:
+            return RoutingMethodType.Unspecified
+
    if scoring_func == "sigmoid":
        if top_k == 1:
            return RoutingMethodType.Llama4
        else:
-            return RoutingMethodType.DeepSeekV3
-    elif scoring_func == "softmax":
+            return RoutingMethodType.Unspecified
+
+    if scoring_func == "softmax":
        if renormalize:
            return RoutingMethodType.Renormalize
        else:
            return RoutingMethodType.Default
-    else:
-        return RoutingMethodType.Unspecified
+
+    return RoutingMethodType.Unspecified


@dataclass
@@ -402,7 +402,7 @@ class CPUFusedMOE:
            input,
            topk_weights,
            topk_ids,
-            activation,
+            activation.value,
            global_num_experts,
            skip_weighted,
        )
@@ -129,25 +129,28 @@ def is_supported_config_trtllm_fp8(
        return f"kernel does not support {reason}"

    if not _supports_current_device():
-        return False, _make_reason("current device")
+        return False, _make_reason(f"current device {current_platform.device_name}")
    elif not (moe_config.is_act_and_mul or _supports_no_act_and_mul()):
        return False, _make_reason("no act_and_mul MLP layer")
    elif not _supports_activation(moe_config.activation):
        return False, _make_reason(f"{moe_config.activation} activation")
    elif not _supports_quant_scheme(weight_key, activation_key):
-        return False, _make_reason("quantization scheme")
+        return False, _make_reason(f"quantization scheme {weight_key}x{activation_key}")
    elif not _supports_parallel_config(moe_config.moe_parallel_config):
-        return False, _make_reason("parallel config")
+        return False, _make_reason(f"parallel config {moe_config.moe_parallel_config}")
    elif not _supports_routing_method(
        weight_key, activation_key, moe_config.routing_method
    ):
-        return False, _make_reason("routing method")
+        return False, _make_reason(f"routing method {moe_config.routing_method}")
    elif activation_format != mk.FusedMoEActivationFormat.Standard:
-        return False, _make_reason("activation format")
+        return False, _make_reason(f"activation format {activation_format}")
    elif not _supports_router_logits_dtype(
        moe_config.router_logits_dtype, moe_config.routing_method
    ):
-        return False, _make_reason("float32 router_logits with non-DeepSeekV3 routing")
+        return False, _make_reason(
+            "float32 router_logits with non-DeepSeekV3 routing "
+            f"{moe_config.router_logits_dtype}x{moe_config.routing_method}"
+        )

    return True, None

@@ -165,17 +168,17 @@ def is_supported_config_trtllm_bf16(
        return f"kernel does not support {reason}"

    if not _supports_current_device():
-        return False, _make_reason("current device")
+        return False, _make_reason(f"current device {current_platform.device_name}")
    elif not (moe_config.is_act_and_mul or _supports_no_act_and_mul()):
        return False, _make_reason("no act_and_mul MLP layer")
    elif not _supports_activation(moe_config.activation):
        return False, _make_reason(f"{moe_config.activation} activation")
    elif not _supports_parallel_config(moe_config.moe_parallel_config):
-        return False, _make_reason("parallel config")
+        return False, _make_reason(f"parallel config {moe_config.moe_parallel_config}")
    elif not _supports_routing_method_bf16(moe_config.routing_method):
-        return False, _make_reason("routing method")
+        return False, _make_reason(f"routing method {moe_config.routing_method}")
    elif activation_format != mk.FusedMoEActivationFormat.Standard:
-        return False, _make_reason("activation format")
+        return False, _make_reason(f"activation format {activation_format}")

    return True, None

@@ -30,6 +30,7 @@ from vllm.model_executor.layers.fused_moe.utils import (
    disable_inplace,
 )
 from vllm.model_executor.layers.quantization.utils.marlin_utils import (
+    get_marlin_input_dtype,
    marlin_make_workspace_new,
    marlin_moe_intermediate_size,
    marlin_quant_input,
@@ -550,6 +551,8 @@ class MarlinExpertsBase(mk.FusedMoEPermuteExpertsUnpermute):
        self.w13_g_idx_sort_indices = w13_g_idx_sort_indices
        self.w2_g_idx_sort_indices = w2_g_idx_sort_indices
        self.is_k_full = is_k_full
+        self.input_dtype = get_marlin_input_dtype()
+
        super().__init__(
            moe_config=moe_config,
            quant_config=quant_config,
@@ -736,6 +739,7 @@ class MarlinExperts(MarlinExpertsBase):
            sort_indices1=self.w13_g_idx_sort_indices,
            sort_indices2=self.w2_g_idx_sort_indices,
            is_k_full=self.is_k_full,
+            input_dtype=self.input_dtype,
        )

    def moe_sum(self, input: torch.Tensor, output: torch.Tensor) -> None:
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
Zhuohan LiandGitHub	cfc1df2184	Simplify output assignment condition Refactor conditional check for output assignment. Signed-off-by: Zhuohan Li <zhuohan123@gmail.com>	2026-02-20 15:27:56 -08:00
Wei Zhao GitHub Cyrus Leung	ea5f903f80	Bump Flashinfer Version and Re-enable DeepSeek NVFP4 AR+Norm Fusion (#34899 ) Signed-off-by: wzhao18 <wzhao18.sz@gmail.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>	2026-02-20 13:37:31 -08:00
Ryan RockandGitHub	0632ed8778	[AMD][CI] Fix test_custom_allreduce for A100 testgroup (#34735 ) Signed-off-by: Ryan Rock <ryan.rock@amd.com>	2026-02-20 21:33:04 +00:00
Lucas WilkinsonandGitHub	aaefc58ee0	[CI] Revert PRs 34818 and 33600 (#34979 )	2026-02-20 13:25:50 -08:00
Wei ZhaoandGitHub	f24b2de3d3	[Test] Add FP8 KV Cache Testing for MLA Backends (#34473 ) Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>	2026-02-20 18:51:58 +00:00
Michael GoinandGitHub	fac1507f03	[CI] Remove failing prime-rl integration test (#34843 ) Signed-off-by: mgoin <mgoin64@gmail.com> Signed-off-by: Michael Goin <mgoin64@gmail.com>	2026-02-20 10:17:42 -08:00
Zhengxu Chen GitHub Cyrus Leung	f863994084	[compile] Fix torch.compile time discrepancy in logging. (#34912 ) Signed-off-by: zhxchen17 <zhxchen17@fb.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>	2026-02-20 08:47:14 -08:00
Zhengxu ChenandGitHub	e4a5d8c653	[compile] Move torch_aot_compile directory under torch_compile_cache (#34831 ) Signed-off-by: zhxchen17 <zhxchen17@fb.com>	2026-02-20 08:46:45 -08:00
Yanan CaoandGitHub	a6d0299c75	[Kernel] [Helion] [6/N] Add num_tokens dimension to silu_mul autotuning and dispatching (#34185 ) Signed-off-by: Yanan Cao <gmagogsfm@gmail.com>	2026-02-20 08:36:51 -08:00
Harry MellorandGitHub	6ce80f7071	Ensure that MkDocs v2 does not get installed (#34958 ) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>	2026-02-20 15:38:11 +00:00
Huamin Li GitHub Cyrus Leung	1fe462168c	[perf] Avoid dtype promotion sync in mamba_get_block_table_tensor (#34870 ) Signed-off-by: Huamin Li <3ericli@gmail.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>	2026-02-20 06:21:56 -08:00
Flora Feng GitHub Cyrus Leung	ed31a020ee	[Refactor] Extract Harmony streaming SSE event builders into streaming_events.py (#34909 ) Signed-off-by: sfeng33 <4florafeng@gmail.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>	2026-02-20 06:20:46 -08:00
Cyrus LeungandGitHub	f9ac19204f	[V0 Deprecation] Remove unused MM placeholders in request output (#34944 ) Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>	2026-02-20 06:19:23 -08:00
Vadim GimpelsonandGitHub	59965affbd	[BUGFIX] Fix `_dummy_run` missing `prepare_inputs_event` synchronization (#34866 ) Signed-off-by: Vadim Gimpelson <vadim.gimpelson@gmail.com>	2026-02-20 05:54:27 -08:00
Xin YangandGitHub	b1c4f0b265	[Kernel] Optimize grouped topk kernel (#34206 ) Signed-off-by: Xin Yang <xyangx@amazon.com>	2026-02-20 01:34:45 -08:00
Kevin McKay GitHub Lu Fang	8de7c636cc	[Bugfix][Hardware][AMD] Fix ROCM_AITER_FA speculative decoding support (#32877 ) Signed-off-by: c0de128 <kevin.mckay@outlook.com> Co-authored-by: Lu Fang <30275821+houseroad@users.noreply.github.com>	2026-02-19 22:25:46 -08:00
Frank Wang GitHub Wentao Ye	059779231f	[Minor] Add logging when using MXFP4 MXFP8 TRTLLM backend (#34916 ) Signed-off-by: frankwang28 <frank.wbb@hotmail.com> Signed-off-by: Frank Wang <41319051+frankwang28@users.noreply.github.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>	2026-02-19 22:07:57 -08:00
tianshu-Michael-yu GitHub Piotr Mazurek Cyrus Leung	ea37530b47	[Models] LFM2: Support LoRA (#34921 ) Co-authored-by: Piotr Mazurek <piotr635@gmail.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>	2026-02-19 22:07:23 -08:00
Micah WilliamsonandGitHub	f5432e35a3	[ROCm][CI] Loosen RemoteOpenAIServer Startup Timeout (#34922 ) Signed-off-by: Micah Williamson <micah.williamson@amd.com>	2026-02-20 05:37:49 +00:00
杨朱 · Kiki GitHub Claude Opus 4.5 Harry Mellor	07cab212f0	[Misc] Add deprecated environment variable utilities (#33677 ) Signed-off-by: carlory <baofa.fan@daocloud.io> Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>	2026-02-19 21:33:25 -08:00
rasmithandGitHub	0c1dc42748	[CI][AMD][BugFix][P/D] Add default_vllm_config to test_moriio_connector.py so tests pass (#33739 ) Signed-off-by: Randall Smith <Randall.Smith@amd.com>	2026-02-19 21:32:40 -08:00
Varun ChawlaandGitHub	676f82ae81	Add validation to reject non-text content in system messages (#34072 ) Signed-off-by: Varun Chawla <varun_6april@hotmail.com>	2026-02-19 21:30:33 -08:00
Elizabeth Thomas GitHub Robert Shaw Robert Shaw	81bfc21a6a	[Model Bash]: Improve FP8 Oracle for Config Specific Kernel Selection (#34260 ) Signed-off-by: Elizabeth Thomas <email2eliza@gmail.com> Signed-off-by: Robert Shaw <robertgshaw2-redhat@h100-02.nemg-001.lab.rdu2.dc.redhat.com> Signed-off-by: Robert Shaw <robertgshaw2@gmail.com> Co-authored-by: Robert Shaw <robertgshaw2-redhat@h100-02.nemg-001.lab.rdu2.dc.redhat.com> Co-authored-by: Robert Shaw <robertgshaw2@gmail.com>	2026-02-19 21:29:08 -08:00
Matthias GehreandGitHub	4e2c7caf2d	[Bugfix] Add regression test for MoE quant_config under torch.compile (#34335 ) Signed-off-by: Matthias Gehre <matthias.gehre@amd.com>	2026-02-20 13:27:26 +08:00
Bowen BaoandGitHub	d9e62c03eb	[Quark] Fix MoE fp8 activation scale handling on mi300 (#34386 ) Signed-off-by: Bowen Bao <bowenbao@amd.com>	2026-02-19 21:27:14 -08:00
Kevin H. LuuandGitHub	a1a2d79442	[ci] Use the right tag for CPU arm64 image (#34915 ) Signed-off-by: Kevin H. Luu <khluu000@gmail.com>	2026-02-19 19:59:15 -08:00
Cyrus LeungandGitHub	ac900c89bb	[Refactor] Implement output type check in LLM (#34794 ) Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>	2026-02-19 19:57:55 -08:00
Mark McLoughlinandGitHub	76df6072ff	[Core] Fix state names in pause_scheduler() (#34840 ) Signed-off-by: Mark McLoughlin <markmc@redhat.com>	2026-02-19 17:21:46 -08:00
Michael GoinandGitHub	16f24e8797	[CI] Add GPT-OSS Eval job for H100 (#34359 ) Signed-off-by: Michael Goin <mgoin64@gmail.com>	2026-02-19 17:14:54 -08:00
Nick HillandGitHub	40b2f1c3d9	[Model Runner V2] Minor CPU optimizations (#34856 ) Signed-off-by: Nick Hill <nickhill123@gmail.com>	2026-02-19 16:05:37 -08:00
Mayank Ketkar GitHub Mayank Ketkar Claude Opus 4.6	648951a9c3	[Bugfix] Fix benchmark_fused_collective crash on CustomOp init (#34665 ) Signed-off-by: Mayank Ketkar <mketkar@zoox.com> Signed-off-by: Mayank Ketkar <mayket04@gmail.com> Co-authored-by: Mayank Ketkar <mketkar@zoox.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>	2026-02-19 19:01:00 -05:00
Michael GoinandGitHub	f72061a19a	[UX] More descriptive reasons in is_supported_config for MoE (#34908 ) Signed-off-by: mgoin <mgoin64@gmail.com>	2026-02-19 15:20:52 -08:00
Matthew Bonanni GitHub Robert Shaw	662205d34e	[Bugfix] Fix Basic Models Test (#34818 ) Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>	2026-02-19 14:49:07 -08:00
Roger WangandGitHub	4fb8beefaa	[Bugfix] Fix cutlass fp8 kernel on hopper for Qwen3.5 (#34914 ) Signed-off-by: Roger Wang <hey@rogerw.io>	2026-02-19 13:34:55 -08:00
Alexei-V-Ivanov-AMDandGitHub	304319c4ed	Change targets for AMD build in the "CI" pipeline (#34918 ) Signed-off-by: Alexei V. Ivanov <alexei.ivanov@amd.com>	2026-02-19 21:26:53 +00:00
Wentao YeandGitHub	c683d11c94	[Refactor] Deprecate `head_first` for `chunk_gated_delta_rule` (#34263 ) Signed-off-by: yewentao256 <zhyanwentao@126.com>	2026-02-19 13:23:49 -05:00
roikoren755 GitHub Michael Goin	3eff45d793	Revert "[NemotronH] Do not force router to run in fp32 (#34582 )" (#34808 ) Signed-off-by: Roi Koren <roik@nvidia.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>	2026-02-19 09:47:05 -08:00
Robert Shaw GitHub Robert Shaw	4685a630a2	[Model Bash][DeepSeekR1] Remove Shared Expert Clone (#34344 ) Signed-off-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Robert Shaw <robshaw@redhat.com>	2026-02-19 07:56:14 -08:00
Eldar Kurtić GitHub Your Name Robert Shaw	ee1d25f199	[Llama4,Quantization] Simplify and generalize logic for Q/K permutations in quantized self-attn layers (#34471 ) Signed-off-by: Your Name <you@example.com> Co-authored-by: Your Name <you@example.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>	2026-02-19 07:55:41 -08:00
LindaandGitHub	6fff24f30f	[Bugfix] Qwen3.5 kv-scale weight remapping (#34719 ) Signed-off-by: Linda-Stadter <57756729+Linda-Stadter@users.noreply.github.com>	2026-02-19 04:13:37 -08:00
Cyrus LeungandGitHub	23210a911e	[CI/Build] Try to make beam search test less flaky (#34885 ) Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>	2026-02-19 19:16:58 +08:00
Cyrus LeungandGitHub	1391378861	[Bugfix] Fix edge case in UUID data parsing (#34884 ) Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>	2026-02-19 02:24:30 -08:00
Andreas KaratzasandGitHub	f6220f9877	[ROCm][Test] Fix beam search determinism failures from batch-size-dependent FP divergence and removed wrong marker (#34878 ) Signed-off-by: Andreas Karatzas <akaratza@amd.com>	2026-02-19 08:25:26 +00:00
Andreas KaratzasandGitHub	2df2bb27b0	[ROCm][CI] Removing all blocking labels from MI355 until stable infra (#34879 ) Signed-off-by: Andreas Karatzas <akaratza@amd.com>	2026-02-19 07:53:08 +00:00
Tal Nir GitHub Claude Opus 4.6	f75b61a9e9	[Voxtral Realtime] Fix engine crash on empty multimodal embeddings (#34862 ) Signed-off-by: Tal Nir <tal@nervexneurotech.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>	2026-02-18 23:21:47 -08:00
Wei ZhaoandGitHub	7f51e93864	[Bug] Fix DeepSeek V3 weight loading caused by incorrect prefix (#34876 ) Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>	2026-02-18 23:20:30 -08:00
Alex BrooksandGitHub	4611af1663	[Bugfix] Add Quant Config to Llava Next Projector (#34847 ) Signed-off-by: Alex Brooks <albrooks@redhat.com>	2026-02-18 23:18:23 -08:00
Manrique VargasandGitHub	ad5aa6bd9f	fix(docs): fix typos in comments and docstrings (#34836 ) Signed-off-by: machov <mv1742@nyu.edu>	2026-02-18 23:17:41 -08:00
Jaeyeon Kim(김재연)andGitHub	9681068cf9	[Frontend] Fix reasoning_tokens for text-based parsers in Responses API (#33513 ) Signed-off-by: Jaeyeon Kim <anencore94@gmail.com>	2026-02-18 23:16:41 -08:00
Kevin H. LuuandGitHub	b6101d384d	Deprecate test-pipeline.yaml (#34864 ) Signed-off-by: Kevin H. Luu <khluu000@gmail.com>	2026-02-19 02:15:27 +00:00
Woosuk KwonandGitHub	5fcb0cdd68	[Model Runner V2] Use FP32 for Gumbel Noise (#34854 ) Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-02-18 17:07:37 -08:00
Woosuk KwonandGitHub	c878b43b64	[Model Runner V2] Remove unnecessary copies in PW CUDA graph capture (#34849 ) Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-02-18 15:52:50 -08:00
rasmithandGitHub	2b84ac669c	[CI][AMD][BugFix] Use torch.testing.assert_close instead of assert torch.allclose in test_rocm_skinny_gemms.py (#34181 ) Signed-off-by: Randall Smith <Randall.Smith@amd.com>	2026-02-18 23:10:19 +00:00
zhrrrandGitHub	11d3976b88	[Model Runner V2] support piecewise & mixed cudagraph (#32771 ) Signed-off-by: zhuhaoran <zhuhaoran.zhr@alibaba-inc.com>	2026-02-18 15:03:17 -08:00
Yongye Zhu GitHub Robert Shaw	40da9625a1	[MoE Refactor] Convert mxfp4 marlin into modular kernel format (#34588 ) Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>	2026-02-18 14:37:14 -08:00
Flora Feng GitHub Flora Feng Nicolò Lucchesi	8d9babd4de	Fix empty tool_call_id in Anthropic messages API tool result conversion (#34745 ) Signed-off-by: <> Signed-off-by: sfeng33 <4florafeng@gmail.com> Co-authored-by: Flora Feng <sfeng33@h100-01.nemg-001.lab.rdu2.dc.redhat.com> Co-authored-by: Nicolò Lucchesi <nlucches@redhat.com>	2026-02-18 14:31:59 -08:00
Aaron HaoandGitHub	e99ba957ec	[BUG] Fixing Weight Sync unit test (#34841 ) Signed-off-by: ahao-anyscale <ahao@anyscale.com>	2026-02-18 17:20:10 -05:00
Kyle SayersandGitHub	64ac1395e8	[Docs] Clean up speculators docs (#34065 ) Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>	2026-02-18 13:48:11 -08:00
Cyrus Leung GitHub Michael Goin	61cf087680	[Bugfix] Fix lora tests (#34834 ) Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk> Signed-off-by: Michael Goin <mgoin64@gmail.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>	2026-02-18 13:22:31 -08:00
Wenlong Wang GitHub Robert Shaw Robert Shaw	847a57cd12	[Bugfix][MoE Kernel] Fix incorrect routing selection for models without expert groups (e.g., MiniMax-M2.1) (#34673 ) Signed-off-by: wwl2755 <wangwenlong2755@gmail.com> Signed-off-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>	2026-02-18 13:03:24 -08:00
rasmithandGitHub	fcd6ac97ed	[CI][AMD][BugFix] Skip tests in test_unquantized_backend_selection that should not run on ROCm (#34655 ) Signed-off-by: Randall Smith <Randall.Smith@amd.com>	2026-02-18 15:00:40 -05:00
Woosuk KwonandGitHub	95be2a7f22	[Model Runner V2] Minor simplification for DCP (#34786 ) Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-02-18 11:04:53 -08:00
Jaden MathiasandGitHub	0e60c925cf	[Bugfix] Remove assert causing hipErrorStreamCaptureUnsupported (#34455 ) Signed-off-by: Jaden Mathias <jaden.mathias@amd.com>	2026-02-18 18:54:54 +00:00
Teng MaandGitHub	d7ff22204a	[Misc] Add mooncake-transfer-engine to kv_connectors requirements (#34826 ) Signed-off-by: Teng Ma <teng-ma@linux.alibaba.com>	2026-02-18 18:26:24 +00:00
Isotr0py GitHub JJJYmmm	c0bd8b13da	[Bugfix] Redo Qwen3.5/Qwen3-Next GDN projector fusion (#34697 ) Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Co-authored-by: JJJYmmm <92386084+JJJYmmm@users.noreply.github.com>	2026-02-18 09:46:53 -08:00
Michael GoinandGitHub	caeb887bf6	[Bugfix] Fix NVFP4 TRTLLM MoE non-gated support; add gsm8k for Nemotron-3-Nano FP8+NVFP4 (#34725 ) Signed-off-by: mgoin <mgoin64@gmail.com>	2026-02-18 09:39:22 -08:00
Ilya MarkovandGitHub	6b3166a7c7	[CI][Bugfix] Fix multinode test script (#34820 ) Signed-off-by: ilmarkov <markovilya197@gmail.com>	2026-02-18 11:45:10 -05:00
Robert Shaw GitHub Robert Shaw	25e2e136ef	[CI] temporarily disable multi-node tests (#34825 ) Signed-off-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Robert Shaw <robshaw@redhat.com>	2026-02-18 11:32:44 -05:00
Robert Shaw GitHub Robert Shaw	6874638bc4	[Model Bash] DeepSeek R1 BF16 Min Latency QKV A GEMM (0.5% E2E Speedup) (#34758 ) Signed-off-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Robert Shaw <robshaw@redhat.com>	2026-02-18 07:42:36 -08:00
Burkhard Ringlein GitHub Luka Govedič	e24663c5a9	Add unit tests for fp8 output fusion of triton_attn (#34228 ) Signed-off-by: Burkhard Ringlein <ngl@zurich.ibm.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>	2026-02-18 06:22:49 -05:00
Nick HillandGitHub	c50e105a88	[Model Runner V2] Avoid prepare prefill kernel launch overhead (#34780 ) Signed-off-by: Nick Hill <nickhill123@gmail.com>	2026-02-18 00:49:21 -08:00
Cyrus LeungandGitHub	a766b30349	[Renderer] Deprecate code paths for old input processing (#34775 ) Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>	2026-02-18 00:35:04 -08:00
Asaf Joseph GardinandGitHub	1faa8cb73c	[Quantization] - Added uses_meta_device_weights to quant config (#34645 ) Signed-off-by: Josephasafg <ajgard7@gmail.com>	2026-02-17 23:43:44 -08:00
Marek MichalowskiandGitHub	e89a91d927	[Bugfix] fix activation in cpu_fused_moe_torch call (#34696 ) Signed-off-by: Marek Michalowski <marek.michalowski@arm.com>	2026-02-17 23:39:46 -08:00
Michael GoinandGitHub	909b147197	[Bugfix] Fix prefix creation for Qwen3.5 (#34723 ) Signed-off-by: mgoin <mgoin64@gmail.com>	2026-02-17 23:39:15 -08:00
ElizaWszola GitHub Luka Govedič	a88b3be7c4	[Bugfix] Fix quant RMS norm fusion for quantization with TMA-aligned scales (#33255 ) Signed-off-by: ElizaWszola <ewszola@redhat.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>	2026-02-17 23:35:04 -08:00
Nick HillandGitHub	a49ea5a58f	[Model Runner V2] A bit more PP simplification (#34766 ) Signed-off-by: Nick Hill <nickhill123@gmail.com>	2026-02-17 21:39:07 -08:00
Cyrus LeungandGitHub	30ebe0dc3c	[CI/Build] Remove use of `skip_v1` (#34699 ) Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>	2026-02-18 12:19:11 +08:00
Andreas KaratzasandGitHub	cef65f0715	[ROCm][CI] Removed hard-coded attn backend requirement for Qwen VL (#34753 ) Signed-off-by: Andreas Karatzas <akaratza@amd.com>	2026-02-18 03:59:53 +00:00
Russell Bryant GitHub isotr0py	6f3b2047ab	[Core] Fix SSRF bypass via backslash-@ URL parsing inconsistency (#34743 ) Signed-off-by: Russell Bryant <rbryant@redhat.com> Co-authored-by: isotr0py <2037008807@qq.com>	2026-02-18 03:53:35 +00:00
Luka GovedičandGitHub	02e8f26cea	[torch.compile] Turn on silu+fp4 quant fusion by default for O1+ (#34718 ) Signed-off-by: Luka Govedič <lgovedic@redhat.com>	2026-02-18 03:29:15 +00:00
Hongxia Yang GitHub Hongxia Yang	4a00a511bb	[BugFix] [Build] fix string literals comparison in indexer_k_quant_and_cache calling site (#34653 ) Signed-off-by: Hongxia Yang <hongxiay.yang@amd.com> Co-authored-by: Hongxia Yang <hongxiay.yang@amd.com>	2026-02-17 19:19:41 -08:00
Cyrus LeungandGitHub	a0d8d944e2	[Renderer] Move MM Hash parsing into Renderer (#34711 ) Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>	2026-02-17 19:18:55 -08:00