Merge branch 'main' into wentao-fix-dcp-IMA-for-v2

Signed-off-by: yewentao256 <zhyanwentao@126.com>

.buildkite/hardware_tests/amd.yaml

@@ -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

.buildkite/image_build/image_build_cpu_arm64.sh

@@ -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

.buildkite/scripts/run-prime-rl-test.sh

@@ -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!"

.buildkite/test-amd.yaml

@@ -67,7 +67,7 @@ steps:
   timeout_in_minutes: 30
   mirror_hardwares: [amdexperimental, amdproduction, amdtentative]
   agent_pool: mi325_1
-  grade: Blocking
+  # grade: Blocking
   source_file_dependencies:
   - vllm/
   - tests/test_inputs.py
@@ -1622,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]
@@ -3201,20 +3186,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: mi355_2
-  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]

.buildkite/test_areas/distributed.yaml

@@ -209,7 +209,7 @@ steps:
     - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
     - CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
 
-- label: Pipeline + Context Parallelism (4 GPUs))
+- label: Pipeline + Context Parallelism (4 GPUs)
   timeout_in_minutes: 60
   working_dir: "/vllm-workspace/tests"
   num_devices: 4

.buildkite/test_areas/e2e_integration.yaml

@@ -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

.buildkite/test_areas/misc.yaml

@@ -147,6 +147,19 @@ steps:
   - pytest -v -s transformers_utils
   - pytest -v -s config
 
+- label: GPT-OSS Eval (H100)
+  timeout_in_minutes: 60
+  working_dir: "/vllm-workspace/"
+  device: h100
+  optional: true
+  source_file_dependencies:
+  - tests/evals/gpt_oss
+  - vllm/model_executor/models/gpt_oss.py
+  - vllm/model_executor/layers/quantization/mxfp4.py
+  commands:
+    - uv pip install --system 'gpt-oss[eval]==0.0.5'
+    - pytest -s -v tests/evals/gpt_oss/test_gpqa_correctness.py --model openai/gpt-oss-20b --metric 0.58
+
 - label: GPT-OSS Eval (B200)
   timeout_in_minutes: 60
   working_dir: "/vllm-workspace/"

benchmarks/kernels/benchmark_fused_collective.py

@@ -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,

csrc/activation_kernels.cu

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

csrc/moe/grouped_topk_kernels.cu

@@ -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};
 }

csrc/moe/moeTopKFuncs.cuh

@@ -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

docker/Dockerfile

@@ -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} \

docker/Dockerfile.nightly_torch

@@ -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..." \

docker/versions.json

@@ -68,7 +68,7 @@
       "default": "true"
     },
     "FLASHINFER_VERSION": {
-      "default": "0.6.3"
+      "default": "0.6.4"
     },
     "GDRCOPY_CUDA_VERSION": {
       "default": "12.8"

docs/features/nixl_connector_usage.md

@@ -197,8 +197,8 @@ For multi-host DP deployment, only need to provide the host/port of the head ins
 
 The `kv_load_failure_policy` setting controls how the system handles failures when the decoder instance loads KV cache blocks from the prefiller instance:
 
-- **fail** (recommended): Immediately fail the request with an error when KV load fails. This prevents performance degradation by avoiding recomputation of prefill work on the decode instance.
-- **recompute** (default): Recompute failed blocks locally on the decode instance. This may cause performance _jitter_ on decode instances as the scheduled prefill will delay and interfere with other decodes. Furthermore, decode instances are typically configured with low-latency optimizations.
+- **fail** (default): Immediately fail the request with an error when KV load fails. This prevents performance degradation by avoiding recomputation of prefill work on the decode instance.
+- **recompute**: Recompute failed blocks locally on the decode instance. This may cause performance _jitter_ on decode instances as the scheduled prefill will delay and interfere with other decodes. Furthermore, decode instances are typically configured with low-latency optimizations.
 
 !!! warning
     Using `kv_load_failure_policy="recompute"` can lead to performance degradation in production deployments. When KV loads fail, the decode instance will execute prefill work with decode-optimized configurations, which is inefficient and defeats the purpose of disaggregated prefilling. This also increases tail latency for other ongoing decode requests.

docs/features/quantization/bnb.md

@@ -7,7 +7,7 @@ Compared to other quantization methods, BitsAndBytes eliminates the need for cal
 Below are the steps to utilize BitsAndBytes with vLLM.
 
 ```bash
-pip install bitsandbytes>=0.46.1
+pip install bitsandbytes>=0.49.2
 ```
 
 vLLM reads the model's config file and supports both in-flight quantization and pre-quantized checkpoint.

docs/features/speculative_decoding/eagle.md

@@ -44,7 +44,7 @@ llm = LLM(
         "model": "RedHatAI/Llama-3.1-8B-Instruct-speculator.eagle3",
         "draft_tensor_parallel_size": 2,
         "num_speculative_tokens": 2,
-        "method": "eagle",
+        "method": "eagle3",
     },
 )
 

docs/models/pooling_models.md

@@ -382,6 +382,7 @@ ColQwen3 is based on [ColPali](https://arxiv.org/abs/2407.01449), which extends
 |---|---|---|
 | `ColQwen3` | Qwen3-VL | `TomoroAI/tomoro-colqwen3-embed-4b`, `TomoroAI/tomoro-colqwen3-embed-8b` |
 | `OpsColQwen3Model` | Qwen3-VL | `OpenSearch-AI/Ops-Colqwen3-4B`, `OpenSearch-AI/Ops-Colqwen3-8B` |
+| `Qwen3VLNemotronEmbedModel` | Qwen3-VL | `nvidia/nemotron-colembed-vl-4b-v2`, `nvidia/nemotron-colembed-vl-8b-v2` |
 
 Start the server:
 
@@ -389,7 +390,9 @@ Start the server:
 vllm serve TomoroAI/tomoro-colqwen3-embed-4b --max-model-len 4096
 ```
 
-Then you can use the rerank endpoint:
+#### Text-only scoring and reranking
+
+Use the `/rerank` endpoint:
 
 ```shell
 curl -s http://localhost:8000/rerank -H "Content-Type: application/json" -d '{
@@ -403,7 +406,7 @@ curl -s http://localhost:8000/rerank -H "Content-Type: application/json" -d '{
 }'
 ```
 
-Or the score endpoint:
+Or the `/score` endpoint:
 
 ```shell
 curl -s http://localhost:8000/score -H "Content-Type: application/json" -d '{
@@ -413,7 +416,57 @@ curl -s http://localhost:8000/score -H "Content-Type: application/json" -d '{
 }'
 ```
 
-You can also get the raw token embeddings using the pooling endpoint with `token_embed` task:
+#### Multi-modal scoring and reranking (text query × image documents)
+
+The `/score` and `/rerank` endpoints also accept multi-modal inputs directly.
+Pass image documents using the `data_1`/`data_2` (for `/score`) or `documents` (for `/rerank`) fields
+with a `content` list containing `image_url` and `text` parts — the same format used by the
+OpenAI chat completion API:
+
+Score a text query against image documents:
+
+```shell
+curl -s http://localhost:8000/score -H "Content-Type: application/json" -d '{
+    "model": "TomoroAI/tomoro-colqwen3-embed-4b",
+    "data_1": "Retrieve the city of Beijing",
+    "data_2": [
+        {
+            "content": [
+                {"type": "image_url", "image_url": {"url": "data:image/png;base64,<BASE64>"}},
+                {"type": "text", "text": "Describe the image."}
+            ]
+        }
+    ]
+}'
+```
+
+Rerank image documents by a text query:
+
+```shell
+curl -s http://localhost:8000/rerank -H "Content-Type: application/json" -d '{
+    "model": "TomoroAI/tomoro-colqwen3-embed-4b",
+    "query": "Retrieve the city of Beijing",
+    "documents": [
+        {
+            "content": [
+                {"type": "image_url", "image_url": {"url": "data:image/png;base64,<BASE64_1>"}},
+                {"type": "text", "text": "Describe the image."}
+            ]
+        },
+        {
+            "content": [
+                {"type": "image_url", "image_url": {"url": "data:image/png;base64,<BASE64_2>"}},
+                {"type": "text", "text": "Describe the image."}
+            ]
+        }
+    ],
+    "top_n": 2
+}'
+```
+
+#### Raw token embeddings
+
+You can also get the raw token embeddings using the `/pooling` endpoint with `token_embed` task:
 
 ```shell
 curl -s http://localhost:8000/pooling -H "Content-Type: application/json" -d '{
@@ -423,7 +476,7 @@ curl -s http://localhost:8000/pooling -H "Content-Type: application/json" -d '{
 }'
 ```
 
-For **image inputs**, use the chat-style `messages` field so that the vLLM multimodal processor handles them correctly:
+For **image inputs** via the pooling endpoint, use the chat-style `messages` field:
 
 ```shell
 curl -s http://localhost:8000/pooling -H "Content-Type: application/json" -d '{
@@ -440,10 +493,10 @@ curl -s http://localhost:8000/pooling -H "Content-Type: application/json" -d '{
 }'
 ```
 
-Examples can be found here:
+#### Examples
 
 - Multi-vector retrieval: [examples/pooling/token_embed/colqwen3_token_embed_online.py](../../examples/pooling/token_embed/colqwen3_token_embed_online.py)
-- Reranking: [examples/pooling/score/colqwen3_rerank_online.py](../../examples/pooling/score/colqwen3_rerank_online.py)
+- Reranking (text + multi-modal): [examples/pooling/score/colqwen3_rerank_online.py](../../examples/pooling/score/colqwen3_rerank_online.py)
 
 ### BAAI/bge-m3
 

docs/models/supported_models.md

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

examples/offline_inference/kv_load_failure_recovery/decode_example.py

@@ -42,6 +42,7 @@ def main():
                 "async_load": args.async_load,
             },
             kv_connector_module_path="load_recovery_example_connector",
+            kv_load_failure_policy="recompute",
         )
         out_file = (
             "async_decode_recovered_output.txt"

examples/online_serving/kv_events_subscriber.py

@@ -37,6 +37,12 @@ class BlockStored(KVCacheEvent):
     medium: str | None
     lora_name: str | None
 
+    extra_keys: list[tuple[Any, ...] | None] | None = None
+    """Extra keys used in block hash computation, one entry per block in
+    block_hashes. Each entry contains MM identifiers, LoRA name, cache_salt,
+    prompt embeddings data, etc. for that specific block.
+    """
+
 
 class BlockRemoved(KVCacheEvent):
     block_hashes: list[ExternalBlockHash]

examples/online_serving/opentelemetry/README.md

@@ -1,14 +1,6 @@
 # Setup OpenTelemetry POC
 
-1. Install OpenTelemetry packages:
-
-    ```bash
-    pip install \
-      'opentelemetry-sdk>=1.26.0,<1.27.0' \
-      'opentelemetry-api>=1.26.0,<1.27.0' \
-      'opentelemetry-exporter-otlp>=1.26.0,<1.27.0' \
-      'opentelemetry-semantic-conventions-ai>=0.4.1,<0.5.0'
-    ```
+> **Note:** The core OpenTelemetry packages (`opentelemetry-sdk`, `opentelemetry-api`, `opentelemetry-exporter-otlp`, `opentelemetry-semantic-conventions-ai`) are bundled with vLLM. Manual installation is not required.
 
 1. Start Jaeger in a docker container:
 

examples/pooling/score/colmodernvbert_rerank_online.py

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

examples/pooling/score/colqwen3_rerank_online.py

@@ -1,7 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# ruff: noqa: E501
 """
-Example of using ColQwen3 late interaction model for reranking.
+Example of using ColQwen3 late interaction model for reranking and scoring.
 
 ColQwen3 is a multi-modal ColBERT-style model based on Qwen3-VL.
 It produces per-token embeddings and uses MaxSim scoring for retrieval
@@ -14,13 +15,65 @@ Then run this script:
     python colqwen3_rerank_online.py
 """
 
+import base64
+from io import BytesIO
+
 import requests
+from PIL import Image
 
 MODEL = "TomoroAI/tomoro-colqwen3-embed-4b"
 BASE_URL = "http://127.0.0.1:8000"
 
 headers = {"accept": "application/json", "Content-Type": "application/json"}
 
+# ── Image helpers ──────────────────────────────────────────
+
+
+def load_image(url: str) -> Image.Image:
+    """Download an image from URL (handles Wikimedia 403)."""
+    for hdrs in (
+        {},
+        {"User-Agent": "Mozilla/5.0 (compatible; ColQwen3-demo/1.0)"},
+    ):
+        resp = requests.get(url, headers=hdrs, timeout=15)
+        if resp.status_code == 403:
+            continue
+        resp.raise_for_status()
+        return Image.open(BytesIO(resp.content)).convert("RGB")
+    raise RuntimeError(f"Could not fetch image from {url}")
+
+
+def encode_image_base64(image: Image.Image) -> str:
+    """Encode a PIL image to a base64 data URI."""
+    buf = BytesIO()
+    image.save(buf, format="PNG")
+    return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()
+
+
+def make_image_content(image_url: str, text: str = "Describe the image.") -> dict:
+    """Build a ScoreMultiModalParam dict from an image URL."""
+    image = load_image(image_url)
+    return {
+        "content": [
+            {
+                "type": "image_url",
+                "image_url": {"url": encode_image_base64(image)},
+            },
+            {"type": "text", "text": text},
+        ]
+    }
+
+
+# ── Sample image URLs ─────────────────────────────────────
+
+IMAGE_URLS = {
+    "beijing": "https://upload.wikimedia.org/wikipedia/commons/6/61/Beijing_skyline_at_night.JPG",
+    "london": "https://upload.wikimedia.org/wikipedia/commons/4/49/London_skyline.jpg",
+    "singapore": "https://upload.wikimedia.org/wikipedia/commons/2/27/Singapore_skyline_2022.jpg",
+}
+
+# ── Text-only examples ────────────────────────────────────
+
 
 def rerank_text():
     """Text-only reranking via /rerank endpoint."""
@@ -120,11 +173,86 @@ def score_text_top_n():
         print(f"  {response.text[:300]}")
 
 
+# ── Multi-modal examples (text query × image documents) ──
+
+
+def score_text_vs_images():
+    """Score a text query against image documents via /score."""
+    print()
+    print("=" * 60)
+    print("4. Multi-modal scoring: text query vs image docs (/score)")
+    print("=" * 60)
+
+    query = "Retrieve the city of Beijing"
+    labels = list(IMAGE_URLS.keys())
+    print(f"\n  Loading {len(labels)} images...")
+    image_contents = [make_image_content(IMAGE_URLS[name]) for name in labels]
+
+    data = {
+        "model": MODEL,
+        "data_1": query,
+        "data_2": image_contents,
+    }
+
+    response = requests.post(f"{BASE_URL}/score", headers=headers, json=data)
+
+    if response.status_code == 200:
+        result = response.json()
+        print(f'\n  Query: "{query}"\n')
+        for item in result["data"]:
+            idx = item["index"]
+            print(f"    Doc {idx} [{labels[idx]}] score={item['score']:.4f}")
+    else:
+        print(f"  Request failed: {response.status_code}")
+        print(f"  {response.text[:300]}")
+
+
+def rerank_text_vs_images():
+    """Rerank image documents by a text query via /rerank."""
+    print()
+    print("=" * 60)
+    print("5. Multi-modal reranking: text query vs image docs (/rerank)")
+    print("=" * 60)
+
+    query = "Retrieve the city of London"
+    labels = list(IMAGE_URLS.keys())
+    print(f"\n  Loading {len(labels)} images...")
+    image_contents = [make_image_content(IMAGE_URLS[name]) for name in labels]
+
+    data = {
+        "model": MODEL,
+        "query": query,
+        "documents": image_contents,
+        "top_n": 2,
+    }
+
+    response = requests.post(f"{BASE_URL}/rerank", headers=headers, json=data)
+
+    if response.status_code == 200:
+        result = response.json()
+        print(f'\n  Query: "{query}"')
+        print(f"  Top {data['top_n']} results:\n")
+        for item in result["results"]:
+            idx = item["index"]
+            print(f"    [{item['relevance_score']:.4f}] {labels[idx]}")
+    else:
+        print(f"  Request failed: {response.status_code}")
+        print(f"  {response.text[:300]}")
+
+
+# ── Main ──────────────────────────────────────────────────
+
+
 def main():
+    # Text-only
     rerank_text()
     score_text()
     score_text_top_n()
 
+    # Multi-modal (text query × image documents)
+    score_text_vs_images()
+    rerank_text_vs_images()
+
 
 if __name__ == "__main__":
     main()

requirements/common.txt

@@ -53,3 +53,7 @@ model-hosting-container-standards >= 0.1.13, < 1.0.0
 mcp
 grpcio
 grpcio-reflection
+opentelemetry-sdk >= 1.27.0
+opentelemetry-api >= 1.27.0
+opentelemetry-exporter-otlp >= 1.27.0
+opentelemetry-semantic-conventions-ai >= 0.4.1

requirements/cuda.txt

@@ -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

requirements/docs.txt

@@ -1,4 +1,4 @@
-mkdocs
+mkdocs<2.0.0
 mkdocs-api-autonav
 mkdocs-material
 mkdocstrings-python

requirements/nightly_torch_test.txt

@@ -28,12 +28,12 @@ num2words # required for smolvlm test
 opencv-python-headless >= 4.13.0 # required for video test
 datamodel_code_generator # required for minicpm3 test
 lm-eval[api]>=0.4.11 # required for model evaluation test
-mteb>=1.38.11, <2 # required for mteb test
+mteb[bm25s]>=2, <3 # required for mteb test
 transformers==4.57.5
 tokenizers==0.22.0
 schemathesis>=3.39.15 # Required for openai schema test.
 # quantization
-bitsandbytes>=0.46.1
+bitsandbytes>=0.49.2
 buildkite-test-collector==0.1.9
 
 

requirements/rocm-test.txt

@@ -70,7 +70,7 @@ ray[cgraph,default]>=2.48.0
 torchgeo==0.7.0
     # via terratorch
 # MTEB Benchmark Test
-mteb==2.1.2
+mteb[bm25s]>=2, <3
 
 # Utilities
 num2words==0.5.14
@@ -102,3 +102,7 @@ terratorch==1.2.2
 segmentation-models-pytorch==0.5.0
 # Required for Prithvi tests
 imagehash==4.3.2
+# Required for bitsandbytes quantization test
+bitsandbytes==0.49.2
+# Examples (tensorizer) tests
+tensorizer==2.10.1

requirements/rocm.txt

@@ -1,6 +1,11 @@
 # Common dependencies
 -r common.txt
 
+# The version of gRPC libraries should be consistent with each other
+grpcio==1.78.0
+grpcio-reflection==1.78.0
+grpcio-tools==1.78.0
+
 numba == 0.61.2 # Required for N-gram speculative decoding
 
 # Dependencies for AMD GPUs
@@ -14,5 +19,4 @@ setuptools>=77.0.3,<80.0.0
 setuptools-scm>=8
 runai-model-streamer[s3,gcs]==0.15.3
 conch-triton-kernels==1.2.1
-timm>=1.0.17
-grpcio-tools==1.78.0 # Should match `build.txt`
+timm>=1.0.17

requirements/test.in

@@ -41,14 +41,18 @@ transformers==4.57.5
 tokenizers==0.22.0
 schemathesis>=3.39.15 # Required for openai schema test.
 # quantization
-bitsandbytes==0.46.1
+bitsandbytes==0.49.2
 buildkite-test-collector==0.1.9
 
 
 genai_perf>=0.0.8
 tritonclient>=2.51.0
 
-grpcio-tools==1.78.0 # Should match `build.txt`
+# The version of gRPC libraries should be consistent with each other
+grpcio==1.78.0
+grpcio-reflection==1.78.0
+grpcio-tools==1.78.0
+
 arctic-inference == 0.1.1 # Required for suffix decoding test
 numba == 0.61.2 # Required for N-gram speculative decoding
 numpy

requirements/test.txt

@@ -66,7 +66,7 @@ backoff==2.2.1
     # via
     #   -r requirements/test.in
     #   schemathesis
-bitsandbytes==0.46.1
+bitsandbytes==0.49.2
     # via
     #   -r requirements/test.in
     #   lightning
@@ -287,9 +287,13 @@ greenlet==3.2.3
     # via sqlalchemy
 grpcio==1.78.0
     # via
+    #   -r requirements/test.in
+    #   grpcio-reflection
     #   grpcio-tools
     #   ray
     #   tensorboard
+grpcio-reflection==1.78.0
+    # via -r requirements/test.in
 grpcio-tools==1.78.0
     # via -r requirements/test.in
 h11==0.14.0
@@ -487,7 +491,7 @@ msgpack==1.1.0
     # via
     #   librosa
     #   ray
-mteb==2.1.2
+mteb==2.8.3
     # via -r requirements/test.in
 multidict==6.1.0
     # via
@@ -653,6 +657,7 @@ orjson==3.11.5
 packaging==24.2
     # via
     #   accelerate
+    #   bitsandbytes
     #   black
     #   datamodel-code-generator
     #   datasets
@@ -757,6 +762,7 @@ protobuf==6.33.2
     # via
     #   google-api-core
     #   googleapis-common-protos
+    #   grpcio-reflection
     #   grpcio-tools
     #   opentelemetry-proto
     #   proto-plus

tests/distributed/test_custom_all_reduce.py

@@ -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)

tests/entrypoints/openai/responses/conftest.py

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

tests/entrypoints/openai/responses/test_mcp_tools.py

@@ -1,6 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Integration tests for MCP tool support in the Responses API."""
 
+from __future__ import annotations
 
 import pytest
 import pytest_asyncio
@@ -10,11 +12,31 @@ from openai_harmony import ToolDescription, ToolNamespaceConfig
 from vllm.entrypoints.mcp.tool_server import MCPToolServer
 
 from ....utils import RemoteOpenAIServer
+from .conftest import (
+    BASE_TEST_ENV,
+    events_contain_type,
+    log_response_diagnostics,
+    retry_for_tool_call,
+    retry_streaming_for,
+    validate_streaming_event_stack,
+)
 
 MODEL_NAME = "openai/gpt-oss-20b"
 
+_BASE_SERVER_ARGS = [
+    "--enforce-eager",
+    "--tool-server",
+    "demo",
+    "--max_model_len",
+    "5000",
+]
 
-def test_get_tool_description():
+_PYTHON_TOOL_INSTRUCTION = (
+    "You must use the Python tool to execute code. Never simulate execution."
+)
+
+
+class TestMCPToolServerUnit:
     """Test MCPToolServer.get_tool_description filtering logic.
 
     Note: The wildcard "*" is normalized to None by
@@ -22,283 +44,240 @@ def test_get_tool_description():
     so we only test None and specific tool filtering here.
     See test_serving_responses.py for "*" normalization tests.
     """
-    pytest.importorskip("mcp")
 
-    server = MCPToolServer()
-    tool1 = ToolDescription.new(
-        name="tool1", description="First", parameters={"type": "object"}
-    )
-    tool2 = ToolDescription.new(
-        name="tool2", description="Second", parameters={"type": "object"}
-    )
-    tool3 = ToolDescription.new(
-        name="tool3", description="Third", parameters={"type": "object"}
-    )
+    def test_get_tool_description(self):
+        pytest.importorskip("mcp")
 
-    server.harmony_tool_descriptions = {
-        "test_server": ToolNamespaceConfig(
-            name="test_server", description="test", tools=[tool1, tool2, tool3]
+        server = MCPToolServer()
+        tool1 = ToolDescription.new(
+            name="tool1", description="First", parameters={"type": "object"}
+        )
+        tool2 = ToolDescription.new(
+            name="tool2", description="Second", parameters={"type": "object"}
+        )
+        tool3 = ToolDescription.new(
+            name="tool3", description="Third", parameters={"type": "object"}
         )
-    }
 
-    # Nonexistent server
-    assert server.get_tool_description("nonexistent") is None
+        server.harmony_tool_descriptions = {
+            "test_server": ToolNamespaceConfig(
+                name="test_server",
+                description="test",
+                tools=[tool1, tool2, tool3],
+            )
+        }
 
-    # None (no filter) - returns all tools
-    result = server.get_tool_description("test_server", allowed_tools=None)
-    assert len(result.tools) == 3
+        # Nonexistent server
+        assert server.get_tool_description("nonexistent") is None
 
-    # Filter to specific tools
-    result = server.get_tool_description(
-        "test_server", allowed_tools=["tool1", "tool3"]
-    )
-    assert len(result.tools) == 2
-    assert result.tools[0].name == "tool1"
-    assert result.tools[1].name == "tool3"
+        # None (no filter) - returns all tools
+        result = server.get_tool_description("test_server", allowed_tools=None)
+        assert len(result.tools) == 3
 
-    # Single tool
-    result = server.get_tool_description(
-        "test_server",
-        allowed_tools=["tool2"],
-    )
-    assert len(result.tools) == 1
-    assert result.tools[0].name == "tool2"
+        # Filter to specific tools
+        result = server.get_tool_description(
+            "test_server", allowed_tools=["tool1", "tool3"]
+        )
+        assert len(result.tools) == 2
+        assert result.tools[0].name == "tool1"
+        assert result.tools[1].name == "tool3"
 
-    # No matching tools - returns None
-    result = server.get_tool_description("test_server", allowed_tools=["nonexistent"])
-    assert result is None
+        # Single tool
+        result = server.get_tool_description("test_server", allowed_tools=["tool2"])
+        assert len(result.tools) == 1
+        assert result.tools[0].name == "tool2"
 
-    # Empty list - returns None
-    assert server.get_tool_description("test_server", allowed_tools=[]) is None
+        # No matching tools - returns None
+        result = server.get_tool_description(
+            "test_server", allowed_tools=["nonexistent"]
+        )
+        assert result is None
+
+        # Empty list - returns None
+        assert server.get_tool_description("test_server", allowed_tools=[]) is None
+
+    def test_builtin_tools_consistency(self):
+        """MCP_BUILTIN_TOOLS must match _BUILTIN_TOOL_TO_MCP_SERVER_LABEL values."""
+        from vllm.entrypoints.openai.parser.harmony_utils import (
+            _BUILTIN_TOOL_TO_MCP_SERVER_LABEL,
+            MCP_BUILTIN_TOOLS,
+        )
+
+        assert set(_BUILTIN_TOOL_TO_MCP_SERVER_LABEL.values()) == MCP_BUILTIN_TOOLS, (
+            f"MCP_BUILTIN_TOOLS {MCP_BUILTIN_TOOLS} does not match "
+            f"_BUILTIN_TOOL_TO_MCP_SERVER_LABEL values "
+            f"{set(_BUILTIN_TOOL_TO_MCP_SERVER_LABEL.values())}"
+        )
 
 
 class TestMCPEnabled:
     """Tests that require MCP tools to be enabled via environment variable."""
 
     @pytest.fixture(scope="class")
-    def monkeypatch_class(self):
-        from _pytest.monkeypatch import MonkeyPatch
-
-        mpatch = MonkeyPatch()
-        yield mpatch
-        mpatch.undo()
-
-    @pytest.fixture(scope="class")
-    def mcp_enabled_server(self, monkeypatch_class: pytest.MonkeyPatch):
-        args = ["--enforce-eager", "--tool-server", "demo"]
-
-        with monkeypatch_class.context() as m:
-            m.setenv("VLLM_ENABLE_RESPONSES_API_STORE", "1")
-            m.setenv("PYTHON_EXECUTION_BACKEND", "dangerously_use_uv")
-            m.setenv(
-                "VLLM_GPT_OSS_SYSTEM_TOOL_MCP_LABELS", "code_interpreter,container"
-            )
-            # Helps the model follow instructions better
-            m.setenv("VLLM_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS", "1")
-            with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
-                yield remote_server
+    def mcp_enabled_server(self):
+        env_dict = {
+            **BASE_TEST_ENV,
+            "VLLM_ENABLE_RESPONSES_API_STORE": "1",
+            "PYTHON_EXECUTION_BACKEND": "dangerously_use_uv",
+            "VLLM_GPT_OSS_SYSTEM_TOOL_MCP_LABELS": ("code_interpreter,container"),
+            "VLLM_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS": "1",
+        }
+        with RemoteOpenAIServer(
+            MODEL_NAME, list(_BASE_SERVER_ARGS), env_dict=env_dict
+        ) as remote_server:
+            yield remote_server
 
     @pytest_asyncio.fixture
-    async def mcp_enabled_client(self, mcp_enabled_server):
+    async def client(self, mcp_enabled_server):
         async with mcp_enabled_server.get_async_client() as async_client:
             yield async_client
 
+    @staticmethod
+    def _mcp_tools_payload(*, allowed_tools: list[str] | None = None) -> list[dict]:
+        tool: dict = {
+            "type": "mcp",
+            "server_label": "code_interpreter",
+            "server_url": "http://localhost:8888",
+        }
+        if allowed_tools is not None:
+            tool["allowed_tools"] = allowed_tools
+        return [tool]
+
+    @staticmethod
+    def _python_exec_input(code: str = "") -> str:
+        if not code:
+            code = "import random; print(random.randint(1, 1000000))"
+        return f"Execute the following code: {code}"
+
     @pytest.mark.asyncio
     @pytest.mark.parametrize("model_name", [MODEL_NAME])
-    async def test_mcp_tool_env_flag_enabled(
-        self, mcp_enabled_client: OpenAI, model_name: str
-    ):
-        response = await mcp_enabled_client.responses.create(
+    async def test_mcp_tool_env_flag_enabled(self, client: OpenAI, model_name: str):
+        response = await retry_for_tool_call(
+            client,
             model=model_name,
-            input=(
-                "Execute the following code: "
-                "import random; print(random.randint(1, 1000000))"
-            ),
-            instructions=(
-                "You must use the Python tool to execute code. "
-                "Never simulate execution."
-            ),
-            tools=[
-                {
-                    "type": "mcp",
-                    "server_label": "code_interpreter",
-                    # URL unused for DemoToolServer
-                    "server_url": "http://localhost:8888",
-                }
-            ],
+            expected_tool_type="mcp_call",
+            input=self._python_exec_input(),
+            instructions=_PYTHON_TOOL_INSTRUCTION,
+            tools=self._mcp_tools_payload(),
+            temperature=0.0,
             extra_body={"enable_response_messages": True},
         )
-        assert response is not None
+
         assert response.status == "completed"
-        # Verify output messages: Tool calls and responses on analysis channel
+        log_response_diagnostics(response, label="MCP Enabled")
+
         tool_call_found = False
         tool_response_found = False
         for message in response.output_messages:
             recipient = message.get("recipient")
             if recipient and recipient.startswith("python"):
                 tool_call_found = True
-                assert message.get("channel") == "analysis", (
-                    "Tool call should be on analysis channel"
-                )
+                assert message.get("channel") == "analysis"
             author = message.get("author", {})
-            if (
-                author.get("role") == "tool"
-                and author.get("name")
-                and author.get("name").startswith("python")
+            if author.get("role") == "tool" and (author.get("name") or "").startswith(
+                "python"
             ):
                 tool_response_found = True
-                assert message.get("channel") == "analysis", (
-                    "Tool response should be on analysis channel"
-                )
+                assert message.get("channel") == "analysis"
 
-        assert tool_call_found, "Should have found at least one Python tool call"
-        assert tool_response_found, (
-            "Should have found at least one Python tool response"
+        assert tool_call_found, (
+            f"No Python tool call found. "
+            f"Output types: "
+            f"{[getattr(o, 'type', None) for o in response.output]}"
         )
+        assert tool_response_found, "No Python tool response found"
+
         for message in response.input_messages:
-            assert message.get("author").get("role") != "developer", (
-                "No developer messages should be present with valid mcp tool"
-            )
+            assert message.get("author", {}).get("role") != "developer"
 
-    @pytest.mark.flaky(reruns=3)
     @pytest.mark.asyncio
     @pytest.mark.parametrize("model_name", [MODEL_NAME])
     async def test_mcp_tool_with_allowed_tools_star(
-        self, mcp_enabled_client: OpenAI, model_name: str
+        self, client: OpenAI, model_name: str
     ):
-        """Test MCP tool with allowed_tools=['*'] to select all available
-        tools.
-
-        This E2E test verifies that the "*" wildcard works end-to-end.
-        See test_serving_responses.py for detailed unit tests of "*"
-        normalization.
-        """
-        response = await mcp_enabled_client.responses.create(
+        response = await retry_for_tool_call(
+            client,
             model=model_name,
-            input=(
-                "Execute the following code: "
-                "import random; print(random.randint(1, 1000000))"
-            ),
-            instructions=(
-                "You must use the Python tool to execute code. "
-                "Never simulate execution."
-            ),
-            tools=[
-                {
-                    "type": "mcp",
-                    "server_label": "code_interpreter",
-                    "server_url": "http://localhost:8888",
-                    # Using "*" to allow all tools from this MCP server
-                    "allowed_tools": ["*"],
-                }
-            ],
+            expected_tool_type="mcp_call",
+            input=self._python_exec_input(),
+            instructions=_PYTHON_TOOL_INSTRUCTION,
+            tools=self._mcp_tools_payload(allowed_tools=["*"]),
+            temperature=0.0,
             extra_body={"enable_response_messages": True},
         )
-        assert response is not None
+
         assert response.status == "completed"
-        # Verify tool calls work with allowed_tools=["*"]
-        tool_call_found = False
-        for message in response.output_messages:
-            recipient = message.get("recipient")
-            if recipient and recipient.startswith("python"):
-                tool_call_found = True
-                break
+        log_response_diagnostics(response, label="MCP Allowed Tools *")
+
+        tool_call_found = any(
+            (msg.get("recipient") or "").startswith("python")
+            for msg in response.output_messages
+        )
         assert tool_call_found, (
-            "Should have found at least one Python tool call with '*'"
+            f"No Python tool call with '*'. "
+            f"Output types: "
+            f"{[getattr(o, 'type', None) for o in response.output]}"
         )
 
-    @pytest.mark.flaky(reruns=3)
     @pytest.mark.asyncio
     @pytest.mark.parametrize("model_name", [MODEL_NAME])
     async def test_mcp_tool_calling_streaming_types(
         self,
         pairs_of_event_types: dict[str, str],
-        mcp_enabled_client: OpenAI,
+        client: OpenAI,
         model_name: str,
     ):
-        tools = [
-            {
-                "type": "mcp",
-                "server_label": "code_interpreter",
-            }
-        ]
-        input_text = "What is 123 * 456? Use python to calculate the result."
+        def _has_mcp_events(events: list) -> bool:
+            return events_contain_type(events, "mcp_call")
 
-        stream_response = await mcp_enabled_client.responses.create(
+        events = await retry_streaming_for(
+            client,
             model=model_name,
-            input=input_text,
-            tools=tools,
-            stream=True,
-            instructions=(
-                "You must use the Python tool to execute code. "
-                "Never simulate execution."
-            ),
+            validate_events=_has_mcp_events,
+            input=("What is 123 * 456? Use Python to calculate the result."),
+            tools=[{"type": "mcp", "server_label": "code_interpreter"}],
+            instructions=_PYTHON_TOOL_INSTRUCTION,
+            temperature=0.0,
         )
 
-        stack_of_event_types = []
-        saw_mcp_type = False
-        async for event in stream_response:
-            if event.type == "response.created":
-                stack_of_event_types.append(event.type)
-            elif event.type == "response.completed":
-                assert stack_of_event_types[-1] == pairs_of_event_types[event.type]
-                stack_of_event_types.pop()
-            elif (
-                event.type.endswith("added")
-                or event.type == "response.mcp_call.in_progress"
-            ):
-                stack_of_event_types.append(event.type)
-            elif event.type.endswith("delta"):
-                if stack_of_event_types[-1] == event.type:
-                    continue
-                stack_of_event_types.append(event.type)
-            elif (
-                event.type.endswith("done")
-                or event.type == "response.mcp_call.completed"
-            ):
-                assert stack_of_event_types[-1] == pairs_of_event_types[event.type]
-                if "mcp_call" in event.type:
-                    saw_mcp_type = True
-                stack_of_event_types.pop()
+        validate_streaming_event_stack(events, pairs_of_event_types)
 
-        assert len(stack_of_event_types) == 0
-        assert saw_mcp_type, "Should have seen at least one mcp call"
+        assert events_contain_type(events, "mcp_call"), (
+            f"No mcp_call events after retries. "
+            f"Event types: {sorted({e.type for e in events})}"
+        )
 
 
 class TestMCPDisabled:
-    """Tests that verify behavior when MCP tools are disabled."""
+    """Tests that MCP tools are not executed when the env flag is unset."""
 
     @pytest.fixture(scope="class")
-    def monkeypatch_class(self):
-        from _pytest.monkeypatch import MonkeyPatch
-
-        mpatch = MonkeyPatch()
-        yield mpatch
-        mpatch.undo()
-
-    @pytest.fixture(scope="class")
-    def mcp_disabled_server(self, monkeypatch_class: pytest.MonkeyPatch):
-        args = ["--enforce-eager", "--tool-server", "demo"]
-
-        with monkeypatch_class.context() as m:
-            m.setenv("VLLM_ENABLE_RESPONSES_API_STORE", "1")
-            m.setenv("PYTHON_EXECUTION_BACKEND", "dangerously_use_uv")
-            # Helps the model follow instructions better
-            m.setenv("VLLM_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS", "1")
-            with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
-                yield remote_server
+    def mcp_disabled_server(self):
+        env_dict = {
+            **BASE_TEST_ENV,
+            "VLLM_ENABLE_RESPONSES_API_STORE": "1",
+            "PYTHON_EXECUTION_BACKEND": "dangerously_use_uv",
+            "VLLM_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS": "1",
+        }
+        with RemoteOpenAIServer(
+            MODEL_NAME, list(_BASE_SERVER_ARGS), env_dict=env_dict
+        ) as remote_server:
+            yield remote_server
 
     @pytest_asyncio.fixture
-    async def mcp_disabled_client(self, mcp_disabled_server):
+    async def client(self, mcp_disabled_server):
         async with mcp_disabled_server.get_async_client() as async_client:
             yield async_client
 
     @pytest.mark.asyncio
     @pytest.mark.parametrize("model_name", [MODEL_NAME])
-    async def test_mcp_tool_env_flag_disabled(
-        self, mcp_disabled_client: OpenAI, model_name: str
+    async def test_mcp_disabled_server_does_not_execute(
+        self, client: OpenAI, model_name: str
     ):
-        response = await mcp_disabled_client.responses.create(
+        """When MCP is disabled the model may still attempt tool calls
+        (tool descriptions can remain in the prompt), but the server
+        must NOT execute them."""
+        response = await client.responses.create(
             model=model_name,
             input=(
                 "Execute the following code if the tool is present: "
@@ -308,38 +287,35 @@ class TestMCPDisabled:
                 {
                     "type": "mcp",
                     "server_label": "code_interpreter",
-                    # URL unused for DemoToolServer
                     "server_url": "http://localhost:8888",
                 }
             ],
+            temperature=0.0,
             extra_body={"enable_response_messages": True},
         )
         assert response is not None
         assert response.status == "completed"
-        # Verify output messages: No tool calls and responses
-        tool_call_found = False
-        tool_response_found = False
+
+        log_response_diagnostics(response, label="MCP Disabled")
+
+        # Server must not have executed any tool calls
         for message in response.output_messages:
-            recipient = message.get("recipient")
-            if recipient and recipient.startswith("python"):
-                tool_call_found = True
-                assert message.get("channel") == "analysis", (
-                    "Tool call should be on analysis channel"
-                )
             author = message.get("author", {})
-            if (
+            assert not (
                 author.get("role") == "tool"
-                and author.get("name")
-                and author.get("name").startswith("python")
-            ):
-                tool_response_found = True
-                assert message.get("channel") == "analysis", (
-                    "Tool response should be on analysis channel"
+                and (author.get("name") or "").startswith("python")
+            ), (
+                "Server executed a python tool call even though MCP is "
+                f"disabled. Message: {message}"
+            )
+
+        # No completed mcp_call output items
+        for item in response.output:
+            if getattr(item, "type", None) == "mcp_call":
+                assert getattr(item, "status", None) != "completed", (
+                    "MCP call should not be completed when MCP is disabled"
                 )
 
-        assert not tool_call_found, "Should not have a python call"
-        assert not tool_response_found, "Should not have a tool response"
+        # No developer messages injected
         for message in response.input_messages:
-            assert message.get("author").get("role") != "developer", (
-                "No developer messages should be present without a valid tool"
-            )
+            assert message.get("author", {}).get("role") != "developer"

tests/entrypoints/openai/responses/test_parsable_context.py

@@ -3,15 +3,29 @@
 
 import importlib.util
 import json
+import logging
 
 import pytest
 import pytest_asyncio
 from openai import OpenAI
 
 from ....utils import RemoteOpenAIServer
+from .conftest import (
+    BASE_TEST_ENV,
+    has_output_type,
+    log_response_diagnostics,
+    retry_for_tool_call,
+)
+
+logger = logging.getLogger(__name__)
 
 MODEL_NAME = "Qwen/Qwen3-8B"
 
+_PYTHON_TOOL_INSTRUCTION = (
+    "You must use the Python tool to execute code. "
+    "Never simulate execution. You must print the final answer."
+)
+
 
 @pytest.fixture(scope="module")
 def server():
@@ -32,12 +46,12 @@ def server():
         "--tool-server",
         "demo",
     ]
-    env_dict = dict(
-        VLLM_ENABLE_RESPONSES_API_STORE="1",
-        VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT="1",
-        PYTHON_EXECUTION_BACKEND="dangerously_use_uv",
-    )
-
+    env_dict = {
+        **BASE_TEST_ENV,
+        "VLLM_ENABLE_RESPONSES_API_STORE": "1",
+        "VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT": "1",
+        "PYTHON_EXECUTION_BACKEND": "dangerously_use_uv",
+    }
     with RemoteOpenAIServer(MODEL_NAME, args, env_dict=env_dict) as remote_server:
         yield remote_server
 
@@ -54,6 +68,7 @@ async def test_basic(client: OpenAI, model_name: str):
     response = await client.responses.create(
         model=model_name,
         input="What is 123 * 456?",
+        temperature=0.0,
     )
     assert response is not None
     print("response: ", response)
@@ -99,10 +114,15 @@ async def test_reasoning_and_function_items(client: OpenAI, model_name: str):
     )
     assert response is not None
     assert response.status == "completed"
-    # make sure we get a reasoning and text output
-    assert response.output[0].type == "reasoning"
-    assert response.output[1].type == "message"
-    assert type(response.output[1].content[0].text) is str
+
+    output_types = [getattr(o, "type", None) for o in response.output]
+    assert "reasoning" in output_types, (
+        f"Expected reasoning in output, got: {output_types}"
+    )
+    assert "message" in output_types, f"Expected message in output, got: {output_types}"
+
+    msg = next(o for o in response.output if o.type == "message")
+    assert type(msg.content[0].text) is str
 
 
 def get_horoscope(sign):
@@ -110,10 +130,10 @@ def get_horoscope(sign):
 
 
 def call_function(name, args):
+    logger.info("Calling function %s with args %s", name, args)
     if name == "get_horoscope":
         return get_horoscope(**args)
-    else:
-        raise ValueError(f"Unknown function: {name}")
+    raise ValueError(f"Unknown function: {name}")
 
 
 @pytest.mark.asyncio
@@ -136,61 +156,111 @@ async def test_function_call_first_turn(client: OpenAI, model_name: str):
         }
     ]
 
-    response = await client.responses.create(
+    response = await retry_for_tool_call(
+        client,
         model=model_name,
+        expected_tool_type="function_call",
         input="What is the horoscope for Aquarius today?",
         tools=tools,
         temperature=0.0,
     )
     assert response is not None
     assert response.status == "completed"
-    assert len(response.output) == 2
-    assert response.output[0].type == "reasoning"
-    assert response.output[1].type == "function_call"
 
-    function_call = response.output[1]
+    output_types = [getattr(o, "type", None) for o in response.output]
+    assert "reasoning" in output_types, (
+        f"Expected reasoning in output, got: {output_types}"
+    )
+    assert has_output_type(response, "function_call"), (
+        f"Expected function_call in output, got: {output_types}"
+    )
+
+    function_call = next(o for o in response.output if o.type == "function_call")
     assert function_call.name == "get_horoscope"
     assert function_call.call_id is not None
 
     args = json.loads(function_call.arguments)
     assert "sign" in args
 
-    # the multi turn function call is tested above in
-    # test_reasoning_and_function_items
-
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("model_name", [MODEL_NAME])
 async def test_mcp_tool_call(client: OpenAI, model_name: str):
-    response = await client.responses.create(
+    """MCP tool calling with code_interpreter.
+
+    The model may make one or more tool calls before producing a final
+    message.  We validate server invariants (mcp_call items have correct
+    fields) with hard assertions.  Output indices are never hardcoded
+    since the model can produce multiple tool-call rounds.
+    """
+    # MCP + container init + code execution can be slow
+    client_with_timeout = client.with_options(timeout=client.timeout * 3)
+
+    response = await retry_for_tool_call(
+        client_with_timeout,
         model=model_name,
-        input="What is 123 * 456? Use python to calculate the result.",
+        expected_tool_type="mcp_call",
+        input=(
+            "What is 123 * 456? Use python to calculate the result. "
+            "Print the result with print()."
+        ),
         tools=[{"type": "code_interpreter", "container": {"type": "auto"}}],
-        extra_body={"enable_response_messages": True},
+        instructions=_PYTHON_TOOL_INSTRUCTION,
         temperature=0.0,
+        extra_body={"enable_response_messages": True},
     )
 
     assert response is not None
-    assert response.status == "completed"
 
-    # The model may produce multiple reasoning/mcp_call rounds before the
-    # final message, so validate structurally rather than by exact index.
-    output_types = [o.type for o in response.output]
-    assert "reasoning" in output_types
-    mcp_calls = [o for o in response.output if o.type == "mcp_call"]
-    assert len(mcp_calls) >= 1
-    assert type(mcp_calls[0].arguments) is str
-    assert type(mcp_calls[0].output) is str
+    output_types = [getattr(o, "type", None) for o in response.output]
+    log_response_diagnostics(response, label="test_mcp_tool_call")
 
-    # The final output should be a message containing the correct answer
-    assert response.output[-1].type == "message"
-    assert any(s in response.output[-1].content[0].text for s in ("56088", "56,088"))
+    assert response.status == "completed", (
+        f"Response status={response.status} "
+        f"(details={getattr(response, 'incomplete_details', None)}). "
+        f"Output types: {output_types}."
+    )
 
-    # Test raw input_messages / output_messages
-    assert len(response.input_messages) == 1
-    assert len(response.output_messages) >= 3
+    assert "reasoning" in output_types, (
+        f"Expected reasoning in output, got: {output_types}"
+    )
+    assert "mcp_call" in output_types, (
+        f"Expected mcp_call in output, got: {output_types}"
+    )
+
+    # Every mcp_call item must have well-typed fields
+    for item in response.output:
+        if getattr(item, "type", None) == "mcp_call":
+            assert type(item.arguments) is str, (
+                f"mcp_call.arguments should be str, got {type(item.arguments)}"
+            )
+            assert type(item.output) is str, (
+                f"mcp_call.output should be str, got {type(item.output)}"
+            )
+
+    # The model may make 1+ tool-call rounds but must still produce
+    # a final message for a trivial calculation like 123 * 456.
+    message_outputs = [
+        o for o in response.output if getattr(o, "type", None) == "message"
+    ]
+    assert message_outputs, (
+        f"Model did not produce a final message. Output types: {output_types}"
+    )
+
+    final_message = message_outputs[-1]
+    assert any(s in final_message.content[0].text for s in ("56088", "56,088")), (
+        f"Expected 56088 in final message, got: {final_message.content[0].text!r}"
+    )
+
+    # Validate raw input_messages / output_messages
+    assert len(response.input_messages) >= 1, "Expected at least 1 input message"
+    assert len(response.output_messages) >= 1, "Expected at least 1 output message"
     assert any(
-        s in response.output_messages[-1]["message"] for s in ("56088", "56,088")
+        any(s in str(msg) for s in ("56088", "56,088"))
+        for msg in response.output_messages
+    ), (
+        f"Expected 56088 in at least one output_message, "
+        f"got {len(response.output_messages)} messages"
     )
 
 
@@ -202,6 +272,7 @@ async def test_max_tokens(client: OpenAI, model_name: str):
         input="What is the first paragraph of Moby Dick?",
         reasoning={"effort": "low"},
         max_output_tokens=30,
+        temperature=0.0,
     )
     assert response is not None
     assert response.status == "incomplete"

tests/entrypoints/openai/responses/test_simple.py

@@ -12,13 +12,15 @@ MODEL_NAME = "Qwen/Qwen3-8B"
 
 @pytest.fixture(scope="module")
 def server():
-    args = ["--reasoning-parser", "qwen3", "--max_model_len", "5000"]
-    env_dict = dict(
-        VLLM_ENABLE_RESPONSES_API_STORE="1",
-        # uncomment for tool calling
-        # PYTHON_EXECUTION_BACKEND="dangerously_use_uv",
-    )
+    from .conftest import BASE_TEST_ENV
 
+    args = ["--reasoning-parser", "qwen3", "--max_model_len", "5000"]
+    env_dict = {
+        **BASE_TEST_ENV,
+        "VLLM_ENABLE_RESPONSES_API_STORE": "1",
+        # uncomment for tool calling
+        # PYTHON_EXECUTION_BACKEND: "dangerously_use_uv",
+    }
     with RemoteOpenAIServer(MODEL_NAME, args, env_dict=env_dict) as remote_server:
         yield remote_server
 

tests/entrypoints/openai/test_chat_error.py

@@ -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

tests/entrypoints/openai/test_realtime_validation.py

@@ -4,6 +4,7 @@
 import asyncio
 import base64
 import json
+import warnings
 
 import librosa
 import numpy as np
@@ -85,7 +86,41 @@ async def test_multi_chunk_streaming(
 
             await send_event(ws, {"type": "session.update", "model": model_name})
 
-            # Send commit to start transcription
+            # Wait for the server to acknowledge the session update.
+            try:
+                while True:
+                    event = await receive_event(ws, timeout=5.0)
+                    if event["type"] == "session.updated":
+                        break
+            except TimeoutError:
+                warnings.warn(
+                    f"session.updated not received within {5.0}s after "
+                    "session.update. The server may not implement this event.",
+                    stacklevel=2,
+                )
+
+            # (ROCm) Warm-up: send a non-final commit (required to start
+            # transcription) with a small audio chunk to trigger aiter
+            # compilation on first use.
+            await send_event(ws, {"type": "input_audio_buffer.commit"})
+            await send_event(
+                ws,
+                {
+                    "type": "input_audio_buffer.append",
+                    "audio": mary_had_lamb_audio_chunks[0],
+                },
+            )
+            await send_event(ws, {"type": "input_audio_buffer.commit", "final": True})
+
+            # (ROCm) Drain all warm-up responses with generous timeout for
+            # JIT compilation
+            warmup_done = False
+            while not warmup_done:
+                event = await receive_event(ws, timeout=360.0)
+                if event["type"] in ("transcription.done", "error"):
+                    warmup_done = True
+
+            # Now send the real test audio
             await send_event(ws, {"type": "input_audio_buffer.commit"})
 
             # Send multiple audio chunks
@@ -153,6 +188,18 @@ async def test_empty_commit_does_not_crash_engine(
 
             await send_event(ws, {"type": "session.update", "model": model_name})
 
+            try:
+                while True:
+                    event = await receive_event(ws, timeout=5.0)
+                    if event["type"] == "session.updated":
+                        break
+            except TimeoutError:
+                warnings.warn(
+                    f"session.updated not received within {5.0}s after "
+                    "session.update. The server may not implement this event.",
+                    stacklevel=2,
+                )
+
             # Start generation without sending any audio
             await send_event(ws, {"type": "input_audio_buffer.commit"})
 
@@ -161,7 +208,8 @@ async def test_empty_commit_does_not_crash_engine(
 
             # We should get *some* response (error or empty transcription),
             # but the engine must NOT crash.
-            event = await receive_event(ws, timeout=30.0)
+            # (ROCm) Use generous timeout for first request (aiter JIT compilation)
+            event = await receive_event(ws, timeout=360.0)
             assert event["type"] in (
                 "error",
                 "transcription.done",
@@ -176,6 +224,19 @@ async def test_empty_commit_does_not_crash_engine(
 
             await send_event(ws, {"type": "session.update", "model": model_name})
 
+            try:
+                while True:
+                    event = await receive_event(ws, timeout=5.0)
+                    if event["type"] == "session.updated":
+                        break
+            except TimeoutError:
+                warnings.warn(
+                    f"session.updated not received within {5.0}s after "
+                    "session.update. The server may not implement this event.",
+                    stacklevel=2,
+                )
+
+            # Start transcription
             await send_event(ws, {"type": "input_audio_buffer.commit"})
 
             for chunk in mary_had_lamb_audio_chunks:

tests/entrypoints/openai/test_serving_chat.py

@@ -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
 

tests/entrypoints/openai/test_transcription_validation_whisper.py

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

tests/entrypoints/pooling/embed/test_online.py

@@ -58,13 +58,19 @@ if current_platform.is_rocm():
     torch.backends.cuda.enable_mem_efficient_sdp(False)
     torch.backends.cuda.enable_math_sdp(True)
 
+# On ROCm, floating-point reductions in attention and GEMM kernels are
+# non-associative and sensitive to batch geometry. Force LLM instances
+# into an identical, deterministic execution mode:
+ROCM_DETERMINISM_ARGS: list[str] = (
+    ["--max-num-seqs", "1"] if current_platform.is_rocm() else []
+)
+
 
 @pytest.fixture(scope="module")
 def server():
     args = [
         "--runner",
         "pooling",
-        # use half precision for speed and memory savings in CI environment
         "--dtype",
         DTYPE,
         "--enforce-eager",
@@ -72,12 +78,9 @@ def server():
         "512",
         "--chat-template",
         DUMMY_CHAT_TEMPLATE,
+        *ROCM_DETERMINISM_ARGS,
     ]
 
-    # ROCm: Use Flex Attention to support encoder-only self-attention.
-    if current_platform.is_rocm():
-        args.extend(["--attention-backend", "FLEX_ATTENTION"])
-
     with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
         yield remote_server
 
@@ -343,8 +346,15 @@ async def test_chat_request(
     assert chat_embeddings.id is not None
     assert completion_embeddings.id is not None
     assert chat_embeddings.created <= completion_embeddings.created
-    assert chat_embeddings.model_dump(exclude={"id", "created"}) == (
-        completion_embeddings.model_dump(exclude={"id", "created"})
+    # Use tolerance-based comparison for embeddings
+    check_embeddings_close(
+        embeddings_0_lst=[d.embedding for d in chat_embeddings.data],
+        embeddings_1_lst=[d.embedding for d in completion_embeddings.data],
+        name_0="chat",
+        name_1="completion",
+    )
+    assert chat_embeddings.model_dump(exclude={"id", "created", "data"}) == (
+        completion_embeddings.model_dump(exclude={"id", "created", "data"})
     )
 
     # test add_generation_prompt

tests/kernels/helion/test_silu_mul_fp8.py

@@ -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])

tests/kernels/helion/test_utils.py

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

tests/kernels/moe/test_grouped_topk.py

@@ -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)

tests/models/language/pooling/test_colbert.py

@@ -20,6 +20,12 @@ COLBERT_MODELS = {
         "colbert_dim": 96,
         "max_model_len": 512,
         "extra_kwargs": {},
+        "hf_comparison": {
+            "weights_file": "model.safetensors",
+            "weights_key": "linear.weight",
+            "trust_remote_code": False,
+            "model_cls": "BertModel",
+        },
     },
     "modernbert": {
         "model": "lightonai/GTE-ModernColBERT-v1",
@@ -30,6 +36,12 @@ COLBERT_MODELS = {
                 "architectures": ["ColBERTModernBertModel"],
             },
         },
+        "hf_comparison": {
+            "weights_file": "1_Dense/model.safetensors",
+            "weights_key": "linear.weight",
+            "trust_remote_code": False,
+            "model_cls": "AutoModel",
+        },
     },
     "jina": {
         "model": "jinaai/jina-colbert-v2",
@@ -40,9 +52,16 @@ COLBERT_MODELS = {
                 "architectures": ["ColBERTJinaRobertaModel"],
             },
         },
+        "hf_comparison": {
+            "weights_file": "model.safetensors",
+            "weights_key": "linear.weight",
+            "trust_remote_code": True,
+            "model_cls": "AutoModel",
+        },
     },
 }
 
+
 TEXTS_1 = [
     "What is the capital of France?",
     "What is the capital of Germany?",
@@ -56,9 +75,68 @@ TEXTS_2 = [
 DTYPE = "half"
 
 
-# -----------------------------------------------------------------------
-# Fixtures
-# -----------------------------------------------------------------------
+def _load_hf_model(model_name: str, hf_spec: dict, device: torch.device):
+    """Load HF model on the given device with a compatible attention impl."""
+    from transformers import AutoModel, BertModel
+
+    cls = BertModel if hf_spec["model_cls"] == "BertModel" else AutoModel
+    trust = hf_spec.get("trust_remote_code", False)
+
+    # Flash / Triton kernels require GPU tensors; fall back to eager on CPU.
+    extra = {}
+    if device.type == "cpu":
+        extra["attn_implementation"] = "eager"
+
+    model = cls.from_pretrained(
+        model_name,
+        trust_remote_code=trust,
+        **extra,
+    ).to(device)
+    model.eval()
+    return model
+
+
+def _load_projection_weight(model_name: str, hf_spec: dict, device: torch.device):
+    """Download and return the ColBERT linear projection weight."""
+    from huggingface_hub import hf_hub_download
+    from safetensors.torch import load_file
+
+    path = hf_hub_download(model_name, filename=hf_spec["weights_file"])
+    weights = load_file(path)
+    return weights[hf_spec["weights_key"]].to(device)
+
+
+def _compute_hf_colbert_embeddings(model, tokenizer, linear_weight, texts, device):
+    """Run HF model + projection and return L2-normalised token embeddings."""
+    import torch.nn.functional as F
+
+    embeddings = []
+    for text in texts:
+        inputs = tokenizer(text, return_tensors="pt").to(device)
+        with torch.no_grad():
+            hidden = model(**inputs).last_hidden_state.float()
+            projected = F.linear(hidden, linear_weight.float())
+            normalised = F.normalize(projected, p=2, dim=-1)
+            embeddings.append(normalised.squeeze(0).cpu())
+    return embeddings
+
+
+def _assert_embeddings_close(vllm_outputs, hf_embeddings):
+    """Assert that vLLM and HuggingFace embeddings match."""
+    for i, (hf_emb, vllm_out) in enumerate(zip(hf_embeddings, vllm_outputs)):
+        vllm_emb = torch.as_tensor(vllm_out).float()
+
+        assert hf_emb.shape == vllm_emb.shape, (
+            f"Shape mismatch for text {i}: HF {hf_emb.shape} vs vLLM {vllm_emb.shape}"
+        )
+
+        torch.testing.assert_close(
+            vllm_emb,
+            hf_emb,
+            rtol=1e-2,
+            atol=1e-2,
+            msg=f"Embedding mismatch for text {i}",
+        )
 
 
 @pytest.fixture(params=list(COLBERT_MODELS.keys()), scope="module")
@@ -87,11 +165,6 @@ def colbert_extra_kwargs(colbert_spec):
     return colbert_spec["extra_kwargs"]
 
 
-# -----------------------------------------------------------------------
-# Tests
-# -----------------------------------------------------------------------
-
-
 def test_colbert_token_embed(
     vllm_runner,
     colbert_model_name,
@@ -111,7 +184,7 @@ def test_colbert_token_embed(
         outputs = vllm_model.token_embed([TEXTS_1[0]])
 
         assert len(outputs) == 1
-        emb = torch.tensor(outputs[0])
+        emb = torch.as_tensor(outputs[0])
         assert emb.dim() == 2
         assert emb.shape[1] == colbert_dim
         assert emb.shape[0] > 1
@@ -135,8 +208,8 @@ def test_colbert_late_interaction_1_to_1(
         q_outputs = vllm_model.token_embed([TEXTS_1[0]])
         d_outputs = vllm_model.token_embed([TEXTS_2[0]])
 
-        q_emb = torch.tensor(q_outputs[0])
-        d_emb = torch.tensor(d_outputs[0])
+        q_emb = torch.as_tensor(q_outputs[0])
+        d_emb = torch.as_tensor(d_outputs[0])
 
         manual_score = compute_maxsim_score(q_emb, d_emb).item()
 
@@ -164,11 +237,11 @@ def test_colbert_late_interaction_1_to_N(
         q_outputs = vllm_model.token_embed([TEXTS_1[0]])
         d_outputs = vllm_model.token_embed(TEXTS_2)
 
-        q_emb = torch.tensor(q_outputs[0])
+        q_emb = torch.as_tensor(q_outputs[0])
 
         manual_scores = []
         for d_out in d_outputs:
-            d_emb = torch.tensor(d_out)
+            d_emb = torch.as_tensor(d_out)
             manual_scores.append(compute_maxsim_score(q_emb, d_emb).item())
 
         vllm_scores = vllm_model.score(TEXTS_1[0], TEXTS_2)
@@ -198,8 +271,8 @@ def test_colbert_late_interaction_N_to_N(
 
         manual_scores = []
         for q_out, d_out in zip(q_outputs, d_outputs):
-            q_emb = torch.tensor(q_out)
-            d_emb = torch.tensor(d_out)
+            q_emb = torch.as_tensor(q_out)
+            d_emb = torch.as_tensor(d_out)
             manual_scores.append(compute_maxsim_score(q_emb, d_emb).item())
 
         vllm_scores = vllm_model.score(TEXTS_1, TEXTS_2)
@@ -259,79 +332,16 @@ def test_colbert_embed_not_supported(
         vllm_model.embed([TEXTS_1[0]])
 
 
-# -----------------------------------------------------------------------
-# Per-model HuggingFace comparison tests
-# -----------------------------------------------------------------------
+@pytest.mark.parametrize("backend", list(COLBERT_MODELS.keys()))
+def test_colbert_hf_comparison(vllm_runner, backend):
+    """Test that vLLM ColBERT embeddings match HuggingFace for each backend."""
+    from transformers import AutoTokenizer
 
-
-def _assert_embeddings_close(vllm_outputs, hf_embeddings):
-    """Assert that vLLM and HuggingFace embeddings match."""
-    for i, (hf_emb, vllm_out) in enumerate(zip(hf_embeddings, vllm_outputs)):
-        vllm_emb = torch.tensor(vllm_out).float()
-
-        assert hf_emb.shape == vllm_emb.shape, (
-            f"Shape mismatch for text {i}: HF {hf_emb.shape} vs vLLM {vllm_emb.shape}"
-        )
-
-        torch.testing.assert_close(
-            vllm_emb,
-            hf_emb,
-            rtol=1e-2,
-            atol=1e-2,
-            msg=f"Embedding mismatch for text {i}",
-        )
-
-
-def test_colbert_hf_comparison_bert(vllm_runner):
-    """Test that vLLM ColBERT produces same embeddings as HuggingFace (BERT)."""
-    import torch.nn.functional as F
-    from huggingface_hub import hf_hub_download
-    from safetensors.torch import load_file
-    from transformers import AutoTokenizer, BertModel
-
-    model_name = COLBERT_MODELS["bert"]["model"]
-    test_texts = [TEXTS_1[0], TEXTS_2[0]]
-
-    with vllm_runner(
-        model_name,
-        runner="pooling",
-        dtype="float32",
-        max_model_len=512,
-        enforce_eager=True,
-    ) as vllm_model:
-        vllm_outputs = vllm_model.token_embed(test_texts)
-
-    hf_tokenizer = AutoTokenizer.from_pretrained(model_name)
-    hf_bert = BertModel.from_pretrained(model_name)
-    hf_bert.eval()
-
-    weights_path = hf_hub_download(model_name, filename="model.safetensors")
-    weights = load_file(weights_path)
-    linear_weight = weights["linear.weight"]  # [96, 384]
-
-    hf_embeddings = []
-    for text in test_texts:
-        inputs = hf_tokenizer(text, return_tensors="pt")
-        with torch.no_grad():
-            outputs = hf_bert(**inputs)
-            hidden_states = outputs.last_hidden_state
-            token_emb = F.linear(hidden_states, linear_weight)
-            token_emb = F.normalize(token_emb, p=2, dim=-1)
-            hf_embeddings.append(token_emb.squeeze(0).float())
-
-    _assert_embeddings_close(vllm_outputs, hf_embeddings)
-
-
-def test_colbert_hf_comparison_modernbert(vllm_runner):
-    """Test that vLLM ColBERT produces same embeddings as HuggingFace
-    (ModernBERT)."""
-    import torch.nn.functional as F
-    from huggingface_hub import hf_hub_download
-    from safetensors.torch import load_file
-    from transformers import AutoModel, AutoTokenizer
-
-    spec = COLBERT_MODELS["modernbert"]
+    spec = COLBERT_MODELS[backend]
+    hf_spec = spec["hf_comparison"]
     model_name = spec["model"]
+    assert isinstance(model_name, str)
+    assert isinstance(hf_spec, dict)
     test_texts = [TEXTS_1[0], TEXTS_2[0]]
 
     with vllm_runner(
@@ -344,73 +354,21 @@ def test_colbert_hf_comparison_modernbert(vllm_runner):
     ) as vllm_model:
         vllm_outputs = vllm_model.token_embed(test_texts)
 
-    hf_tokenizer = AutoTokenizer.from_pretrained(model_name)
-    hf_model = AutoModel.from_pretrained(model_name)
-    hf_model.eval()
-
-    # Load projection from sentence-transformers 1_Dense layer
-    dense_path = hf_hub_download(model_name, filename="1_Dense/model.safetensors")
-    dense_weights = load_file(dense_path)
-    linear_weight = dense_weights["linear.weight"]  # [128, 768]
-
-    hf_embeddings = []
-    for text in test_texts:
-        inputs = hf_tokenizer(text, return_tensors="pt")
-        with torch.no_grad():
-            outputs = hf_model(**inputs)
-            hidden_states = outputs.last_hidden_state
-            token_emb = F.linear(hidden_states, linear_weight)
-            token_emb = F.normalize(token_emb, p=2, dim=-1)
-            hf_embeddings.append(token_emb.squeeze(0).float())
-
-    _assert_embeddings_close(vllm_outputs, hf_embeddings)
-
-
-def test_colbert_hf_comparison_jina(vllm_runner):
-    """Test that vLLM ColBERT produces same embeddings as HuggingFace
-    (Jina XLM-RoBERTa)."""
-    import torch.nn.functional as F
-    from huggingface_hub import hf_hub_download
-    from safetensors.torch import load_file
-    from transformers import AutoModel, AutoTokenizer
-
-    spec = COLBERT_MODELS["jina"]
-    model_name = spec["model"]
-    test_texts = [TEXTS_1[0], TEXTS_2[0]]
-
-    with vllm_runner(
-        model_name,
-        runner="pooling",
-        dtype="float32",
-        max_model_len=spec["max_model_len"],
-        enforce_eager=True,
-        **spec["extra_kwargs"],
-    ) as vllm_model:
-        vllm_outputs = vllm_model.token_embed(test_texts)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
     hf_tokenizer = AutoTokenizer.from_pretrained(
         model_name,
-        trust_remote_code=True,
+        trust_remote_code=hf_spec.get("trust_remote_code", False),
     )
-    hf_model = AutoModel.from_pretrained(
-        model_name,
-        trust_remote_code=True,
+    hf_model = _load_hf_model(model_name, hf_spec, device)
+    linear_weight = _load_projection_weight(model_name, hf_spec, device)
+
+    hf_embeddings = _compute_hf_colbert_embeddings(
+        hf_model,
+        hf_tokenizer,
+        linear_weight,
+        test_texts,
+        device,
     )
-    hf_model.eval()
-
-    # Load projection from main checkpoint
-    weights_path = hf_hub_download(model_name, filename="model.safetensors")
-    weights = load_file(weights_path)
-    linear_weight = weights["linear.weight"]  # [128, 1024]
-
-    hf_embeddings = []
-    for text in test_texts:
-        inputs = hf_tokenizer(text, return_tensors="pt")
-        with torch.no_grad():
-            outputs = hf_model(**inputs)
-            hidden_states = outputs.last_hidden_state
-            token_emb = F.linear(hidden_states.float(), linear_weight.float())
-            token_emb = F.normalize(token_emb, p=2, dim=-1)
-            hf_embeddings.append(token_emb.squeeze(0).float())
 
     _assert_embeddings_close(vllm_outputs, hf_embeddings)

tests/models/language/pooling_mteb_test/mteb_score_utils.py

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

tests/models/multimodal/generation/test_whisper.py

@@ -111,6 +111,47 @@ def check_model_available(model: str) -> None:
     model_info.check_transformers_version(on_fail="skip")
 
 
+def test_parse_language_detection_output():
+    """Unit test for WhisperForConditionalGeneration.parse_language_detection_output.
+
+    No GPU or model loading required.
+    """
+    from unittest.mock import MagicMock
+
+    from vllm.model_executor.models.whisper import (
+        WhisperForConditionalGeneration,
+    )
+
+    cls = WhisperForConditionalGeneration
+
+    def make_tokenizer(return_value: str) -> MagicMock:
+        tok = MagicMock()
+        tok.decode = MagicMock(return_value=return_value)
+        return tok
+
+    # English
+    assert (
+        cls.parse_language_detection_output([50259], make_tokenizer("<|en|>")) == "en"
+    )
+
+    # German
+    assert (
+        cls.parse_language_detection_output([50261], make_tokenizer("<|de|>")) == "de"
+    )
+
+    # Unsupported language code
+    with pytest.raises(AssertionError):
+        cls.parse_language_detection_output([99999], make_tokenizer("<|xx|>"))
+
+    # No special token format
+    with pytest.raises(AssertionError):
+        cls.parse_language_detection_output([1], make_tokenizer("hello"))
+
+    # Empty token_ids
+    with pytest.raises((AssertionError, IndexError)):
+        cls.parse_language_detection_output([], make_tokenizer("anything"))
+
+
 @pytest.mark.core_model
 @pytest.mark.cpu_model
 @pytest.mark.parametrize("model", ["openai/whisper-large-v3-turbo"])

tests/models/multimodal/pooling/test_colmodernvbert.py

@@ -0,0 +1,115 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for ColModernVBERT multimodal late-interaction model.
+
+ColModernVBERT combines SigLIP vision encoder + ModernBERT text encoder
+with a pixel shuffle connector and ColBERT-style 128-dim per-token
+embeddings for visual document retrieval.
+"""
+
+import pytest
+import torch
+
+from vllm.entrypoints.pooling.score.utils import compute_maxsim_score
+
+MODEL_NAME = "ModernVBERT/colmodernvbert-merged"
+COLBERT_DIM = 128
+DTYPE = "half"
+
+
+# -----------------------------------------------------------------------
+# Text-only tests
+# -----------------------------------------------------------------------
+
+
+def test_colmodernvbert_text_token_embed(vllm_runner):
+    """Text query produces per-token embeddings with shape (seq_len, 128)."""
+    with vllm_runner(
+        MODEL_NAME,
+        runner="pooling",
+        dtype=DTYPE,
+        enforce_eager=True,
+    ) as vllm_model:
+        outputs = vllm_model.token_embed(["What is machine learning?"])
+
+        assert len(outputs) == 1
+        emb = torch.tensor(outputs[0])
+        assert emb.dim() == 2
+        assert emb.shape[1] == COLBERT_DIM
+        assert emb.shape[0] > 1
+
+
+def test_colmodernvbert_text_relevance_ordering(vllm_runner):
+    """Relevant documents score higher than irrelevant ones."""
+    query = "What is machine learning?"
+    documents = [
+        "Machine learning is a subset of artificial intelligence.",
+        "The weather in Paris is mild in spring.",
+    ]
+
+    with vllm_runner(
+        MODEL_NAME,
+        runner="pooling",
+        dtype=DTYPE,
+        enforce_eager=True,
+    ) as vllm_model:
+        scores = vllm_model.score(query, documents)
+
+        assert len(scores) == 2
+        assert scores[0] > scores[1], "ML doc should score higher than weather doc"
+
+
+def test_colmodernvbert_text_late_interaction(vllm_runner):
+    """MaxSim scoring via vLLM matches manual computation."""
+    query = "What is the capital of France?"
+    doc = "The capital of France is Paris."
+
+    with vllm_runner(
+        MODEL_NAME,
+        runner="pooling",
+        dtype=DTYPE,
+        enforce_eager=True,
+    ) as vllm_model:
+        q_out = vllm_model.token_embed([query])
+        d_out = vllm_model.token_embed([doc])
+
+        q_emb = torch.tensor(q_out[0])
+        d_emb = torch.tensor(d_out[0])
+        manual_score = compute_maxsim_score(q_emb, d_emb).item()
+
+        vllm_scores = vllm_model.score(query, doc)
+
+        assert len(vllm_scores) == 1
+        assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01)
+
+
+# -----------------------------------------------------------------------
+# Image tests
+# -----------------------------------------------------------------------
+
+
+def test_colmodernvbert_image_token_embed(vllm_runner, image_assets):
+    """Image input produces per-token embeddings including vision tokens."""
+    with vllm_runner(
+        MODEL_NAME,
+        runner="pooling",
+        dtype=DTYPE,
+        enforce_eager=True,
+    ) as vllm_model:
+        image = image_assets[0].pil_image
+        inputs = vllm_model.get_inputs(
+            [""],
+            images=[image],
+        )
+        req_outputs = vllm_model.llm.encode(
+            inputs,
+            pooling_task="token_embed",
+        )
+        outputs = [req_output.outputs.data for req_output in req_outputs]
+
+        assert len(outputs) == 1
+        emb = torch.tensor(outputs[0])
+        assert emb.dim() == 2
+        assert emb.shape[1] == COLBERT_DIM
+        # Should have at least the image tokens (64 after pixel shuffle)
+        assert emb.shape[0] >= 64

tests/models/multimodal/pooling/test_colqwen3.py

@@ -7,19 +7,31 @@ ColBERT-style late interaction scoring (MaxSim). It produces per-token
 embeddings for both text and image inputs.
 """
 
+import base64
+from io import BytesIO
+
 import pytest
 import torch
+from PIL import Image
+
+from vllm.entrypoints.chat_utils import (
+    ChatCompletionContentPartImageParam,
+    ChatCompletionContentPartTextParam,
+)
+from vllm.entrypoints.pooling.score.utils import ScoreMultiModalParam
 
 from ....conftest import VllmRunner
 
 MODELS = [
     "TomoroAI/tomoro-colqwen3-embed-4b",
     "OpenSearch-AI/Ops-Colqwen3-4B",
+    "nvidia/nemotron-colembed-vl-4b-v2",
 ]
 
 EMBED_DIMS = {
     "TomoroAI/tomoro-colqwen3-embed-4b": 320,
     "OpenSearch-AI/Ops-Colqwen3-4B": 2560,
+    "nvidia/nemotron-colembed-vl-4b-v2": 2560,
 }
 
 TEXT_QUERIES = [
@@ -33,6 +45,43 @@ TEXT_DOCUMENTS = [
 ]
 
 DTYPE = "half"
+GPU_MEMORY_UTILIZATION = 0.7
+
+
+def _make_base64_image(
+    width: int = 64, height: int = 64, color: tuple[int, int, int] = (255, 0, 0)
+) -> str:
+    """Create a small solid-color PNG image and return its base64 data URI."""
+    img = Image.new("RGB", (width, height), color)
+    buf = BytesIO()
+    img.save(buf, format="PNG")
+    b64 = base64.b64encode(buf.getvalue()).decode()
+    return f"data:image/png;base64,{b64}"
+
+
+def _make_image_mm_param(
+    image_uri: str,
+    text: str | None = None,
+) -> ScoreMultiModalParam:
+    """Build a ScoreMultiModalParam containing an image (and optional text)."""
+    content: list = [
+        ChatCompletionContentPartImageParam(
+            type="image_url",
+            image_url={"url": image_uri},
+        ),
+    ]
+    if text is not None:
+        content.append(
+            ChatCompletionContentPartTextParam(type="text", text=text),
+        )
+    return ScoreMultiModalParam(content=content)
+
+
+def _make_text_mm_param(text: str) -> ScoreMultiModalParam:
+    """Build a ScoreMultiModalParam containing only text."""
+    return ScoreMultiModalParam(
+        content=[ChatCompletionContentPartTextParam(type="text", text=text)],
+    )
 
 
 def _run_token_embed_test(
@@ -48,6 +97,7 @@ def _run_token_embed_test(
         dtype=dtype,
         max_model_len=4096,
         enforce_eager=True,
+        gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
     ) as vllm_model:
         outputs = vllm_model.token_embed([TEXT_QUERIES[0]])
 
@@ -83,6 +133,7 @@ def _run_late_interaction_test(
         dtype=dtype,
         max_model_len=4096,
         enforce_eager=True,
+        gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
     ) as vllm_model:
         q_outputs = vllm_model.token_embed([TEXT_QUERIES[0]])
         d_outputs = vllm_model.token_embed([TEXT_DOCUMENTS[0]])
@@ -118,6 +169,7 @@ def _run_relevance_test(
         dtype=dtype,
         max_model_len=4096,
         enforce_eager=True,
+        gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
     ) as vllm_model:
         scores = vllm_model.score(query, documents)
 
@@ -154,3 +206,142 @@ def test_colqwen3_relevance_ordering(
     dtype: str,
 ) -> None:
     _run_relevance_test(vllm_runner, model, dtype=dtype)
+
+
+# ── Multimodal scoring tests ────────────────────────────────
+
+
+def _run_multimodal_text_query_image_docs_test(
+    vllm_runner: type[VllmRunner],
+    model: str,
+    *,
+    dtype: str,
+) -> None:
+    """Score a text query against image documents via the multimodal path.
+
+    Verifies that score_data_to_prompts correctly handles image content
+    and produces valid MaxSim scores.
+    """
+    red_image = _make_base64_image(64, 64, color=(255, 0, 0))
+    blue_image = _make_base64_image(64, 64, color=(0, 0, 255))
+
+    query = "Describe the red object"
+    image_docs = [
+        _make_image_mm_param(red_image),
+        _make_image_mm_param(blue_image),
+    ]
+
+    with vllm_runner(
+        model,
+        runner="pooling",
+        dtype=dtype,
+        max_model_len=4096,
+        enforce_eager=True,
+        gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
+    ) as vllm_model:
+        scores = vllm_model.llm.score(query, image_docs)
+
+        assert len(scores) == 2
+        for s in scores:
+            assert isinstance(s.outputs.score, float)
+
+
+def _run_multimodal_mixed_docs_test(
+    vllm_runner: type[VllmRunner],
+    model: str,
+    *,
+    dtype: str,
+) -> None:
+    """Score a text query against a mix of text and image documents.
+
+    Ensures the late-interaction path handles heterogeneous document
+    types (plain strings alongside ScoreMultiModalParam images) in
+    a single call.
+    """
+    red_image = _make_base64_image(64, 64, color=(255, 0, 0))
+
+    query = "What is the capital of France?"
+    documents: list = [
+        "The capital of France is Paris.",
+        _make_image_mm_param(red_image),
+    ]
+
+    with vllm_runner(
+        model,
+        runner="pooling",
+        dtype=dtype,
+        max_model_len=4096,
+        enforce_eager=True,
+        gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
+    ) as vllm_model:
+        scores = vllm_model.llm.score(query, documents)
+
+        assert len(scores) == 2
+        for s in scores:
+            assert isinstance(s.outputs.score, float)
+        # Text document about France should score higher than a random image
+        assert scores[0].outputs.score > scores[1].outputs.score
+
+
+def _run_multimodal_image_query_text_docs_test(
+    vllm_runner: type[VllmRunner],
+    model: str,
+    *,
+    dtype: str,
+) -> None:
+    """Score an image query against text documents.
+
+    Verifies the reverse direction: multimodal query with text-only
+    documents through the late-interaction scoring path.
+    """
+    red_image = _make_base64_image(64, 64, color=(255, 0, 0))
+    image_query = _make_image_mm_param(red_image, text="red color")
+
+    documents = [
+        "A bright red sports car.",
+        "The weather forecast shows rain tomorrow.",
+    ]
+
+    with vllm_runner(
+        model,
+        runner="pooling",
+        dtype=dtype,
+        max_model_len=4096,
+        enforce_eager=True,
+        gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
+    ) as vllm_model:
+        scores = vllm_model.llm.score(image_query, documents)
+
+        assert len(scores) == 2
+        for s in scores:
+            assert isinstance(s.outputs.score, float)
+
+
+@pytest.mark.parametrize("model", MODELS)
+@pytest.mark.parametrize("dtype", [DTYPE])
+def test_colqwen3_multimodal_text_query_image_docs(
+    vllm_runner,
+    model: str,
+    dtype: str,
+) -> None:
+    _run_multimodal_text_query_image_docs_test(vllm_runner, model, dtype=dtype)
+
+
+@pytest.mark.parametrize("model", MODELS)
+@pytest.mark.parametrize("dtype", [DTYPE])
+def test_colqwen3_multimodal_mixed_docs(
+    vllm_runner,
+    model: str,
+    dtype: str,
+) -> None:
+    _run_multimodal_mixed_docs_test(vllm_runner, model, dtype=dtype)
+
+
+@pytest.mark.parametrize("model", MODELS)
+@pytest.mark.parametrize("dtype", [DTYPE])
+def test_colqwen3_multimodal_image_query_text_docs(
+    vllm_runner,
+    model: str,
+    dtype: str,
+) -> None:
+    _run_multimodal_image_query_text_docs_test(vllm_runner, model, dtype=dtype)

tests/models/registry.py

@@ -592,6 +592,9 @@ _EMBEDDING_EXAMPLE_MODELS = {
     ),
     # [Multimodal]
     "CLIPModel": _HfExamplesInfo("openai/clip-vit-base-patch32"),
+    "ColModernVBertForRetrieval": _HfExamplesInfo(
+        "ModernVBERT/colmodernvbert-merged",
+    ),
     "LlavaNextForConditionalGeneration": _HfExamplesInfo("royokong/e5-v"),
     "Phi3VForCausalLM": _HfExamplesInfo(
         "TIGER-Lab/VLM2Vec-Full", trust_remote_code=True
@@ -603,6 +606,9 @@ _EMBEDDING_EXAMPLE_MODELS = {
     "OpsColQwen3Model": _HfExamplesInfo(
         "OpenSearch-AI/Ops-Colqwen3-4B", trust_remote_code=True
     ),
+    "Qwen3VLNemotronEmbedModel": _HfExamplesInfo(
+        "nvidia/nemotron-colembed-vl-4b-v2",
+    ),
     "SiglipModel": _HfExamplesInfo("google/siglip-base-patch16-224"),
     "PrithviGeoSpatialMAE": _HfExamplesInfo(
         "ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11",
@@ -1023,6 +1029,14 @@ _MULTIMODAL_EXAMPLE_MODELS = {
         min_transformers_version="4.57",
         is_available_online=False,
     ),
+    "Qwen3ASRRealtimeGeneration": _HfExamplesInfo(
+        "Qwen/Qwen3-ASR-1.7B",
+        max_model_len=4096,
+        min_transformers_version="4.57",
+        enforce_eager=True,
+        hf_overrides={"architectures": ["Qwen3ASRRealtimeGeneration"]},
+        is_available_online=False,
+    ),
     "RForConditionalGeneration": _HfExamplesInfo("YannQi/R-4B", trust_remote_code=True),
     "SkyworkR1VChatModel": _HfExamplesInfo(
         "Skywork/Skywork-R1V-38B", trust_remote_code=True

tests/models/test_transformers.py

@@ -6,8 +6,6 @@ from typing import Any
 
 import pytest
 
-from vllm.platforms import current_platform
-
 from ..conftest import HfRunner, VllmRunner
 from ..utils import multi_gpu_test, prep_prompts
 from .registry import HF_EXAMPLE_MODELS
@@ -131,6 +129,7 @@ def test_distributed(
                 "quantization": "bitsandbytes",
             },
         ),
+        ("unsloth/tinyllama-bnb-4bit", {}),
     ],
 )
 @pytest.mark.parametrize("max_tokens", [32])
@@ -143,12 +142,6 @@ def test_quantization(
     max_tokens: int,
     num_logprobs: int,
 ) -> None:
-    if (
-        current_platform.is_rocm()
-        and quantization_kwargs.get("quantization", "") == "bitsandbytes"
-    ):
-        pytest.skip("bitsandbytes quantization is currently not supported in rocm.")
-
     with vllm_runner(
         model,
         model_impl="auto",

tests/quantization/test_compressed_tensors.py

@@ -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

tests/reasoning/test_qwen3_reasoning_parser.py

@@ -4,46 +4,79 @@
 import pytest
 from transformers import AutoTokenizer
 
-from tests.reasoning.utils import run_reasoning_extraction
+from tests.reasoning.utils import (
+    StreamingReasoningReconstructor,
+    run_reasoning_extraction,
+    run_reasoning_extraction_streaming,
+)
 from vllm.reasoning import ReasoningParser, ReasoningParserManager
 
 parser_name = "qwen3"
 start_token = "<think>"
 end_token = "</think>"
 
-REASONING_MODEL_NAME = "Qwen/Qwen3-0.6B"
+REASONING_MODEL_NAMES = [
+    "Qwen/Qwen3-0.6B",
+    "Qwen/Qwen3.5-397B-A17B",
+    "Qwen/Qwen3-4B-Thinking-2507",
+]
 
 
-@pytest.fixture(scope="module")
-def qwen3_tokenizer():
-    return AutoTokenizer.from_pretrained(REASONING_MODEL_NAME)
+@pytest.fixture(scope="module", params=REASONING_MODEL_NAMES)
+def qwen3_tokenizer(request):
+    return AutoTokenizer.from_pretrained(request.param)
 
 
-# 带 <think></think>，非stream
+# --- <think> in prompt, only </think> in output (typical) ---
+
+WITHOUT_START_TOKEN = {
+    "output": "This is a reasoning section</think>This is the rest",
+    "reasoning": "This is a reasoning section",
+    "content": "This is the rest",
+}
+WITHOUT_START_TOKEN_STREAM = {
+    "output": "This is a reasoning section</think>This is the rest",
+    "reasoning": "This is a reasoning section",
+    "content": "This is the rest",
+}
+WITHOUT_START_TOKEN_COMPLETE_REASONING = {
+    "output": "This is a reasoning section</think>",
+    "reasoning": "This is a reasoning section",
+    "content": None,
+}
+
+# --- <think> present in output (old template / edge case) ---
+
 WITH_THINK = {
     "output": "<think>This is a reasoning section</think>This is the rest",
     "reasoning": "This is a reasoning section",
     "content": "This is the rest",
 }
-# 带 <think></think>，stream
 WITH_THINK_STREAM = {
     "output": "<think>This is a reasoning section</think>This is the rest",
     "reasoning": "This is a reasoning section",
     "content": "This is the rest",
 }
-# 不带 <think></think>，非stream
+
+# --- No think tokens at all (thinking disabled) ---
+
 WITHOUT_THINK = {
     "output": "This is the rest",
     "reasoning": None,
     "content": "This is the rest",
 }
-# 不带 <think></think>，stream
+# In streaming, the parser cannot distinguish "thinking disabled" from
+# "reasoning in progress" when no think tokens have appeared yet.
+# It assumes reasoning. The serving layer handles the "thinking disabled"
+# case by checking prompt_is_reasoning_end_arr before calling the parser.
 WITHOUT_THINK_STREAM = {
     "output": "This is the rest",
-    "reasoning": None,
-    "content": "This is the rest",
+    "reasoning": "This is the rest",
+    "content": None,
 }
 
+# --- Edge cases ---
+
 COMPLETE_REASONING = {
     "output": "<think>This is a reasoning section</think>",
     "reasoning": "This is a reasoning section",
@@ -57,7 +90,7 @@ MULTILINE_REASONING = {
 ONLY_OPEN_TAG = {
     "output": "<think>This is a reasoning section",
     "reasoning": None,
-    "content": "<think>This is a reasoning section",
+    "content": "This is a reasoning section",
 }
 
 ONLY_OPEN_TAG_STREAM = {
@@ -67,6 +100,26 @@ ONLY_OPEN_TAG_STREAM = {
 }
 
 TEST_CASES = [
+    pytest.param(
+        False,
+        WITHOUT_START_TOKEN,
+        id="without_start_token",
+    ),
+    pytest.param(
+        True,
+        WITHOUT_START_TOKEN_STREAM,
+        id="without_start_token_stream",
+    ),
+    pytest.param(
+        False,
+        WITHOUT_START_TOKEN_COMPLETE_REASONING,
+        id="without_start_token_complete_reasoning",
+    ),
+    pytest.param(
+        True,
+        WITHOUT_START_TOKEN_COMPLETE_REASONING,
+        id="without_start_token_complete_reasoning_stream",
+    ),
     pytest.param(
         False,
         WITH_THINK,
@@ -140,3 +193,59 @@ def test_reasoning(
 
     assert reasoning == param_dict["reasoning"]
     assert content == param_dict["content"]
+
+
+# Multi-token delta tests: simulate real-world streaming where a single
+# delta can contain multiple tokens (e.g., speculative decoding).
+MULTI_TOKEN_DELTA_CASES = [
+    pytest.param(
+        # <think> grouped with following text in one delta
+        ["<think>This is a reasoning section", "</think>", "This is the rest"],
+        "This is a reasoning section",
+        "This is the rest",
+        id="start_token_grouped_with_text",
+    ),
+    pytest.param(
+        # </think> grouped with following content in one delta
+        ["reasoning section", "</think>This is the rest"],
+        "reasoning section",
+        "This is the rest",
+        id="end_token_grouped_with_content",
+    ),
+    pytest.param(
+        # <think> and </think> in the same delta, no content after
+        ["<think>reasoning</think>"],
+        "reasoning",
+        None,
+        id="start_and_end_in_one_delta_no_content",
+    ),
+    pytest.param(
+        # No start token, end grouped with content (Qwen3.5 style)
+        ["reasoning section", "</think>content"],
+        "reasoning section",
+        "content",
+        id="no_start_end_grouped_with_content",
+    ),
+]
+
+
+@pytest.mark.parametrize(
+    "deltas, expected_reasoning, expected_content", MULTI_TOKEN_DELTA_CASES
+)
+def test_reasoning_streaming_multi_token_deltas(
+    deltas: list[str],
+    expected_reasoning: str | None,
+    expected_content: str | None,
+    qwen3_tokenizer,
+):
+    """Test that multi-token deltas don't leak <think> into reasoning."""
+    parser: ReasoningParser = ReasoningParserManager.get_reasoning_parser(parser_name)(
+        qwen3_tokenizer
+    )
+
+    reconstructor: StreamingReasoningReconstructor = run_reasoning_extraction_streaming(
+        parser, deltas
+    )
+
+    assert reconstructor.reasoning == expected_reasoning
+    assert (reconstructor.other_content or None) == expected_content

tests/utils.py

@@ -128,6 +128,9 @@ class RemoteOpenAIServer:
             env=env,
             stdout=sys.stdout,
             stderr=sys.stderr,
+            # Create a dedicated process group so we can kill
+            # the entire tree (parent + EngineCore + workers) at once.
+            start_new_session=True,
         )
 
     def __init__(
@@ -189,8 +192,17 @@ class RemoteOpenAIServer:
             model_loader = get_model_loader(load_config)
             model_loader.download_model(model_config)
 
+        # Record GPU memory before server start so we know what
+        # "released" looks like.
+        self._pre_server_gpu_memory = self._get_gpu_memory_used()
+        if self._pre_server_gpu_memory is not None:
+            pre_gb = self._pre_server_gpu_memory / 1e9
+            print(
+                f"[RemoteOpenAIServer] GPU memory before server start: {pre_gb:.2f} GB"
+            )
+
         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):
@@ -198,27 +210,69 @@ class RemoteOpenAIServer:
 
     def __exit__(self, exc_type, exc_value, traceback):
         pid = self.proc.pid
-        # Graceful shutdown
-        self.proc.terminate()
+
+        # Get the process group ID. Because we used
+        # start_new_session=True the pgid equals the server's pid.
+        try:
+            pgid = os.getpgid(pid)
+        except (ProcessLookupError, OSError):
+            pgid = None
+
+        # Phase 1: graceful SIGTERM to the entire process group
+        if pgid is not None:
+            with contextlib.suppress(ProcessLookupError, OSError):
+                os.killpg(pgid, signal.SIGTERM)
+                print(f"[RemoteOpenAIServer] Sent SIGTERM to process group {pgid}")
+        else:
+            self.proc.terminate()
+
         try:
             self.proc.wait(timeout=15)
             print(f"[RemoteOpenAIServer] Server {pid} terminated gracefully")
         except subprocess.TimeoutExpired:
+            # Phase 2: SIGKILL the entire process group
             print(
                 f"[RemoteOpenAIServer] Server {pid} did not respond "
-                "to SIGTERM, sending SIGKILL"
+                "to SIGTERM, sending SIGKILL to process group"
             )
-            self.proc.kill()
+            if pgid is not None:
+                with contextlib.suppress(ProcessLookupError, OSError):
+                    os.killpg(pgid, signal.SIGKILL)
+            else:
+                self.proc.kill()
+
             try:
-                self.proc.wait(timeout=5)
+                self.proc.wait(timeout=10)
                 print(f"[RemoteOpenAIServer] Server {pid} killed")
-            except subprocess.TimeoutExpired as err:
-                raise RuntimeError(
-                    f"[RemoteOpenAIServer] Failed to kill server process {pid}"
-                ) from err
-        # Wait for GPU memory to be released
+            except subprocess.TimeoutExpired:
+                # Phase 3: last resort - find and kill any orphaned children
+                self._kill_orphaned_children(pid)
+
+        # Wait for GPU memory to actually be *freed*, not just
+        # "stabilized at whatever level it's at".
         self._wait_for_gpu_memory_release()
 
+    def _kill_orphaned_children(self, parent_pid: int) -> None:
+        """Best-effort cleanup of any lingering child processes."""
+        try:
+            import psutil
+
+            parent = psutil.Process(parent_pid)
+            children = parent.children(recursive=True)
+            for child in children:
+                print(
+                    f"[RemoteOpenAIServer] Killing orphaned child "
+                    f"pid={child.pid} name={child.name()}"
+                )
+                child.kill()
+            psutil.wait_procs(children, timeout=5)
+        except Exception as e:
+            # psutil may not be installed, or processes already gone
+            print(f"[RemoteOpenAIServer] Orphan cleanup failed: {e}")
+            # Fallback: try to kill by pgid one more time
+            with contextlib.suppress(ProcessLookupError, OSError):
+                os.killpg(parent_pid, signal.SIGKILL)
+
     def _get_gpu_memory_used(self) -> float | None:
         """Get total GPU memory used across all visible devices in bytes."""
         try:
@@ -244,10 +298,26 @@ class RemoteOpenAIServer:
             return None
         return None
 
-    def _wait_for_gpu_memory_release(self, timeout: float = 30.0):
-        """Poll GPU memory until it stabilizes, indicating cleanup is complete."""
+    def _wait_for_gpu_memory_release(self, timeout: float = 60.0):
+        """Wait for GPU memory to drop back toward pre-server levels.
+
+        Two-phase strategy:
+          1. Try to wait for memory to return close to pre-server baseline.
+          2. If that doesn't happen, fall back to waiting for stabilization
+             and log a warning (the next server might still OOM).
+        """
+        baseline = self._pre_server_gpu_memory
+        if baseline is None:
+            # Can't query GPU memory - nothing to do
+            return
+
+        # Allow up to 2 GiB overhead above baseline for driver/context state
+        # that may persist between server instances.
+        headroom_bytes = 2 * 1024 * 1024 * 1024
+        target = baseline + headroom_bytes
+
         start = time.time()
-        prev_used: float | None = None
+        last_used: float | None = None
         stable_count = 0
 
         while time.time() - start < timeout:
@@ -256,26 +326,49 @@ class RemoteOpenAIServer:
             if used is None:
                 return  # Can't query, assume ok
 
-            if prev_used is not None and abs(used - prev_used) < 100 * 1024 * 1024:
-                stable_count += 1
-                if stable_count >= 3:
-                    used_gb = used / 1e9
-                    print(
-                        f"[RemoteOpenAIServer] GPU memory stabilized "
-                        f"at {used_gb:.2f} GB"
-                    )
-                    return
-            else:
-                stable_count = 0
+            used_gb = used / 1e9
+            target_gb = target / 1e9
+            elapsed = time.time() - start
 
-            prev_used = used
-            time.sleep(0.1)
+            # Phase 1: memory dropped to near baseline - we're done.
+            if used <= target:
+                print(
+                    f"[RemoteOpenAIServer] GPU memory released to "
+                    f"{used_gb:.2f} GB (target: {target_gb:.2f} GB) "
+                    f"in {elapsed:.1f}s"
+                )
+                return
 
-        last_reading = prev_used / 1e9 if prev_used is not None else 0.0
+            # Phase 2 (after 40s): fall back to stabilization check.
+            # This handles cases where another process is using GPU memory
+            # and we'll never reach baseline.
+            if elapsed > 40.0 and last_used is not None:
+                delta = abs(used - last_used)
+                if delta < 200 * 1024 * 1024:  # 200 MB
+                    stable_count += 1
+                    if stable_count >= 3:
+                        print(
+                            f"[RemoteOpenAIServer] WARNING: GPU memory "
+                            f"stabilized at {used_gb:.2f} GB "
+                            f"(target was {target_gb:.2f} GB). "
+                            f"Proceeding - next server may OOM."
+                        )
+                        return
+                else:
+                    stable_count = 0
+
+            last_used = used
+            time.sleep(1.0)
+
+        # Timeout - log clearly so CI failures are diagnosable
+        final_used = self._get_gpu_memory_used()
+        final_gb = final_used / 1e9 if final_used else 0.0
         raise RuntimeError(
-            f"[RemoteOpenAIServer] GPU memory did not stabilize within {timeout}s. "
-            f"Last reading: {last_reading:.2f} GB. "
-            "Child processes may still be holding GPU memory."
+            f"[RemoteOpenAIServer] GPU memory did not release within "
+            f"{timeout}s. Current: {final_gb:.2f} GB, "
+            f"target: {target / 1e9:.2f} GB, "
+            f"baseline: {baseline / 1e9:.2f} GB. "
+            f"Child processes may still be holding GPU memory."
         )
 
     def _poll(self) -> int | None:

tests/v1/attention/test_mla_backends.py

@@ -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

tests/v1/core/test_kv_cache_utils.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import hashlib
 import importlib
 from collections.abc import Callable
 from typing import Any
@@ -498,14 +499,41 @@ def test_generate_block_hash_extra_keys_prompt_embeds():
     # Test with prompt embeds for the first block
     extra_keys, _ = generate_block_hash_extra_keys(request, 0, 5, 0)
     expected_embeds = prompt_embeds[0:5]
-    expected_bytes = kv_cache_utils.tensor_data(expected_embeds).tobytes()
-    assert extra_keys == (expected_bytes,)
+    expected_hash = hashlib.sha256(kv_cache_utils.tensor_data(expected_embeds)).digest()
+    assert extra_keys == (expected_hash,)
 
     # Test with prompt embeds for the second block
     extra_keys, _ = generate_block_hash_extra_keys(request, 5, 10, 0)
     expected_embeds = prompt_embeds[5:10]
-    expected_bytes = kv_cache_utils.tensor_data(expected_embeds).tobytes()
-    assert extra_keys == (expected_bytes,)
+    expected_hash = hashlib.sha256(kv_cache_utils.tensor_data(expected_embeds)).digest()
+    assert extra_keys == (expected_hash,)
+
+
+def test_generate_block_hash_extra_keys_prompt_embeds_cached(monkeypatch):
+    prompt_embeds = torch.randn(10, 3)
+    request = make_request(
+        request_id="0",
+        prompt_token_ids=None,
+        mm_positions=None,
+        mm_hashes=None,
+        prompt_embeds=prompt_embeds,
+        block_size=20,
+    )
+
+    num_tensor_data_calls = 0
+    original_tensor_data = kv_cache_utils.tensor_data
+
+    def counting_tensor_data(tensor: torch.Tensor):
+        nonlocal num_tensor_data_calls
+        num_tensor_data_calls += 1
+        return original_tensor_data(tensor)
+
+    monkeypatch.setattr(kv_cache_utils, "tensor_data", counting_tensor_data)
+
+    extra_keys_1, _ = generate_block_hash_extra_keys(request, 0, 5, 0)
+    extra_keys_2, _ = generate_block_hash_extra_keys(request, 0, 5, 0)
+    assert extra_keys_1 == extra_keys_2
+    assert num_tensor_data_calls == 1
 
 
 def test_generate_block_hash_extra_keys_different_prompt_embeds():
@@ -1858,22 +1886,26 @@ def test_request_block_hasher_with_prompt_embeds(hash_fn: Callable[[Any], bytes]
     block_hashes = request.block_hashes
     assert len(block_hashes) == 2
 
-    block1_embeds_bytes = tensor_data(prompt_embeds[:block_size]).tobytes()
+    block1_embeds_hash = hashlib.sha256(
+        tensor_data(prompt_embeds[:block_size])
+    ).digest()
     expected_hash1 = hash_fn(
         (
             kv_cache_utils.NONE_HASH,
             tuple(prompt_token_ids[:block_size]),
-            (block1_embeds_bytes,),
+            (block1_embeds_hash,),
         )
     )
     assert block_hashes[0] == expected_hash1
 
-    block2_embeds_bytes = tensor_data(prompt_embeds[block_size:num_tokens]).tobytes()
+    block2_embeds_hash = hashlib.sha256(
+        tensor_data(prompt_embeds[block_size:num_tokens])
+    ).digest()
     expected_hash2 = hash_fn(
         (
             block_hashes[0],
             tuple(prompt_token_ids[block_size:num_tokens]),
-            (block2_embeds_bytes,),
+            (block2_embeds_hash,),
         )
     )
     assert block_hashes[1] == expected_hash2
@@ -1903,22 +1935,26 @@ def test_request_with_prompt_embeds_and_mm_inputs(hash_fn: Callable[[Any], bytes
     block_hashes = request.block_hashes
     assert len(block_hashes) == 2
 
-    block1_embeds_bytes = tensor_data(prompt_embeds[:block_size]).tobytes()
+    block1_embeds_hash = hashlib.sha256(
+        tensor_data(prompt_embeds[:block_size])
+    ).digest()
     expected_hash1 = hash_fn(
         (
             kv_cache_utils.NONE_HASH,
             tuple(prompt_token_ids[:block_size]),
-            ("hash1", block1_embeds_bytes),
+            ("hash1", block1_embeds_hash),
         )
     )
     assert block_hashes[0] == expected_hash1
 
-    block2_embeds_bytes = tensor_data(prompt_embeds[block_size:num_tokens]).tobytes()
+    block2_embeds_hash = hashlib.sha256(
+        tensor_data(prompt_embeds[block_size:num_tokens])
+    ).digest()
     expected_hash2 = hash_fn(
         (
             block_hashes[0],
             tuple(prompt_token_ids[block_size:num_tokens]),
-            ("hash2", block2_embeds_bytes),
+            ("hash2", block2_embeds_hash),
         )
     )
     assert block_hashes[1] == expected_hash2

tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py

@@ -30,7 +30,7 @@ def _make_get_num_new_matched_tokens(
 
 @pytest.fixture
 def scheduler():
-    vllm_config = create_vllm_config()
+    vllm_config = create_vllm_config(kv_load_failure_policy="recompute")
     return create_scheduler(vllm_config)
 
 

tests/v1/kv_connector/unit/test_moriio_connector.py

@@ -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)

tests/v1/kv_connector/unit/utils.py

@@ -5,7 +5,7 @@ from collections import defaultdict
 from collections.abc import Callable
 from dataclasses import dataclass
 from itertools import chain, count
-from typing import Any
+from typing import Any, Literal
 
 import torch
 
@@ -96,6 +96,7 @@ def create_vllm_config(
     cache_dtype: str = "auto",
     hf_overrides: dict[str, Any] | None = None,
     attention_backend: str | None = None,
+    kv_load_failure_policy: Literal["recompute", "fail"] = "fail",
 ) -> VllmConfig:
     """Initialize VllmConfig For Testing."""
     model_config = ModelConfig(
@@ -125,6 +126,7 @@ def create_vllm_config(
         kv_role="kv_both",
         enable_permute_local_kv=enable_permute_local_kv,
         kv_connector_extra_config=kv_connector_extra_config or {},
+        kv_load_failure_policy=kv_load_failure_policy,
     )
     attention_config = AttentionConfig(backend=attention_backend)
     return VllmConfig(

tests/v1/sample/test_logprobs.py

@@ -52,7 +52,7 @@ def vllm_model(vllm_runner, request) -> Generator[VllmRunner, None, None]:
         # TODO: enable this once we support it for
         # prompt logprobs.
         enable_prefix_caching=request.param,
-        gpu_memory_utilization=0.4,  # up to 2 alive concurrently
+        gpu_memory_utilization=0.4,
     ) as vllm_model:
         yield vllm_model
 
@@ -366,21 +366,20 @@ def test_max_logprobs():
     Should also fail for `prompt_logprobs > max_logprobs`
     APC should not matter as this test checks basic request validation.
     """
-    runner = VllmRunner(
+    with VllmRunner(
         "facebook/opt-125m",
         max_logprobs=1,
         enable_prefix_caching=False,
-        # 2 other llms alive during whole session
         gpu_memory_utilization=0.15,
         max_model_len=256,
-    )
-    vllm_sampling_params = SamplingParams(logprobs=1)
-    # should pass
-    runner.generate(["Hello world"], sampling_params=vllm_sampling_params)
+    ) as runner:
+        vllm_sampling_params = SamplingParams(logprobs=1)
+        # should pass
+        runner.generate(["Hello world"], sampling_params=vllm_sampling_params)
 
-    bad_sampling_params = SamplingParams(logprobs=2)
-    with pytest.raises(ValueError):
-        runner.generate(["Hello world"], sampling_params=bad_sampling_params)
+        bad_sampling_params = SamplingParams(logprobs=2)
+        with pytest.raises(ValueError):
+            runner.generate(["Hello world"], sampling_params=bad_sampling_params)
 
 
 def test_none_logprobs(vllm_model, example_prompts):
@@ -449,33 +448,31 @@ def test_all_logprobs(example_prompts):
     Args:
       example_prompts: list of example prompts (test fixture)
     """
-    runner = VllmRunner(
+    with VllmRunner(
         "facebook/opt-125m",
         max_logprobs=-1,
         enable_prefix_caching=False,
-        # 2 other llms alive during whole session
         gpu_memory_utilization=0.15,
         max_model_len=256,
-    )
+    ) as runner:
+        sampling_params_logprobs_all = SamplingParams(
+            max_tokens=5, logprobs=-1, prompt_logprobs=-1
+        )
+        results_logprobs_all = runner.llm.generate(
+            example_prompts, sampling_params=sampling_params_logprobs_all
+        )
+        vocab_size = runner.llm.llm_engine.model_config.get_vocab_size()
 
-    sampling_params_logprobs_all = SamplingParams(
-        max_tokens=5, logprobs=-1, prompt_logprobs=-1
-    )
-    results_logprobs_all = runner.llm.generate(
-        example_prompts, sampling_params=sampling_params_logprobs_all
-    )
-    vocab_size = runner.llm.llm_engine.model_config.get_vocab_size()
-
-    for i in range(len(results_logprobs_all)):
-        logprobs = results_logprobs_all[i].outputs[0].logprobs
-        prompt_logprobs = results_logprobs_all[i].prompt_logprobs
-        assert logprobs is not None
-        for logprob in logprobs:
-            assert len(logprob) == vocab_size
-        assert prompt_logprobs is not None
-        assert prompt_logprobs[0] is None
-        for prompt_logprob in prompt_logprobs[1:]:
-            assert len(prompt_logprob) == vocab_size
+        for i in range(len(results_logprobs_all)):
+            logprobs = results_logprobs_all[i].outputs[0].logprobs
+            prompt_logprobs = results_logprobs_all[i].prompt_logprobs
+            assert logprobs is not None
+            for logprob in logprobs:
+                assert len(logprob) == vocab_size
+            assert prompt_logprobs is not None
+            assert prompt_logprobs[0] is None
+            for prompt_logprob in prompt_logprobs[1:]:
+                assert len(prompt_logprob) == vocab_size
 
 
 @pytest.mark.parametrize("logprobs_mode", get_args(LogprobsMode))
@@ -495,24 +492,28 @@ def test_logprobs_mode(logprobs_mode: LogprobsMode):
         max_model_len=16,
         logprobs_mode=logprobs_mode,
     )
-    vllm_sampling_params = SamplingParams(logprobs=1)
-    results = llm.generate(["Hello world"], sampling_params=vllm_sampling_params)
+    try:
+        vllm_sampling_params = SamplingParams(logprobs=1)
+        results = llm.generate(["Hello world"], sampling_params=vllm_sampling_params)
 
-    total_token_with_logprobs = 0
-    positive_values = 0
-    for output in results[0].outputs:
-        for logprobs in output.logprobs:
-            for token_id in logprobs:
-                logprob = logprobs[token_id]
-                if logprobs_mode in ("raw_logprobs", "processed_logprobs"):
-                    assert logprob.logprob <= 0
-                if logprob.logprob > 0:
-                    positive_values = positive_values + 1
-                total_token_with_logprobs = total_token_with_logprobs + 1
-    assert total_token_with_logprobs >= len(results[0].outputs)
-    if logprobs_mode in ("raw_logits", "processed_logits"):
-        assert positive_values > 0
-    del llm
+        total_token_with_logprobs = 0
+        positive_values = 0
+        for output in results[0].outputs:
+            for logprobs in output.logprobs:
+                for token_id in logprobs:
+                    logprob = logprobs[token_id]
+                    if logprobs_mode in ("raw_logprobs", "processed_logprobs"):
+                        assert logprob.logprob <= 0
+                    if logprob.logprob > 0:
+                        positive_values = positive_values + 1
+                    total_token_with_logprobs = total_token_with_logprobs + 1
+        assert total_token_with_logprobs >= len(results[0].outputs)
+        if logprobs_mode in ("raw_logits", "processed_logits"):
+            assert positive_values > 0
+    finally:
+        del llm
+        torch.cuda.empty_cache()
+        cleanup_dist_env_and_memory()
 
 
 class TestCorrectDecodedToken:
@@ -767,7 +768,7 @@ class TestCorrectDecodedToken:
             # Simulate cases where individual tokens decode to "�"
             # but combinations decode correctly
             if len(ids) == 1:
-                if ids[0] == 3 or ids[0] == 4 or ids[0] == 8 or ids[0] == 9:
+                if ids[0] in (3, 4, 8, 9):
                     return "�"
             elif len(ids) == 2:
                 if ids == [2, 3]:
@@ -809,42 +810,41 @@ def test_verify_tokens_integration():
     corrects tokens ending with the replacement character "�".
     Uses facebook/opt-125m which is known to produce these issues.
     """
-    runner = VllmRunner(
+    with VllmRunner(
         "facebook/opt-125m",
         max_logprobs=0,
         enable_prefix_caching=False,
         gpu_memory_utilization=0.15,
         max_model_len=256,
-    )
+    ) as runner:
+        # Use a prompt that triggers multi-byte UTF-8 issues
+        # Based on user's example: "In this example,"
+        test_prompts = ["In this example,"]
 
-    # Use a prompt that triggers multi-byte UTF-8 issues
-    # Based on user's example: "In this example,"
-    test_prompts = ["In this example,"]
+        sampling_params = SamplingParams(
+            max_tokens=16,
+            temperature=0,
+            logprobs=0,
+        )
 
-    sampling_params = SamplingParams(
-        max_tokens=16,
-        temperature=0,
-        logprobs=0,
-    )
+        results = runner.llm.generate(test_prompts, sampling_params=sampling_params)
 
-    results = runner.llm.generate(test_prompts, sampling_params=sampling_params)
-
-    # Verify that decoded tokens don't contain replacement characters
-    for result in results:
-        assert result.outputs[0].logprobs is not None
-        for logprob_dict in result.outputs[0].logprobs:
-            for token_id, logprob_info in logprob_dict.items():
-                decoded_token = logprob_info.decoded_token
-                # Decoded tokens should not end with replacement character
-                # They should either be corrected or empty string
-                assert not decoded_token.endswith("�"), (
-                    f"Token {token_id} decoded to '{decoded_token}' which "
-                    f"ends with replacement character"
-                )
-                # Decoded tokens should not contain lone replacement characters
-                assert decoded_token != "�", (
-                    f"Token {token_id} is a lone replacement character"
-                )
+        # Verify that decoded tokens don't contain replacement characters
+        for result in results:
+            assert result.outputs[0].logprobs is not None
+            for logprob_dict in result.outputs[0].logprobs:
+                for token_id, logprob_info in logprob_dict.items():
+                    decoded_token = logprob_info.decoded_token
+                    # Decoded tokens should not end with replacement character
+                    # They should either be corrected or empty string
+                    assert not decoded_token.endswith("�"), (
+                        f"Token {token_id} decoded to '{decoded_token}' which "
+                        f"ends with replacement character"
+                    )
+                    # Decoded tokens should not contain lone replacement characters
+                    assert decoded_token != "�", (
+                        f"Token {token_id} is a lone replacement character"
+                    )
 
 
 def test_utf8_edge_cases_with_real_model():
@@ -853,45 +853,44 @@ def test_utf8_edge_cases_with_real_model():
     Tests prompts that are likely to trigger byte-fallback tokenization
     and multi-byte UTF-8 splitting.
     """
-    runner = VllmRunner(
+    with VllmRunner(
         "facebook/opt-125m",
         max_logprobs=1,
         enable_prefix_caching=False,
         gpu_memory_utilization=0.15,
         max_model_len=256,
-    )
+    ) as runner:
+        # Prompts with various multi-byte UTF-8 characters
+        test_prompts = [
+            'Smart quotes: "Hello"',  # Curly quotes
+            "Em dash — test",  # Em dash
+            "Ellipsis… continues",  # Ellipsis
+            "Chinese: 你好",  # Chinese characters
+            "Emoji: 😀 🎉",  # Emojis
+            'Mixed: "quoted" — with symbols',  # Mixed
+        ]
 
-    # Prompts with various multi-byte UTF-8 characters
-    test_prompts = [
-        'Smart quotes: "Hello"',  # Curly quotes
-        "Em dash — test",  # Em dash
-        "Ellipsis… continues",  # Ellipsis
-        "Chinese: 你好",  # Chinese characters
-        "Emoji: 😀 🎉",  # Emojis
-        'Mixed: "quoted" — with symbols',  # Mixed
-    ]
+        sampling_params = SamplingParams(
+            max_tokens=10,
+            temperature=0,
+            logprobs=1,
+        )
 
-    sampling_params = SamplingParams(
-        max_tokens=10,
-        temperature=0,
-        logprobs=1,
-    )
+        results = runner.llm.generate(test_prompts, sampling_params=sampling_params)
 
-    results = runner.llm.generate(test_prompts, sampling_params=sampling_params)
+        for i, result in enumerate(results):
+            prompt = test_prompts[i]
+            assert result.outputs[0].logprobs is not None
 
-    for i, result in enumerate(results):
-        prompt = test_prompts[i]
-        assert result.outputs[0].logprobs is not None
-
-        # Check that no decoded tokens end with replacement character
-        for logprob_dict in result.outputs[0].logprobs:
-            for token_id, logprob_info in logprob_dict.items():
-                decoded_token = logprob_info.decoded_token
-                assert not decoded_token.endswith("�"), (
-                    f"Prompt: '{prompt}'\n"
-                    f"Token {token_id} decoded to '{decoded_token}' which "
-                    f"ends with replacement character"
-                )
+            # Check that no decoded tokens end with replacement character
+            for logprob_dict in result.outputs[0].logprobs:
+                for token_id, logprob_info in logprob_dict.items():
+                    decoded_token = logprob_info.decoded_token
+                    assert not decoded_token.endswith("�"), (
+                        f"Prompt: '{prompt}'\n"
+                        f"Token {token_id} decoded to '{decoded_token}' which "
+                        f"ends with replacement character"
+                    )
 
 
 def test_correct_decoded_token_preserves_valid_tokens():
@@ -901,36 +900,35 @@ def test_correct_decoded_token_preserves_valid_tokens():
     ending with "�", but this test verifies the broader _verify_tokens
     logic doesn't affect valid tokens.
     """
-    runner = VllmRunner(
+    with VllmRunner(
         "facebook/opt-125m",
         max_logprobs=2,
         enable_prefix_caching=False,
         gpu_memory_utilization=0.15,
         max_model_len=256,
-    )
+    ) as runner:
+        # Simple prompt with standard ASCII characters
+        test_prompts = ["Hello world, this is a test."]
 
-    # Simple prompt with standard ASCII characters
-    test_prompts = ["Hello world, this is a test."]
+        sampling_params = SamplingParams(
+            max_tokens=10,
+            temperature=0,
+            logprobs=2,
+        )
 
-    sampling_params = SamplingParams(
-        max_tokens=10,
-        temperature=0,
-        logprobs=2,
-    )
+        results = runner.llm.generate(test_prompts, sampling_params=sampling_params)
 
-    results = runner.llm.generate(test_prompts, sampling_params=sampling_params)
+        for result in results:
+            assert result.outputs[0].logprobs is not None
 
-    for result in results:
-        assert result.outputs[0].logprobs is not None
-
-        # All decoded tokens should be valid strings
-        for logprob_dict in result.outputs[0].logprobs:
-            for token_id, logprob_info in logprob_dict.items():
-                decoded_token = logprob_info.decoded_token
-                # Valid tokens should be non-empty strings (or empty if corrected)
-                assert isinstance(decoded_token, str)
-                # Should not contain replacement character
-                assert "�" not in decoded_token
+            # All decoded tokens should be valid strings
+            for logprob_dict in result.outputs[0].logprobs:
+                for token_id, logprob_info in logprob_dict.items():
+                    decoded_token = logprob_info.decoded_token
+                    # Valid tokens should be non-empty strings (or empty if corrected)
+                    assert isinstance(decoded_token, str)
+                    # Should not contain replacement character
+                    assert "�" not in decoded_token
 
 
 @pytest.mark.parametrize("logprobs_mode", get_args(LogprobsMode))
@@ -985,16 +983,33 @@ def test_correct_decoded_token_preserves_valid_tokens():
 def test_spec_decode_logprobs(
     logprobs_mode: LogprobsMode,
     model_setup: tuple[str, str, dict, int],
+    monkeypatch,
 ):
     """Spec decode logprobs should match those of the base model.
 
+    Runs the base model and spec decode model sequentially, ensuring
+    only one LLM instance is alive at a time to avoid GPU memory
+    contention. Both use identical chunked prefill settings and eager
+    mode to control for infrastructure differences.
+
     Args:
         logprobs_mode: logprobs mode.
         model_setup: Tuple of (method, base model name,
             speculative_config dict, top_logprobs).
+        monkeypatch: pytest fixture for setting env vars.
     """
     from vllm import LLM
 
+    # The ROCm skinny GEMM kernels (gemm_kernels.cu) are
+    # non-deterministic across LLM instantiations due to persistent
+    # workgroup scheduling and wave-level shuffle reductions, which
+    # causes logprob differences that get misattributed to spec decode.
+    # Disable them so this test isolates spec decode correctness only.
+    # TODO(akaratza): Remove this workaround once the follow-up to
+    # https://github.com/vllm-project/vllm/pull/33493#issuecomment-3906083975
+    # lands with a determinism fix for wvSplitK kernels.
+    monkeypatch.setenv("VLLM_ROCM_USE_SKINNY_GEMM", "0")
+
     method, model_name, spec_config, top_logprobs = model_setup
 
     prompt = "Hello world " * 50
@@ -1068,8 +1083,17 @@ def test_spec_decode_logprobs(
     for ref_logprob, spec_logprob in zip(ref_logprobs, spec_logprobs):
         assert math.isclose(
             ref_logprob.logprob, spec_logprob.logprob, rel_tol=5e-2, abs_tol=1e-1
+        ), (
+            f"Logprob mismatch: ref={ref_logprob.logprob} "
+            f"spec={spec_logprob.logprob} "
+            f"diff={abs(ref_logprob.logprob - spec_logprob.logprob)} "
+            f"(token={ref_logprob.decoded_token!r})"
+        )
+        assert ref_logprob.rank == spec_logprob.rank, (
+            f"Rank mismatch: ref={ref_logprob.rank} "
+            f"spec={spec_logprob.rank} "
+            f"(token={ref_logprob.decoded_token!r})"
         )
-        assert ref_logprob.rank == spec_logprob.rank
         assert ref_logprob.decoded_token == spec_logprob.decoded_token
 
 

tests/v1/sample/test_rejection_sampler.py

@@ -11,7 +11,11 @@ from tests.v1.sample.utils import create_allowed_token_ids
 from vllm.platforms import current_platform
 from vllm.v1.sample.logits_processor import LogitsProcessors
 from vllm.v1.sample.metadata import SamplingMetadata
-from vllm.v1.sample.rejection_sampler import PLACEHOLDER_TOKEN_ID, RejectionSampler
+from vllm.v1.sample.rejection_sampler import (
+    PLACEHOLDER_TOKEN_ID,
+    RejectionSampler,
+    sample_recovered_tokens,
+)
 from vllm.v1.sample.sampler import Sampler, SamplerOutput
 from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
 
@@ -518,6 +522,70 @@ def estimate_rejection_sampling_pdf(
     return hist.hist
 
 
+def native_sample_recovered_tokens(
+    max_spec_len: int,
+    num_draft_tokens: list[int],
+    cu_num_draft_tokens: torch.Tensor,  # [batch_size]
+    draft_token_ids: torch.Tensor,  # [num_tokens]
+    draft_probs: torch.Tensor | None,  # [num_tokens, vocab_size]
+    target_probs: torch.Tensor,  # [num_tokens, vocab_size]
+    sampling_metadata: SamplingMetadata,
+    device: torch.device,
+) -> torch.Tensor:
+    batch_size = len(num_draft_tokens)
+    vocab_size = target_probs.shape[-1]
+
+    q = torch.empty(
+        (batch_size, vocab_size),
+        dtype=torch.float32,
+        device=device,
+    )
+    q.exponential_()
+
+    states = {
+        i: generator.get_state()
+        for i, generator in sampling_metadata.generators.items()
+    }
+    for i, generator in sampling_metadata.generators.items():
+        # Do not generate random numbers for requests with no draft tokens.
+        # This can be important for reproducibility.
+        if num_draft_tokens[i] > 0:
+            q[i].exponential_(generator=generator)
+
+        # In order to generate the same exponential later, reset the CUDA RNG
+        # state because RNG state advances after each call.
+        generator.set_state(states[i])
+
+    inv_q = q.reciprocal()
+
+    out = torch.empty_like(draft_token_ids)
+
+    for req_idx in range(batch_size):
+        start_idx = 0 if req_idx == 0 else int(cu_num_draft_tokens[req_idx - 1].item())
+        end_idx = int(cu_num_draft_tokens[req_idx].item())
+        num_tokens = end_idx - start_idx
+
+        for pos in range(max_spec_len):
+            if pos >= num_tokens:
+                continue
+            token_idx = start_idx + pos
+
+            if draft_probs is None:
+                # prob is target_probs[token_idx] except draft_token_id is zeroed
+                prob = target_probs[token_idx].clone()
+                draft_token_id = draft_token_ids[token_idx]
+                prob[draft_token_id] = 0.0
+            else:
+                prob = (target_probs[token_idx] - draft_probs[token_idx]).clamp_min_(
+                    0.0
+                )
+
+            score = prob * inv_q[req_idx]
+            recovered_id = torch.argmax(score, dim=-1)
+            out[token_idx] = recovered_id
+    return out
+
+
 def _test_masked_logits(
     rejection_sampler,
     batch_size: int,
@@ -778,3 +846,60 @@ def test_allowed_token_ids(rejection_sampler):
         device=logits.device,
     )
     assert torch.equal(output.sampled_token_ids, expected)
+
+
+@pytest.mark.parametrize("batch_size", [1, 100])
+@pytest.mark.parametrize("vocab_size", [100, 8192, 10000])
+@pytest.mark.parametrize("max_spec_len", [1, 3])
+@pytest.mark.parametrize("no_draft_probs", [True, False])
+def test_sample_recovered_tokens(
+    batch_size: int, vocab_size: int, max_spec_len: int, no_draft_probs: bool
+):
+    num_tokens = batch_size * max_spec_len
+
+    # Create random draft probabilities.
+    draft_probs = torch.rand(num_tokens, vocab_size, dtype=torch.float32, device=DEVICE)
+    draft_probs = F.softmax(draft_probs, dim=-1)
+
+    # Create random target probabilities.
+    target_logits = torch.rand(
+        num_tokens, vocab_size, dtype=torch.float32, device=DEVICE
+    )
+    target_probs = F.softmax(target_logits, dim=-1)
+
+    # Randomly sample draft token ids from draft probs
+    draft_token_ids = torch.multinomial(draft_probs, num_samples=1).to(torch.int32)
+
+    temperature = torch.ones(batch_size, dtype=torch.float32, device=DEVICE)
+    generators = {
+        i: torch.Generator(device=DEVICE).manual_seed(i) for i in range(batch_size)
+    }
+    sampling_metadata = create_sampling_metadata(
+        all_greedy=False, temperature=temperature, generators=generators
+    )
+
+    spec_decode_metadata = create_spec_decode_metadata(
+        draft_token_ids.reshape(batch_size, max_spec_len).tolist(), target_logits
+    )
+
+    ref_recovered_token_ids = native_sample_recovered_tokens(
+        max_spec_len,
+        spec_decode_metadata.num_draft_tokens,
+        spec_decode_metadata.cu_num_draft_tokens,
+        draft_token_ids,
+        None if no_draft_probs else draft_probs,
+        target_probs,
+        sampling_metadata,
+        device=DEVICE,
+    )
+    recovered_token_ids = sample_recovered_tokens(
+        max_spec_len,
+        spec_decode_metadata.num_draft_tokens,
+        spec_decode_metadata.cu_num_draft_tokens,
+        draft_token_ids,
+        None if no_draft_probs else draft_probs,
+        target_probs,
+        sampling_metadata,
+        device=DEVICE,
+    )
+    assert torch.equal(recovered_token_ids, ref_recovered_token_ids)

tests/v1/spec_decode/test_tree_attention.py

@@ -13,6 +13,7 @@ from tests.v1.attention.utils import (
     try_get_attention_backend,
 )
 from vllm.config import ParallelConfig, SpeculativeConfig
+from vllm.platforms import current_platform
 from vllm.v1.attention.backend import CommonAttentionMetadata
 from vllm.v1.attention.backends.fa_utils import is_flash_attn_varlen_func_available
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
@@ -23,11 +24,156 @@ if not is_flash_attn_varlen_func_available():
         allow_module_level=True,
     )
 
+# --------------------------------------------------------------------------- #
+#  KV cache layout adaptation
+# --------------------------------------------------------------------------- #
+# Two KV cache layouts exist across backends:
+#
+#   Flash layout: (2, num_blocks, block_size, num_kv_heads, head_size)
+#     - dim 0 separates key (index 0) and value (index 1)
+#     - Used by: FLASH_ATTN, TREE_ATTN, ROCM_AITER_FA, ROCM_ATTN
+#
+#   Block layout: (num_blocks, 2, block_size, num_kv_heads, head_size)
+#     - dim 1 separates key (index 0) and value (index 1)
+#     - Used by: TRITON_ATTN
+#
+# The test creates KV caches in flash layout (the canonical format used by
+# tree attention). When a reference backend needs block layout we transpose
+# dims 0 and 1.
+#
+# Note: ROCM_ATTN uses flash layout for storage but its forward path calls
+# PagedAttention.split_kv_cache which reinterprets the raw memory as paged
+# layout (num_blocks, num_kv_heads, head_size//x, block_size, x). This is
+# a view-level incompatibility, not a transpose - see the TODO in
+# _get_available_reference_backends for details.
+#
+# TODO: Replace this mapping with a `KV_CACHE_LAYOUT` class attribute on each
+# AttentionImpl so the layout is self-documented by the backend itself, e.g.:
+#     class TritonAttentionImpl(AttentionImpl):
+#         KV_CACHE_LAYOUT = "block"
+# --------------------------------------------------------------------------- #
+
+_BLOCK_KV_LAYOUT_BACKENDS = frozenset(
+    {
+        AttentionBackendEnum.TRITON_ATTN,
+    }
+)
+
+# Backends whose do_kv_cache_update requires engine-level state (e.g.
+# ForwardContext) that is not available in this test harness, but whose
+# KV cache is flash layout and can be written with reshape_and_cache_flash.
+# When a backend is listed here, forward_attention() bypasses
+# do_kv_cache_update and writes directly to the cache.
+_NEEDS_DIRECT_CACHE_UPDATE = frozenset(
+    {
+        AttentionBackendEnum.ROCM_AITER_FA,
+    }
+)
+
+# Backends with known test-harness incompatibilities - see the TODOs
+# inside _get_available_reference_backends for details.
+_INCOMPATIBLE_REFERENCE_BACKENDS = frozenset(
+    {
+        AttentionBackendEnum.ROCM_AITER_FA,
+        AttentionBackendEnum.ROCM_ATTN,
+    }
+)
+
+
+def _adapt_kv_cache_for_backend(
+    kv_cache: torch.Tensor,
+    backend: AttentionBackendEnum,
+) -> torch.Tensor:
+    """Convert kv_cache from flash layout ``(2, num_blocks, ...)`` to block
+    layout ``(num_blocks, 2, ...)`` if the backend requires it.  Returns the
+    original tensor unchanged when no conversion is needed."""
+    if backend in _BLOCK_KV_LAYOUT_BACKENDS:
+        return kv_cache.transpose(0, 1).contiguous()
+    return kv_cache
+
+
+def _get_platform_default_backend() -> AttentionBackendEnum:
+    """Ask the platform what backend it would auto-select at runtime."""
+    from vllm.v1.attention.selector import AttentionSelectorConfig
+
+    config = AttentionSelectorConfig(
+        block_size=32,
+        kv_cache_dtype="auto",
+        use_mla=False,
+        use_sparse=False,
+        head_size=128,
+        dtype=torch.bfloat16,
+    )
+    backend_path = current_platform.get_attn_backend_cls(
+        selected_backend=None,
+        attn_selector_config=config,
+    )
+    for backend in AttentionBackendEnum:
+        try:
+            if backend.get_path() == backend_path:
+                return backend
+        except ValueError:
+            continue
+    raise RuntimeError(
+        f"Platform returned backend path '{backend_path}' "
+        f"that doesn't match any AttentionBackendEnum member."
+    )
+
+
+def _get_available_reference_backends() -> list[AttentionBackendEnum]:
+    """Collect all reference backends the current platform can run.
+
+    On CUDA this is just FLASH_ATTN. On ROCm this includes the platform
+    default plus every backend the hardware supports, so the test validates
+    tree attention against all of them.
+    """
+    if current_platform.is_rocm():
+        backends: list[AttentionBackendEnum] = []
+
+        # 1. Whatever the platform would auto-select at runtime.
+        default_backend = _get_platform_default_backend()
+        if default_backend not in _INCOMPATIBLE_REFERENCE_BACKENDS:
+            backends.append(default_backend)
+
+        # 2. TRITON_ATTN - always available on ROCm.
+        if AttentionBackendEnum.TRITON_ATTN not in backends:
+            backends.append(AttentionBackendEnum.TRITON_ATTN)
+
+        # TODO: Enable ROCM_ATTN. Its forward path uses
+        # PagedAttention.split_kv_cache which reinterprets the raw
+        # cache memory as paged layout:
+        #   key:   (num_blocks, num_kv_heads, head_size//x, block_size, x)
+        #   value: (num_blocks, num_kv_heads, head_size, block_size)
+        # Tree attention writes prefix data in NHD flash layout, so the
+        # same bytes produce completely different values when read in
+        # paged format. Supporting ROCM_ATTN would require writing
+        # prefix data via PagedAttention.write_to_paged_cache into a
+        # separate paged-format KV cache.
+
+        # TODO: Enable ROCM_AITER_FA. Its metadata builder reads head
+        # counts from the model config at construction time and
+        # allocates extend_workspace with those dimensions. The test
+        # uses independent head count parameters (num_heads=2/4,
+        # num_kv_heads=2) that don't match the model config
+        # (Llama-3-8B: 32 q heads, 8 kv heads), causing a head count
+        # mismatch in flash_attn_varlen_func during extend_forward.
+        # Fixing this requires either matching test head counts to the
+        # model config or decoupling the builder from model config
+        # head geometry. The direct cache update path
+        # (_NEEDS_DIRECT_CACHE_UPDATE) is already in place for when
+        # this is resolved.
+
+        return backends
+
+    # CUDA: flash attention.
+    return [AttentionBackendEnum.FLASH_ATTN]
+
 
 class MockAttentionLayer(torch.nn.Module):
     _q_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda")
     _k_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda")
     _v_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda")
+    layer_name = "mock_layer"
 
     def __init__(self):
         super().__init__()
@@ -48,6 +194,13 @@ def forward_attention(
     spec_token_tree: str | None = None,
     num_spec_tokens: int = 0,
 ) -> torch.Tensor:
+    """Run a single attention forward pass through the given backend.
+
+    ``kv_cache`` is expected in **flash layout**
+    ``(2, num_blocks, block_size, num_kv_heads, head_size)``.
+    It is automatically converted when the target backend needs a
+    different layout.
+    """
     batch_size, q_len, num_heads, dim_per_head = q.shape
     num_kv_heads = k.shape[-2]
     # Initialize the query and KV sequence lengths.
@@ -116,31 +269,58 @@ def forward_attention(
         kv_cache_dtype="auto",
     )
 
+    # Adapt KV cache layout for this backend.
+    adapted_kv_cache = _adapt_kv_cache_for_backend(kv_cache, backend)
+
     # Run forward pass and return output.
     query = q.view(-1, num_heads, dim_per_head)
     key = k.view(-1, num_kv_heads, dim_per_head)
     value = v.view(-1, num_kv_heads, dim_per_head)
     output = torch.empty_like(query)
     if not try_backend_includes_kv_cache_update(backend):
-        instance.do_kv_cache_update(
-            layer=layer,
-            key=key,
-            value=value,
-            kv_cache=kv_cache,
-            slot_mapping=attn_metadata.slot_mapping,
-        )
+        if backend in _NEEDS_DIRECT_CACHE_UPDATE:
+            # This backend's do_kv_cache_update requires engine-level
+            # ForwardContext that isn't available in this test harness.
+            # Write directly using reshape_and_cache_flash since the
+            # KV cache layout is identical (flash layout, unbind on dim 0).
+            key_cache, value_cache = adapted_kv_cache.unbind(0)
+            torch.ops._C_cache_ops.reshape_and_cache_flash(
+                key,
+                value,
+                key_cache,
+                value_cache,
+                attn_metadata.slot_mapping,
+                "auto",
+                layer._k_scale,
+                layer._v_scale,
+            )
+        else:
+            instance.do_kv_cache_update(
+                layer=layer,
+                key=key,
+                value=value,
+                kv_cache=adapted_kv_cache,
+                slot_mapping=attn_metadata.slot_mapping,
+            )
     return instance.forward(
         layer=layer,
         query=query,
         key=key,
         value=value,
-        kv_cache=kv_cache.clone(),
+        kv_cache=adapted_kv_cache.clone(),
         attn_metadata=attn_metadata,
         output=output,
     )
 
 
-def test_tree_attn_correctness() -> None:
+@pytest.mark.parametrize(
+    "reference_backend",
+    _get_available_reference_backends(),
+    ids=lambda b: b.name,
+)
+def test_tree_attn_correctness(
+    reference_backend: AttentionBackendEnum,
+) -> None:
     torch.manual_seed(42)
     torch.cuda.manual_seed_all(42)
 
@@ -205,7 +385,9 @@ def test_tree_attn_correctness() -> None:
                         dtype=torch.bfloat16,
                     )
 
-                    # Set up the block table and KV cache for paged KV.
+                    # KV cache in flash layout - the canonical format for
+                    # tree attention. forward_attention() handles conversion
+                    # when needed.
                     assert max_sequence_length % block_size == 0
                     max_blocks_per_batch = max_sequence_length // block_size
                     kv_cache = torch.randn(
@@ -263,9 +445,7 @@ def test_tree_attn_correctness() -> None:
                         num_spec_tokens=tree_size_q - 1,
                     ).view(batch_size, -1, num_heads, dim_per_head)
 
-                    # Verify that the chain attention output for each
-                    # branch of the tree (computed using FA3) matches
-                    # the tree attention output.
+                    # Verify each branch against the reference backend.
                     for q_index in range(tree_size_q):
                         # Get the q, k, and v for the branch.
                         branch_mask = tree_attn_mask[q_index, :]
@@ -286,8 +466,8 @@ def test_tree_attn_correctness() -> None:
                             branch_positions, block_table, block_size
                         )
 
-                        # Compute flash attention for the branch.
-                        flash_attn_output = forward_attention(
+                        # Reference attention for this branch.
+                        ref_output = forward_attention(
                             q=q_branch,
                             k=k_branch,
                             v=v_branch,
@@ -295,16 +475,17 @@ def test_tree_attn_correctness() -> None:
                             block_table=block_table,
                             slot_mapping=branch_slot_mapping,
                             seqlen_k=sequence_position + q_len,
-                            backend=AttentionBackendEnum.FLASH_ATTN,
+                            backend=reference_backend,
                         ).view(batch_size, -1, num_heads, dim_per_head)
 
                         # Compare the outputs.
                         assert torch.allclose(
                             tree_attn_output[:, branch_indices],
-                            flash_attn_output,
+                            ref_output,
                             atol=7.81e-3,
                         ), (
                             f"outputs are not close for "
+                            f"reference_backend: {reference_backend.name}, "
                             f"batch_size: {batch_size}, "
                             f"num_heads: {num_heads}, "
                             f"sequence_position: {sequence_position}, "

tools/pre_commit/mypy.py

@@ -43,7 +43,6 @@ EXCLUDE = [
     "vllm/benchmarks",
     "vllm/config",
     "vllm/device_allocator",
-    "vllm/profiler",
     "vllm/reasoning",
     "vllm/tool_parser",
 ]

vllm/benchmarks/datasets.py

@@ -2627,22 +2627,26 @@ class VisionArenaDataset(HuggingFaceDataset):
         no_oversample: bool = False,
         **kwargs,
     ) -> list:
+        parser_fn = self.SUPPORTED_DATASET_PATHS.get(self.hf_name)
+        if parser_fn is None:
+            raise ValueError(f"Unsupported dataset path: {self.hf_name}")
+
         output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN
+
         sampled_requests = []
         for i, item in enumerate(self.data):
             if len(sampled_requests) >= num_requests:
                 break
-            parser_fn = self.SUPPORTED_DATASET_PATHS.get(self.hf_name)
-            if parser_fn is None:
-                raise ValueError(f"Unsupported dataset path: {self.hf_name}")
+
             prompt = parser_fn(item)
             mm_content = process_image(item["images"][0])
-            prompt_len = len(tokenizer(prompt).input_ids)
+            prompt_len = len(tokenizer.encode(prompt))
             if enable_multimodal_chat:
                 # Note: when chat is enabled the request prompt_len is no longer
                 # accurate and we will be using request output to count the
                 # actual prompt len
                 prompt = self.apply_multimodal_chat_transformation(prompt, mm_content)
+
             sampled_requests.append(
                 SampleRequest(
                     prompt=prompt,
@@ -2652,6 +2656,7 @@ class VisionArenaDataset(HuggingFaceDataset):
                     request_id=request_id_prefix + str(i),
                 )
             )
+
         self.maybe_oversample_requests(
             sampled_requests, num_requests, request_id_prefix, no_oversample
         )
@@ -2681,22 +2686,26 @@ class MMVUDataset(HuggingFaceDataset):
         no_oversample: bool = False,
         **kwargs,
     ) -> list:
+        parser_fn = self.SUPPORTED_DATASET_PATHS.get(self.hf_name)
+        if parser_fn is None:
+            raise ValueError(f"Unsupported dataset path: {self.hf_name}")
+
         output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN
+
         sampled_requests = []
         for i, item in enumerate(self.data):
             if len(sampled_requests) >= num_requests:
                 break
-            parser_fn = self.SUPPORTED_DATASET_PATHS.get(self.hf_name)
-            if parser_fn is None:
-                raise ValueError(f"Unsupported dataset path: {self.hf_name}")
+
             prompt = parser_fn(item)
             mm_content = process_video(item["video"])
-            prompt_len = len(tokenizer(prompt).input_ids)
+            prompt_len = len(tokenizer.encode(prompt))
             if enable_multimodal_chat:
                 # Note: when chat is enabled the request prompt_len is no longer
                 # accurate and we will be using request output to count the
                 # actual prompt len
                 prompt = self.apply_multimodal_chat_transformation(prompt, mm_content)
+
             sampled_requests.append(
                 SampleRequest(
                     prompt=prompt,
@@ -2706,6 +2715,7 @@ class MMVUDataset(HuggingFaceDataset):
                     request_id=request_id_prefix + str(i),
                 )
             )
+
         self.maybe_oversample_requests(
             sampled_requests, num_requests, request_id_prefix, no_oversample
         )

vllm/benchmarks/sweep/plot.py

@@ -19,11 +19,17 @@ from .utils import sanitize_filename
 
 try:
     import matplotlib.pyplot as plt
-    import pandas as pd
-    import seaborn as sns
 except ImportError:
     plt = PlaceholderModule("matplotlib").placeholder_attr("pyplot")
+
+try:
+    import pandas as pd
+except ImportError:
     pd = PlaceholderModule("pandas")
+
+try:
+    import seaborn as sns
+except ImportError:
     seaborn = PlaceholderModule("seaborn")
 
 
@@ -340,7 +346,45 @@ def _plot_fig(
         else "(All)"
     )
 
-    g = sns.FacetGrid(df, row="row_group", col="col_group", height=fig_height)
+    if len(curve_by) <= 3:
+        hue, style, size, *_ = (*curve_by, None, None, None)
+
+        g = sns.relplot(
+            df,
+            x=var_x,
+            y=var_y,
+            hue=hue,
+            style=style,
+            size=size,
+            markers=True,
+            errorbar="sd" if error_bars else None,
+            kind="line",
+            row="row_group",
+            col="col_group",
+            height=fig_height,
+        )
+    else:
+        df["curve_group"] = (
+            pd.concat(
+                [k + "=" + df[k].astype(str) for k in curve_by],
+                axis=1,
+            ).agg("\n".join, axis=1)
+            if curve_by
+            else "(All)"
+        )
+
+        g = sns.relplot(
+            df,
+            x=var_x,
+            y=var_y,
+            hue="curve_group",
+            markers=True,
+            errorbar="sd" if error_bars else None,
+            kind="line",
+            row="row_group",
+            col="col_group",
+            height=fig_height,
+        )
 
     if row_by and col_by:
         g.set_titles("{row_name}\n{col_name}")
@@ -356,42 +400,6 @@ def _plot_fig(
     if scale_y:
         g.set(yscale=scale_y)
 
-    if len(curve_by) <= 3:
-        hue, style, size, *_ = (*curve_by, None, None, None)
-
-        g.map_dataframe(
-            sns.lineplot,
-            x=var_x,
-            y=var_y,
-            hue=hue,
-            style=style,
-            size=size,
-            markers=True,
-            errorbar="sd" if error_bars else None,
-        )
-
-        g.add_legend(title=hue)
-    else:
-        df["curve_group"] = (
-            pd.concat(
-                [k + "=" + df[k].astype(str) for k in curve_by],
-                axis=1,
-            ).agg("\n".join, axis=1)
-            if curve_by
-            else "(All)"
-        )
-
-        g.map_dataframe(
-            sns.lineplot,
-            x=var_x,
-            y=var_y,
-            hue="curve_group",
-            markers=True,
-            errorbar="sd" if error_bars else None,
-        )
-
-        g.add_legend()
-
     g.savefig(fig_path, dpi=fig_dpi)
     plt.close(g.figure)
 

vllm/benchmarks/sweep/plot_pareto.py

@@ -16,12 +16,18 @@ from .utils import sanitize_filename
 
 try:
     import matplotlib.pyplot as plt
-    import pandas as pd
-    import seaborn as sns
 except ImportError:
     plt = PlaceholderModule("matplotlib").placeholder_attr("pyplot")
+
+try:
+    import pandas as pd
+except ImportError:
     pd = PlaceholderModule("pandas")
-    sns = PlaceholderModule("seaborn")
+
+try:
+    import seaborn as sns
+except ImportError:
+    seaborn = PlaceholderModule("seaborn")
 
 
 def _first_present(run_data: dict[str, object], keys: list[str]):

vllm/benchmarks/sweep/serve_sla.py

@@ -202,6 +202,7 @@ def solve_sla(
         with path.open("rb") as f:
             past_iter_data = json.load(f)
 
+        sla_data.append(past_iter_data)
         history[past_sla_value] = _compute_margin(sla_comb, past_iter_data)
 
     # NOTE: We don't use equality here to be more robust against noisy results
@@ -264,6 +265,8 @@ def search_sla(
     dry_run: bool,
 ):
     print("[SLA START]")
+    print(f"Serve parameters: {serve_comb.as_text() or '(None)'}")
+    print(f"Bench parameters: {bench_comb.as_text() or '(None)'}")
     print(f"SLA criteria: {sla_comb.as_text()}")
 
     result = solve_sla(

vllm/compilation/backends.py

@@ -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 "
@@ -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"
         )

vllm/compilation/decorators.py

@@ -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,
             )

vllm/compilation/monitor.py

@@ -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

vllm/compilation/piecewise_backend.py

@@ -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:

vllm/config/cache.py

@@ -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

vllm/config/kv_transfer.py

@@ -61,10 +61,10 @@ class KVTransferConfig:
     enable_permute_local_kv: bool = False
     """Experiment feature flag to enable HND to NHD KV Transfer"""
 
-    kv_load_failure_policy: Literal["recompute", "fail"] = "recompute"
+    kv_load_failure_policy: Literal["recompute", "fail"] = "fail"
     """Policy for handling KV cache load failures.
-    'recompute': reschedule the request to recompute failed blocks (default)
-    'fail': immediately fail the request with an error finish reason"""
+    'recompute': reschedule the request to recompute failed blocks
+    'fail': immediately fail the request with an error finish reason (default)"""
 
     def compute_hash(self) -> str:
         """

vllm/config/utils.py

@@ -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)

vllm/distributed/device_communicators/shm_broadcast.py

@@ -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.

vllm/distributed/kv_events.py

@@ -60,6 +60,13 @@ class BlockStored(KVCacheEvent):
     medium: str | None
     lora_name: str | None
 
+    extra_keys: list[tuple[Any, ...] | None] | None = None
+    """Extra keys used in block hash computation, one entry per block in
+    block_hashes. Each entry contains MM identifiers, LoRA name, cache_salt,
+    prompt embedding hashes, etc. for that specific block. Exposed for external
+    KV cache consumers to reconstruct block hashes.
+    """
+
     def __hash__(self) -> int:
         return hash(
             (
@@ -69,6 +76,7 @@ class BlockStored(KVCacheEvent):
                 self.block_size,
                 self.lora_id,
                 self.medium,
+                tuple(self.extra_keys) if self.extra_keys else None,
             )
         )
 

vllm/engine/arg_utils.py

@@ -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

vllm/entrypoints/llm.py

@@ -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,
@@ -50,6 +50,7 @@ from vllm.entrypoints.pooling.score.utils import (
     compress_token_type_ids,
     compute_maxsim_score,
     get_score_prompt,
+    score_data_to_prompts,
     validate_score_input,
 )
 from vllm.entrypoints.utils import log_non_default_args
@@ -94,6 +95,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 +453,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],
@@ -524,23 +529,43 @@ class LLM:
 
         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)
+
+        return self._run_engine(output_type, use_tqdm=use_tqdm)
 
     def _resolve_mm_lora(
         self,
@@ -744,13 +769,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
@@ -987,9 +1012,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,
@@ -1002,8 +1028,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,
@@ -1135,19 +1159,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](
@@ -1160,8 +1181,8 @@ class LLM:
                     finished=True,
                 )
             ]
-        else:
-            return model_outputs
+
+        return outputs
 
     def embed(
         self,
@@ -1353,8 +1374,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,
@@ -1376,25 +1396,13 @@ class LLM:
 
         tokenizer = self.get_tokenizer()
 
-        # Extract text from ScoreData
-        text_1: list[str] = []
-        for text in data_1:
-            if not isinstance(text, str):
-                raise NotImplementedError(
-                    "Late interaction scores currently do not support multimodal input."
-                )
-            text_1.append(text)
-
-        text_2: list[str] = []
-        for text in data_2:
-            if not isinstance(text, str):
-                raise NotImplementedError(
-                    "Late interaction scores currently do not support multimodal input."
-                )
-            text_2.append(text)
+        # Convert ScoreData to PromptType (handles both text and multimodal)
+        model_config = self.model_config
+        prompts_1 = score_data_to_prompts(data_1, "query", model_config)
+        prompts_2 = score_data_to_prompts(data_2, "document", model_config)
 
         encoded_output: list[PoolingRequestOutput] = self.encode(
-            text_1 + text_2,
+            prompts_1 + prompts_2,
             use_tqdm=use_tqdm,
             lora_request=lora_request,
             pooling_params=pooling_params,
@@ -1402,8 +1410,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: list[PoolingRequestOutput] = encoded_output[: len(prompts_1)]
+        encoded_output_2: list[PoolingRequestOutput] = encoded_output[len(prompts_1) :]
 
         if len(encoded_output_1) == 1:
             encoded_output_1 = encoded_output_1 * len(encoded_output_2)
@@ -1434,8 +1442,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,
@@ -1491,13 +1498,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,
@@ -1759,6 +1765,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,
@@ -1790,6 +1797,7 @@ class LLM:
                 )
             ),
             params=seq_params,
+            output_type=output_type,
             use_tqdm=use_tqdm,
             lora_requests=seq_lora_requests,
             priorities=seq_priority,
@@ -1802,6 +1810,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,
@@ -1848,6 +1857,7 @@ class LLM:
                 )
             ),
             params=seq_params,
+            output_type=output_type,
             lora_requests=seq_lora_requests,
             use_tqdm=use_tqdm,
         )
@@ -1856,6 +1866,7 @@ class LLM:
         self,
         prompts: Iterable[ProcessorInputs],
         params: Sequence[SamplingParams | PoolingParams],
+        output_type: type[_O],
         *,
         lora_requests: Sequence[LoRARequest | None] | None = None,
         priorities: Sequence[int] | None = None,
@@ -1878,7 +1889,7 @@ class LLM:
             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,
@@ -1932,9 +1943,10 @@ class LLM:
 
     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()
@@ -1947,14 +1959,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

vllm/entrypoints/openai/chat_completion/protocol.py

@@ -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

vllm/entrypoints/openai/chat_completion/serving.py

@@ -900,6 +900,17 @@ class OpenAIServingChat(OpenAIServing):
                         harmony_tools_streamed[i] |= tools_streamed_flag
                     # handle streaming deltas for tools with named tool_choice
                     elif tool_choice_function_name:
+                        # When encountering think end id in prompt_token_ids
+                        # i.e {"enable_thinking": False},
+                        # check BEFORE calling the parser to avoid a spurious
+                        # reasoning delta on the first chunk.
+                        if (
+                            reasoning_parser
+                            and not reasoning_end_arr[i]
+                            and prompt_is_reasoning_end_arr[i]
+                        ):
+                            reasoning_end_arr[i] = True
+
                         if (
                             reasoning_parser
                             and not reasoning_end_arr[i]
@@ -918,16 +929,11 @@ class OpenAIServingChat(OpenAIServing):
                                     output.token_ids,
                                 )
                             )
-                            # When encountering think end id in delta_token_ids
-                            # or think end id in prompt_token_ids
-                            # i.e {"enable_thinking": False},
+                            # When encountering think end id in delta_token_ids,
                             # set reasoning status to end.
                             # Only keep 'content', remove 'reasoning'.
-                            if (
-                                reasoning_parser.is_reasoning_end(
-                                    as_list(output.token_ids)
-                                )
-                                or prompt_is_reasoning_end_arr[i]
+                            if reasoning_parser.is_reasoning_end(
+                                as_list(output.token_ids)
                             ):
                                 reasoning_end_arr[i] = True
                                 if delta_message and delta_message.content:
@@ -1116,14 +1122,23 @@ class OpenAIServingChat(OpenAIServing):
 
                     # when only reasoning
                     elif reasoning_parser:
-                        delta_message = reasoning_parser.extract_reasoning_streaming(
-                            previous_text,
-                            current_text,
-                            delta_text,
-                            previous_token_ids,
-                            current_token_ids,
-                            output.token_ids,
-                        )
+                        # When encountering think end id in prompt_token_ids
+                        # i.e {"enable_thinking": False},
+                        # set reasoning status to end.
+                        # Route all generated tokens as content directly.
+                        if prompt_is_reasoning_end_arr[i]:
+                            delta_message = DeltaMessage(content=delta_text)
+                        else:
+                            delta_message = (
+                                reasoning_parser.extract_reasoning_streaming(
+                                    previous_text,
+                                    current_text,
+                                    delta_text,
+                                    previous_token_ids,
+                                    current_token_ids,
+                                    output.token_ids,
+                                )
+                            )
                     # handle streaming just a content delta
                     else:
                         delta_message = DeltaMessage(content=delta_text)

vllm/entrypoints/openai/parser/harmony_utils.py

@@ -48,8 +48,11 @@ from vllm.entrypoints.openai.responses.protocol import (
     ResponseInputOutputItem,
     ResponsesRequest,
 )
+from vllm.logger import init_logger
 from vllm.utils import random_uuid
 
+logger = init_logger(__name__)
+
 REASONING_EFFORT = {
     "high": ReasoningEffort.HIGH,
     "medium": ReasoningEffort.MEDIUM,
@@ -62,20 +65,15 @@ _harmony_encoding = None
 # they are available and requested by the user.
 # Tool args are provided by MCP tool descriptions. Output
 # of the tools are stringified.
-MCP_BUILTIN_TOOLS: set[str] = {
-    "web_search_preview",
-    "code_interpreter",
-    "container",
-}
-
-# Mapping from built-in tool recipient names to their MCP server labels.
-# This ensures consistency between streaming and non-streaming responses.
 _BUILTIN_TOOL_TO_MCP_SERVER_LABEL: dict[str, str] = {
     "python": "code_interpreter",
     "browser": "web_search_preview",
     "container": "container",
 }
 
+# Derive MCP_BUILTIN_TOOLS from the canonical mapping
+MCP_BUILTIN_TOOLS: set[str] = set(_BUILTIN_TOOL_TO_MCP_SERVER_LABEL.values())
+
 
 def has_custom_tools(tool_types: set[str]) -> bool:
     """
@@ -116,8 +114,11 @@ def get_system_message(
             REASONING_EFFORT[reasoning_effort]
         )
     if start_date is None:
-        # NOTE(woosuk): This brings non-determinism in vLLM. Be careful.
-        start_date = datetime.datetime.now().strftime("%Y-%m-%d")
+        # NOTE(woosuk): This brings non-determinism in vLLM.
+        # Set VLLM_SYSTEM_START_DATE to pin it.
+        start_date = envs.VLLM_SYSTEM_START_DATE or datetime.datetime.now().strftime(
+            "%Y-%m-%d"
+        )
     sys_msg_content = sys_msg_content.with_conversation_start_date(start_date)
     if browser_description is not None:
         sys_msg_content = sys_msg_content.with_tools(browser_description)
@@ -398,15 +399,60 @@ def parse_chat_input_to_harmony_message(
 
 
 def parse_input_to_harmony_message(chat_msg) -> list[Message]:
-    """
-    Parse a message from request.previous_input_messages in the Responsees API to
-    Harmony messages.
+    """Parse a message from request.previous_input_messages
+    into Harmony messages.
+
+    Supports both OpenAI chat format ({"role": "..."}) and
+    Harmony format ({"author": {"role": "..."}}).
     """
     if not isinstance(chat_msg, dict):
-        # Handle Pydantic models
         chat_msg = chat_msg.model_dump(exclude_none=True)
 
+    if "author" in chat_msg and isinstance(chat_msg.get("author"), dict):
+        return [_parse_harmony_format_message(chat_msg)]
+
+    return _parse_chat_format_message(chat_msg)
+
+
+def _parse_harmony_format_message(chat_msg: dict) -> Message:
+    """Reconstruct a Message from Harmony-format dict,
+    preserving channel, recipient, and content_type."""
+    author_dict = chat_msg["author"]
+    role = author_dict.get("role")
+    name = author_dict.get("name")
+
+    raw_content = chat_msg.get("content", "")
+    if isinstance(raw_content, list):
+        # TODO: Support refusal and non-text content types.
+        contents = [TextContent(text=c.get("text", "")) for c in raw_content]
+    elif isinstance(raw_content, str):
+        contents = [TextContent(text=raw_content)]
+    else:
+        contents = [TextContent(text="")]
+
+    if name:
+        msg = Message.from_author_and_contents(Author.new(Role(role), name), contents)
+    else:
+        msg = Message.from_role_and_contents(Role(role), contents)
+
+    channel = chat_msg.get("channel")
+    if channel:
+        msg = msg.with_channel(channel)
+    recipient = chat_msg.get("recipient")
+    if recipient:
+        msg = msg.with_recipient(recipient)
+    content_type = chat_msg.get("content_type")
+    if content_type:
+        msg = msg.with_content_type(content_type)
+
+    return msg
+
+
+def _parse_chat_format_message(chat_msg: dict) -> list[Message]:
+    """Parse an OpenAI chat-format dict into Harmony messages."""
     role = chat_msg.get("role")
+    if role is None:
+        raise ValueError(f"Message has no 'role' key: {chat_msg}")
 
     # Assistant message with tool calls
     tool_calls = chat_msg.get("tool_calls")
@@ -426,15 +472,21 @@ def parse_input_to_harmony_message(chat_msg) -> list[Message]:
     # Tool role message (tool output)
     if role == "tool":
         name = chat_msg.get("name", "")
+        if name and not name.startswith("functions."):
+            name = f"functions.{name}"
         content = chat_msg.get("content", "") or ""
         content = flatten_chat_text_content(content)
-
-        msg = Message.from_author_and_content(
-            Author.new(Role.TOOL, f"functions.{name}"), content
-        ).with_channel("commentary")
+        # NOTE: .with_recipient("assistant") is required on tool messages
+        # to match parse_chat_input_to_harmony_message behavior and ensure
+        # proper routing in the Harmony protocol.
+        msg = (
+            Message.from_author_and_content(Author.new(Role.TOOL, name), content)
+            .with_channel("commentary")
+            .with_recipient("assistant")
+        )
         return [msg]
 
-    # Default: user/assistant/system messages with content
+    # Default: user/assistant/system messages
     content = chat_msg.get("content", "")
     if isinstance(content, str):
         contents = [TextContent(text=content)]
@@ -497,6 +549,10 @@ def _parse_browser_tool_call(message: Message, recipient: str) -> ResponseOutput
     try:
         browser_call = json.loads(content.text)
     except json.JSONDecodeError:
+        logger.warning(
+            "Invalid JSON in browser tool call, using error placeholder: %s",
+            content.text,
+        )
         json_retry_output_message = (
             f"Invalid JSON args, caught and retried: {content.text}"
         )
@@ -730,22 +786,7 @@ def parse_remaining_state(parser: StreamableParser) -> list[ResponseOutputItem]:
                 )
             ]
 
-    if parser.current_channel == "commentary":
-        return [
-            ResponseReasoningItem(
-                id=f"rs_{random_uuid()}",
-                summary=[],
-                type="reasoning",
-                content=[
-                    ResponseReasoningTextContent(
-                        text=parser.current_content, type="reasoning_text"
-                    )
-                ],
-                status=None,
-            )
-        ]
-
-    if parser.current_channel == "analysis":
+    if parser.current_channel in ("commentary", "analysis"):
         return [
             ResponseReasoningItem(
                 id=f"rs_{random_uuid()}",

vllm/entrypoints/openai/realtime/connection.py

@@ -205,7 +205,7 @@ class RealtimeConnection:
 
             sampling_params = SamplingParams.from_optional(
                 temperature=0.0,
-                max_tokens=1,
+                max_tokens=self.serving.model_cls.realtime_max_tokens,
                 output_kind=RequestOutputKind.DELTA,
                 skip_clone=True,
             )

vllm/entrypoints/openai/responses/context.py

@@ -346,17 +346,17 @@ class ParsableContext(ConversationContext):
         self.parser.response_messages.extend(output)
 
     def need_builtin_tool_call(self) -> bool:
-        """Return true if the last message is a MCP tool call"""
+        """Return true if the last message is a builtin tool call
+        that the request has enabled."""
         last_message = self.parser.response_messages[-1]
-        # TODO(qandrew): figure out which tools are MCP tools
-        if last_message.type == "function_call":  # noqa: SIM102
-            if last_message.name in (
-                "code_interpreter",
-                "python",
-                "web_search_preview",
-            ) or last_message.name.startswith("container"):
-                return True
-
+        if last_message.type != "function_call":
+            return False
+        if last_message.name in ("code_interpreter", "python"):
+            return "python" in self.available_tools
+        if last_message.name == "web_search_preview":
+            return "browser" in self.available_tools
+        if last_message.name.startswith("container"):
+            return "container" in self.available_tools
         return False
 
     async def call_python_tool(
@@ -665,11 +665,15 @@ class HarmonyContext(ConversationContext):
     def need_builtin_tool_call(self) -> bool:
         last_msg = self.messages[-1]
         recipient = last_msg.recipient
-        return recipient is not None and (
-            recipient.startswith("browser.")
-            or recipient.startswith("python")
-            or recipient.startswith("container.")
-        )
+        if recipient is None:
+            return False
+        if recipient.startswith("browser."):
+            return "browser" in self.available_tools
+        if recipient.startswith("python"):
+            return "python" in self.available_tools
+        if recipient.startswith("container."):
+            return "container" in self.available_tools
+        return False
 
     async def call_tool(self) -> list[Message]:
         if not self.messages:

vllm/entrypoints/openai/responses/streaming_events.py

@@ -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

vllm/entrypoints/openai/speech_to_text/speech_to_text.py

@@ -41,7 +41,10 @@ from vllm.exceptions import VLLMValidationError
 from vllm.inputs import ProcessorInputs
 from vllm.logger import init_logger
 from vllm.logprobs import FlatLogprobs, Logprob
-from vllm.model_executor.models import SupportsTranscription, supports_transcription
+from vllm.model_executor.models import (
+    SupportsTranscription,
+    supports_transcription,
+)
 from vllm.outputs import RequestOutput
 from vllm.renderers.inputs import DictPrompt, EncoderDecoderDictPrompt
 from vllm.renderers.inputs.preprocess import parse_enc_dec_prompt, parse_model_prompt
@@ -242,10 +245,57 @@ class OpenAISpeechToText(OpenAIServing):
         model_cls = get_model_cls(self.model_config)
         return cast(type[SupportsTranscription], model_cls)
 
+    async def _detect_language(
+        self,
+        audio_chunk: np.ndarray,
+        request_id: str,
+    ) -> str:
+        """Auto-detect the spoken language from an audio chunk.
+
+        Delegates prompt construction and output parsing to the model class
+        via ``get_language_detection_prompt`` and
+        ``parse_language_detection_output``.
+        """
+        from vllm.sampling_params import SamplingParams
+
+        prompt = self.model_cls.get_language_detection_prompt(
+            audio_chunk,
+            self.asr_config,
+        )
+        allowed_token_ids = self.model_cls.get_language_token_ids(
+            self.tokenizer,
+        )
+        sampling_params = SamplingParams(
+            max_tokens=1,
+            temperature=0.0,
+            allowed_token_ids=allowed_token_ids,
+        )
+
+        result_generator = self.engine_client.generate(
+            prompt,
+            sampling_params,
+            request_id,
+        )
+
+        final_output: RequestOutput
+        async for final_output in result_generator:
+            if final_output.finished:
+                break
+
+        token_ids = list(final_output.outputs[0].token_ids)
+        lang = self.model_cls.parse_language_detection_output(
+            token_ids,
+            self.tokenizer,
+        )
+
+        logger.info("Auto-detected language: '%s'", lang)
+        return lang
+
     async def _preprocess_speech_to_text(
         self,
         request: SpeechToTextRequest,
         audio_data: bytes,
+        request_id: str,
     ) -> tuple[list[ProcessorInputs], float]:
         # Validate request
         language = self.model_cls.validate_language(request.language)
@@ -274,6 +324,15 @@ class OpenAISpeechToText(OpenAIServing):
             and duration > self.asr_config.max_audio_clip_s
         )
         chunks = [y] if not do_split_audio else self._split_audio(y, int(sr))
+
+        if language is None and getattr(
+            self.model_cls, "supports_explicit_language_detection", False
+        ):
+            language = await self._detect_language(
+                chunks[0], f"{request_id}-lang_detect"
+            )
+            request.language = language
+
         parsed_prompts: list[DictPrompt] = []
         for chunk in chunks:
             # The model has control over the construction, as long as it
@@ -435,6 +494,7 @@ class OpenAISpeechToText(OpenAIServing):
             engine_prompts, duration_s = await self._preprocess_speech_to_text(
                 request=request,
                 audio_data=audio_data,
+                request_id=request_id,
             )
 
         except ValueError as e:

vllm/entrypoints/pooling/score/serving.py

@@ -33,6 +33,7 @@ from vllm.entrypoints.pooling.score.utils import (
     compress_token_type_ids,
     compute_maxsim_score,
     get_score_prompt,
+    parse_score_data_single,
     validate_score_input,
 )
 from vllm.inputs.data import ProcessorInputs, TokensPrompt, token_inputs
@@ -174,6 +175,43 @@ class ServingScores(OpenAIServing):
 
         return final_res_batch
 
+    def _preprocess_late_interaction_item(
+        self,
+        data: ScoreData,
+        role: str,
+        request: RerankRequest | ScoreRequest,
+        tokenizer: TokenizerLike,
+        tokenization_kwargs: dict[str, Any],
+    ) -> tuple[str, TokensPrompt]:
+        """Parse a single ScoreData into a text + optional multimodal
+        TokensPrompt for late-interaction encoding.
+
+        For plain strings, tokenises directly.
+        For multimodal content parts, extracts text and multi_modal_data.
+        """
+        model_config = self.model_config
+
+        if isinstance(data, str):
+            text, mm_data, mm_uuids = data, None, None
+        else:
+            text, mm_data, mm_uuids = parse_score_data_single(data, role, model_config)
+
+        prompt_inputs = tokenizer(text, **tokenization_kwargs)
+        self._validate_input(request, prompt_inputs["input_ids"], text)
+
+        engine_prompt = TokensPrompt(
+            prompt_token_ids=prompt_inputs["input_ids"],
+        )
+
+        if mm_data is not None:
+            engine_prompt["multi_modal_data"] = mm_data
+        if mm_uuids is not None:
+            engine_prompt["multi_modal_uuids"] = mm_uuids
+        if request.mm_processor_kwargs is not None:
+            engine_prompt["mm_processor_kwargs"] = request.mm_processor_kwargs
+
+        return text, engine_prompt
+
     async def _late_interaction_score(
         self,
         data_1: list[ScoreData],
@@ -189,37 +227,36 @@ class ServingScores(OpenAIServing):
         Encodes queries and documents into per-token embeddings, then computes
         MaxSim: sum over query tokens of max similarity to any document token.
         """
-        input_texts: list[str] = []
-        for text in data_1 + data_2:
-            if not isinstance(text, str):
-                raise NotImplementedError(
-                    "Late interaction scores currently do not support multimodal input."
-                )
-            input_texts.append(text)
-
         model_config = self.model_config
         tokenizer = self.renderer.get_tokenizer()
+        tokenization_kwargs = request.build_tok_params(model_config).get_encode_kwargs()
 
-        encode_async = make_async(
-            tokenizer.encode,
+        all_data = data_1 + data_2
+        roles = ["query"] * len(data_1) + ["document"] * len(data_2)
+
+        preprocess_async = make_async(
+            self._preprocess_late_interaction_item,
             executor=self._tokenizer_executor,
         )
 
-        tokenization_kwargs = request.build_tok_params(model_config).get_encode_kwargs()
-        tokenized_prompts = await asyncio.gather(
-            *(encode_async(t, **tokenization_kwargs) for t in input_texts)
+        preprocessed = await asyncio.gather(
+            *(
+                preprocess_async(
+                    data=d,
+                    role=r,
+                    request=request,
+                    tokenizer=tokenizer,
+                    tokenization_kwargs=tokenization_kwargs,
+                )
+                for d, r in zip(all_data, roles)
+            )
         )
 
-        engine_prompts: list[ProcessorInputs] = []
-        for tok_result, input_text in zip(tokenized_prompts, input_texts):
-            text_token_prompt = self._validate_input(request, tok_result, input_text)
-
-            engine_prompts.append(
-                token_inputs(
-                    text_token_prompt["prompt_token_ids"],
-                    prompt=input_text,
-                )
-            )
+        input_texts: list[str] = []
+        engine_prompts: list[TokensPrompt] = []
+        for text, engine_prompt in preprocessed:
+            input_texts.append(text)
+            engine_prompts.append(engine_prompt)
 
         # Schedule the request and get the result generator.
         generators: list[AsyncGenerator[PoolingRequestOutput, None]] = []

vllm/entrypoints/pooling/score/utils.py

@@ -21,6 +21,7 @@ from vllm.entrypoints.chat_utils import (
     _parse_chat_message_content_parts,
 )
 from vllm.inputs import TokensPrompt
+from vllm.inputs.data import PromptType, TextPrompt
 from vllm.model_executor.models.interfaces import supports_score_template
 from vllm.multimodal.inputs import MultiModalDataDict, MultiModalUUIDDict
 from vllm.outputs import PoolingRequestOutput
@@ -153,31 +154,91 @@ def validate_score_input(
     return score_input_1, score_input_2
 
 
+def _ensure_str(content: list[ConversationMessage]) -> str:
+    """Extract a single string prompt from parsed conversation content."""
+    assert len(content) == 1
+    prompt = content[0]["content"]
+    if prompt is not None and isinstance(prompt, str):
+        return cast(str, prompt)
+    raise ValueError(f"Only string content is supported, but got {content}.")
+
+
 def parse_score_data(
     data_1: ScoreData,
     data_2: ScoreData,
     model_config: ModelConfig,
 ) -> tuple[str, str, MultiModalDataDict | None, MultiModalUUIDDict | None]:
+    """Parse a query-document pair into text prompts and shared multi-modal
+    data.
+
+    Uses a **single** :class:`MultiModalItemTracker` so that multi-modal
+    items from both inputs are merged into one ``mm_data`` dict.  This is
+    the correct behaviour for cross-encoder scoring, where query and
+    document are concatenated into a single model prompt.
+    """
     mm_tracker = MultiModalItemTracker(model_config)
 
     content_1 = _parse_score_content("query", data_1, mm_tracker)
     content_2 = _parse_score_content("document", data_2, mm_tracker)
 
-    def ensure_str(content: list[ConversationMessage]) -> str:
-        assert len(content) == 1
-        prompt = content[0]["content"]
-        if prompt is not None and isinstance(prompt, str):
-            return cast(str, prompt)
-        else:
-            raise ValueError(f"Only string content is supported, but got {content}.")
-
-    prompt_1 = ensure_str(content_1)
-    prompt_2 = ensure_str(content_2)
+    prompt_1 = _ensure_str(content_1)
+    prompt_2 = _ensure_str(content_2)
     mm_items, mm_uuids = mm_tracker.resolve_items()
 
     return prompt_1, prompt_2, mm_items, mm_uuids
 
 
+def parse_score_data_single(
+    data: ScoreData,
+    role: str,
+    model_config: ModelConfig,
+) -> tuple[str, MultiModalDataDict | None, MultiModalUUIDDict | None]:
+    """Parse **one** ScoreData into a text prompt and its own multi-modal
+    data.
+
+    Unlike :func:`parse_score_data`, each call creates an **independent**
+    :class:`MultiModalItemTracker` so multi-modal items are kept separate.
+    This is the correct behaviour for late-interaction scoring, where
+    query and document are encoded independently.
+    """
+    mm_tracker = MultiModalItemTracker(model_config)
+    content = _parse_score_content(role, data, mm_tracker)
+
+    prompt = _ensure_str(content)
+    mm_items, mm_uuids = mm_tracker.resolve_items()
+    return prompt, mm_items, mm_uuids
+
+
+def score_data_to_prompts(
+    data_list: list[ScoreData],
+    role: str,
+    model_config: ModelConfig,
+) -> list[PromptType]:
+    """Convert a list of ScoreData into PromptType objects.
+
+    For plain text inputs, returns the string directly.
+    For multimodal inputs (list of content parts), parses them into
+    a :class:`TextPrompt` with attached ``multi_modal_data`` /
+    ``multi_modal_uuids``.
+
+    This is used by late-interaction scoring where each query/document
+    is encoded independently.
+    """
+    prompts: list[PromptType] = []
+    for data in data_list:
+        if isinstance(data, str):
+            prompts.append(data)
+        else:
+            text, mm_data, mm_uuids = parse_score_data_single(data, role, model_config)
+            prompt: TextPrompt = TextPrompt(prompt=text)
+            if mm_data is not None:
+                prompt["multi_modal_data"] = mm_data
+            if mm_uuids is not None:
+                prompt["multi_modal_uuids"] = mm_uuids
+            prompts.append(prompt)
+    return prompts
+
+
 def _parse_score_content(
     role: str,
     data: ScoreData,

vllm/envs.py

@@ -209,6 +209,7 @@ if TYPE_CHECKING:
     VLLM_GPT_OSS_SYSTEM_TOOL_MCP_LABELS: set[str] = set()
     VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT: bool = False
     VLLM_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS: bool = False
+    VLLM_SYSTEM_START_DATE: str | None = None
     VLLM_TOOL_JSON_ERROR_AUTOMATIC_RETRY: bool = False
     VLLM_CUSTOM_SCOPES_FOR_PROFILING: bool = False
     VLLM_NVTX_SCOPES_FOR_PROFILING: bool = False
@@ -1458,6 +1459,12 @@ environment_variables: dict[str, Callable[[], Any]] = {
     "VLLM_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS": lambda: bool(
         int(os.getenv("VLLM_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS", "0"))
     ),
+    # Pin the conversation start date injected into the Harmony system
+    # message. When unset the current date is used, which introduces
+    # non-determinism (different tokens -> different model behaviour at
+    # temperature=0). Set to an ISO date string, e.g. "2023-09-12",
+    # for reproducible inference or testing.
+    "VLLM_SYSTEM_START_DATE": lambda: os.getenv("VLLM_SYSTEM_START_DATE", None),
     # Enable automatic retry when tool call JSON parsing fails
     # If enabled, returns an error message to the model to retry
     # If disabled (default), raises an exception and fails the request

vllm/kernels/helion/config_manager.py

@@ -71,10 +71,18 @@ class ConfigSet:
         platform_dict = self._configs.get(platform)
         if platform_dict is None:
             avail_platforms = self.get_platforms()
+            # TODO(@gmagogsfm): add a CLI/env override flag so users can
+            # directly specify a platform name instead of relying on
+            # auto-detection, and suggest it in this error message.
             raise KeyError(
                 f"Config not found for kernel '{self._kernel_name}': "
                 f"platform '{platform}' not found. "
-                f"Available platforms: {avail_platforms or '(none)'}"
+                f"Available platforms: {avail_platforms or '(none)'}. "
+                f"If your GPU is a variant of a supported platform, "
+                f"consider adding a mapping in _GPU_NAME_ALIASES in "
+                f"vllm/kernels/helion/utils.py, or run "
+                f"scripts/autotune_helion_kernels.py to generate configs "
+                f"for your platform."
             )
 
         config = platform_dict.get(config_key)

vllm/kernels/helion/ops/silu_mul_fp8.py

@@ -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:

vllm/kernels/helion/utils.py

@@ -8,6 +8,44 @@ from vllm.platforms import current_platform
 
 logger = logging.getLogger(__name__)
 
+# Maps known variant GPU names (after lowercase/underscore normalization)
+# to their canonical form.
+#
+# Names that are already canonical after normalization are NOT listed here.
+# For example, "NVIDIA H200" normalizes to "nvidia_h200" which needs no
+# further mapping, and AMD ROCm names like "AMD_Instinct_MI300X" come from
+# a controlled lookup table in rocm.py and normalize cleanly to
+# "amd_instinct_mi300x". Only names with variant suffixes (form factor,
+# memory size, memory type, etc.) that should be stripped need entries.
+#
+# To add a new GPU variant: run `canonicalize_gpu_name()` without the alias
+# to see the normalized name, then add a mapping here if it contains variant
+# suffixes that should be stripped (e.g. Blackwell/Rubin variants).
+_GPU_NAME_ALIASES: dict[str, str] = {
+    # H100 variants
+    "nvidia_h100_pcie": "nvidia_h100",
+    "nvidia_h100_sxm5": "nvidia_h100",
+    "nvidia_h100_80gb_hbm3": "nvidia_h100",
+    "nvidia_h100_nvl": "nvidia_h100",
+    # H200 variants
+    "nvidia_h200_nvl": "nvidia_h200",
+    "nvidia_h200_141gb_hbm3e": "nvidia_h200",
+    # A100 variants
+    "nvidia_a100_sxm4_80gb": "nvidia_a100",
+    "nvidia_a100_sxm4_40gb": "nvidia_a100",
+    "nvidia_a100_pcie_80gb": "nvidia_a100",
+    "nvidia_a100_pcie_40gb": "nvidia_a100",
+    "nvidia_a100_80gb_pcie": "nvidia_a100",
+    # V100 variants (Tesla-branded)
+    "tesla_v100_sxm2_32gb": "tesla_v100",
+    "tesla_v100_sxm2_16gb": "tesla_v100",
+    "tesla_v100_pcie_32gb": "tesla_v100",
+    "tesla_v100_pcie_16gb": "tesla_v100",
+    # AMD ROCm variants (from _ROCM_DEVICE_ID_NAME_MAP in rocm.py)
+    "amd_instinct_mi300x_hf": "amd_instinct_mi300x",
+    # ADD MORE HERE
+}
+
 
 def get_gpu_name(device_id: int | None = None) -> str:
     if device_id is None:
@@ -23,17 +61,19 @@ def canonicalize_gpu_name(name: str) -> str:
     """
     Canonicalize GPU name for use as a platform identifier.
 
-    Converts to lowercase and replaces spaces and hyphens with underscores.
-    e.g., "NVIDIA A100-SXM4-80GB" -> "nvidia_a100_sxm4_80gb"
-          "AMD_Instinct_MI300X"   -> "amd_instinct_mi300x"
-
-    Raises ValueError if name is empty.
+    Converts to lowercase, replaces spaces and hyphens with underscores,
+    and maps known variant names to their canonical form via _GPU_NAME_ALIASES.
+    e.g., "NVIDIA H100 80GB HBM3" -> "nvidia_h100"
+          "NVIDIA A100-SXM4-80GB" -> "nvidia_a100"
+          "AMD Instinct MI300X"   -> "amd_instinct_mi300x"
     """
     if not name or not name.strip():
         raise ValueError("GPU name cannot be empty")
     name = name.lower()
     name = name.replace(" ", "_")
     name = name.replace("-", "_")
+    if name in _GPU_NAME_ALIASES:
+        return _GPU_NAME_ALIASES[name]
     return name
 
 

vllm/lora/ops/triton_ops/__init__.py

@@ -2,6 +2,11 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 
+from vllm.lora.ops.triton_ops.fused_moe_lora_fp8_op import (
+    fused_moe_lora_expand_fp8,
+    fused_moe_lora_fp8,
+    fused_moe_lora_shrink_fp8,
+)
 from vllm.lora.ops.triton_ops.fused_moe_lora_op import (
     fused_moe_lora,
     fused_moe_lora_expand,
@@ -18,4 +23,7 @@ __all__ = [
     "fused_moe_lora",
     "fused_moe_lora_shrink",
     "fused_moe_lora_expand",
+    "fused_moe_lora_fp8",
+    "fused_moe_lora_shrink_fp8",
+    "fused_moe_lora_expand_fp8",
 ]

vllm/model_executor/custom_op.py

@@ -1,5 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import functools
+import inspect
+
 import torch
 import torch.nn as nn
 
@@ -205,9 +208,9 @@ class CustomOp(nn.Module):
         NOTE: this does not enable fusion across ops, so opaque custom ops
         should still be unwrapped wherever possible.
         """
-        # Do not compile if compilation disabled
         from vllm.config.compilation import CompilationMode
 
+        # Do not compile if compilation disabled
         if not enable:
             return fn
 
@@ -220,14 +223,42 @@ class CustomOp(nn.Module):
         if compilation_config.backend == "eager":
             return fn
 
+        compile_options = maybe_disable_graph_partition(
+            current_platform.simple_compile_backend
+        )
+        backend = current_platform.simple_compile_backend
+
+        dynamic_arg_dims = getattr(self.__class__, "_dynamic_arg_dims", None)
+        if dynamic_arg_dims is not None:
+            compiled_fn = torch.compile(
+                fn,
+                dynamic=False,
+                backend=backend,
+                options=compile_options,
+            )
+            sig = inspect.signature(fn)
+
+            @functools.wraps(fn)
+            def wrapper(*args, **kwargs):
+                bound = sig.bind(*args, **kwargs)
+                bound.apply_defaults()
+                for name, dims in dynamic_arg_dims.items():
+                    arg = bound.arguments.get(name)
+                    if arg is not None and isinstance(arg, torch.Tensor):
+                        dims_list = [dims] if isinstance(dims, int) else dims
+                        for d in dims_list:
+                            real_d = arg.ndim + d if d < 0 else d
+                            torch._dynamo.mark_dynamic(arg, real_d)
+                return compiled_fn(*args, **kwargs)
+
+            return wrapper
+
         # dynamic=True to avoid recompilations
         return torch.compile(
             fn,
             dynamic=True,
-            backend=current_platform.simple_compile_backend,
-            options=maybe_disable_graph_partition(
-                current_platform.simple_compile_backend
-            ),
+            backend=backend,
+            options=compile_options,
         )
 
     @classmethod
@@ -267,10 +298,15 @@ class CustomOp(nn.Module):
 
     # Decorator to register custom ops.
     @classmethod
-    def register(cls, name: str):
+    def register(
+        cls,
+        name: str,
+        dynamic_arg_dims: dict[str, int | list[int]] | None = None,
+    ):
         def decorator(op_cls):
             assert name not in op_registry, f"Duplicate op name: {name}"
             op_cls.name = name
+            op_cls._dynamic_arg_dims = dynamic_arg_dims
             op_registry[name] = op_cls
             return op_cls
 

vllm/model_executor/layers/attention/mla_attention.py

@@ -950,7 +950,10 @@ def dynamic_per_batched_tensor_quant(
 logger = init_logger(__name__)
 
 
-@CustomOp.register("mla_decode_concat_quant_fp8")
+@CustomOp.register(
+    "mla_decode_concat_quant_fp8",
+    dynamic_arg_dims={"decode_ql_nope": 0, "decode_q_pe": 0},
+)
 class _DecodeConcatQuantFP8(QuantFP8):
     """
     QuantFP8 variant that concatenates decode_ql_nope and decode_q_pe before

vllm/model_executor/layers/fla/ops/chunk.py

@@ -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