[CI][NIXL] Fix PD CI breakage: pin nixl-cu{12,13} versions (#39851)

Signed-off-by: ZhanqiuHu <zhu@redhat.com>
(cherry picked from commit 799973af4e)

Signed-off-by: khluu <khluu000@gmail.com>

.buildkite/scripts/hardware_ci/run-cpu-distributed-smoke-test.sh

@@ -1,9 +1,10 @@
 #!/bin/bash
 set -euox pipefail
 export VLLM_CPU_CI_ENV=0
+export VLLM_CPU_KVCACHE_SPACE=1 # avoid OOM
 
 echo "--- PP+TP"
-vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -pp=2 &
+vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -pp=2 --max-model-len=4096 &
 server_pid=$!
 timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1
 vllm bench serve \
@@ -23,7 +24,7 @@ if [ "$failed_req" -ne 0 ]; then
 fi
 
 echo "--- DP+TP"
-vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -dp=2 &
+vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -dp=2 --max-model-len=4096 &
 server_pid=$!
 timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1
 vllm bench serve \

.buildkite/scripts/hardware_ci/run-cpu-test.sh

@@ -16,5 +16,5 @@ echo "--- :docker: Building Docker image"
 docker build --progress plain --tag "$IMAGE_NAME" --target vllm-test -f docker/Dockerfile.cpu .
 
 # Run the image, setting --shm-size=4g for tensor parallel.
-docker run --rm --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN -e VLLM_CPU_KVCACHE_SPACE=16 -e VLLM_CPU_CI_ENV=1 -e VLLM_CPU_SIM_MULTI_NUMA=1 --shm-size=4g "$IMAGE_NAME" \
+docker run --rm --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN -e VLLM_CPU_KVCACHE_SPACE=16 -e VLLM_CPU_CI_ENV=1 -e VLLM_CPU_SIM_MULTI_NUMA=1 -e VLLM_CPU_ATTN_SPLIT_KV=0 --shm-size=4g "$IMAGE_NAME" \
         timeout "$TIMEOUT_VAL" bash -c "set -euox pipefail; echo \"--- Print packages\"; pip list; echo \"--- Running tests\"; ${TEST_COMMAND}"

.buildkite/test_areas/kernels.yaml

@@ -10,7 +10,20 @@ steps:
   - tests/kernels/test_top_k_per_row.py
   - tests/kernels/test_concat_mla_q.py
   commands:
-    - pytest -v -s kernels/core kernels/test_top_k_per_row.py kernels/test_concat_mla_q.py
+    - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_top_k_per_row.py kernels/test_concat_mla_q.py
+
+- label: Kernels MiniMax Reduce RMS Test (2 GPUs)
+  timeout_in_minutes: 15
+  num_devices: 2
+  device: h100
+  source_file_dependencies:
+  - csrc/minimax_reduce_rms_kernel.cu
+  - csrc/minimax_reduce_rms_kernel.h
+  - vllm/model_executor/layers/mamba/linear_attn.py
+  - vllm/model_executor/layers/mamba/lamport_workspace.py
+  - tests/kernels/core/test_minimax_reduce_rms.py
+  commands:
+    - pytest -v -s kernels/core/test_minimax_reduce_rms.py
 
 - label: Kernels Attention Test %N
   timeout_in_minutes: 35

.buildkite/test_areas/models_basic.yaml

@@ -69,3 +69,18 @@ steps:
     - python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
     # Whisper needs spawn method to avoid deadlock
     - VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper
+
+- label: Transformers Backward Compatibility Models Test
+  working_dir: "/vllm-workspace/"
+  optional: true
+  soft_fail: true
+  commands:
+    - pip install transformers==4.57.5
+    - pytest -v -s tests/models/test_initialization.py
+    - pytest -v -s tests/models/test_transformers.py
+    - pytest -v -s tests/models/multimodal/processing/
+    - pytest -v -s tests/models/multimodal/test_mapping.py
+    - python3 examples/offline_inference/basic/chat.py
+    - python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
+    # Whisper needs spawn method to avoid deadlock
+    - VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper

CMakeLists.txt

@@ -306,6 +306,8 @@ set(VLLM_EXT_SRC
   "csrc/torch_bindings.cpp")
 
 if(VLLM_GPU_LANG STREQUAL "CUDA")
+  list(APPEND VLLM_EXT_SRC "csrc/minimax_reduce_rms_kernel.cu")
+
   SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
 
   # Set CUTLASS_REVISION. Used for FetchContent. Also fixes some bogus messages when building.

csrc/minimax_reduce_rms_kernel.cu

@@ -0,0 +1,879 @@
+
+/*
+ * Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <cooperative_groups.h>
+#include <cuda_runtime.h>
+
+#include <torch/cuda.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+
+#include "cuda_compat.h"
+#include "cuda_utils.h"
+#include "core/registration.h"
+#include "minimax_reduce_rms_kernel.h"
+
+#include <algorithm>
+
+#define FINAL_MASK 0xffffffff
+#define MINIMAX_REDUCE_RMS_WARP_SIZE 32
+
+namespace vllm {
+namespace tensorrt_llm {
+
+template <int NRanks>
+struct LamportComm {
+  __device__ __forceinline__ LamportComm(void** workspace, int rank) {
+    counter_ptr = &reinterpret_cast<int*>(workspace[NRanks * 3])[0];
+    flag_ptr = &reinterpret_cast<int*>(workspace[NRanks * 3])[2];
+    clear_ptr = &reinterpret_cast<int64_t*>(workspace[NRanks * 3 + 1])[0];
+    flag_value = *flag_ptr;
+    auto comm_size = reinterpret_cast<int64_t*>(workspace[NRanks * 3 + 1])[1];
+    clear_size = *clear_ptr;
+    int data_offset = flag_value % 3;
+    int clear_offset = (flag_value + 2) % 3;
+    for (int r = 0; r < NRanks; ++r) {
+      data_bufs[r] = reinterpret_cast<uint8_t*>(workspace[2 * NRanks + r]) +
+                     data_offset * comm_size;
+    }
+    clear_buf = reinterpret_cast<uint8_t*>(workspace[2 * NRanks + rank]) +
+                clear_offset * comm_size;
+    __syncthreads();
+    if (threadIdx.x == 0) {
+      atomicAdd(counter_ptr, 1);
+    }
+  }
+
+  __device__ __forceinline__ void update(int64_t new_clear_size) {
+    if (blockIdx.x == 0 && threadIdx.x == 0) {
+      while (*reinterpret_cast<int volatile*>(counter_ptr) != gridDim.x) {
+      }
+      *flag_ptr = (flag_value + 1) % 3;
+      *clear_ptr = new_clear_size;
+      *counter_ptr = 0;
+    }
+  }
+
+  int* counter_ptr;
+  int* flag_ptr;
+  int64_t* clear_ptr;
+  uint8_t* data_bufs[NRanks];
+  uint8_t* clear_buf;
+  int64_t clear_size;
+  int flag_value;
+};
+
+__device__ __forceinline__ bool is_neg_zero(float v) {
+  return *reinterpret_cast<uint32_t*>(&v) == 0x80000000;
+}
+
+__device__ __forceinline__ bool is_neg_zero(float4 v) {
+  return is_neg_zero(v.x) || is_neg_zero(v.y) || is_neg_zero(v.z) ||
+         is_neg_zero(v.w);
+}
+
+__device__ __forceinline__ float4 get_neg_zero() {
+  float4 vec;
+#pragma unroll
+  for (int i = 0; i < 4; ++i) {
+    reinterpret_cast<uint32_t*>(&vec)[i] = 0x80000000;
+  }
+  return vec;
+}
+
+template <int Dim>
+__device__ __forceinline__ float rms_rsqrt(float& v, float eps) {
+  constexpr float kInvDim = 1.0F / static_cast<float>(Dim);
+  v = rsqrtf((v * kInvDim) + eps);
+  return v;
+}
+
+template <int Dim>
+__device__ __forceinline__ float4 rms_rsqrt(float4& v, float eps) {
+  constexpr float kInvDim = 1.0F / static_cast<float>(Dim);
+  v.x = rsqrtf((v.x * kInvDim) + eps);
+  v.y = rsqrtf((v.y * kInvDim) + eps);
+  v.z = rsqrtf((v.z * kInvDim) + eps);
+  v.w = rsqrtf((v.w * kInvDim) + eps);
+  return v;
+}
+__device__ __forceinline__ float4 ld_global_volatile(float4* addr) {
+  float4 val;
+  asm volatile("ld.volatile.global.v4.f32 {%0, %1, %2, %3}, [%4];"
+               : "=f"(val.x), "=f"(val.y), "=f"(val.z), "=f"(val.w)
+               : "l"(addr));
+  return val;
+}
+
+__device__ __forceinline__ float ld_global_volatile(float* addr) {
+  float val;
+  asm volatile("ld.volatile.global.f32 %0, [%1];" : "=f"(val) : "l"(addr));
+  return val;
+}
+
+// Used by the scalar (non-float4) kernel only
+template <typename T, int NUM>
+__inline__ __device__ T warpReduceSumV2(T* val) {
+#pragma unroll
+  for (int i = 0; i < NUM; i++) {
+#pragma unroll
+    for (int mask = 16; mask > 0; mask >>= 1)
+      val[i] += __shfl_xor_sync(FINAL_MASK, val[i], mask, 32);
+  }
+  return (T)(0.0f);
+}
+
+template <typename T, int NUM>
+__inline__ __device__ T blockReduceSumV2(T* val) {
+  static __shared__ T shared[NUM][33];
+  int lane = threadIdx.x & 0x1f;
+  int wid = threadIdx.x >> 5;
+
+  warpReduceSumV2<T, NUM>(val);
+
+  if (lane == 0) {
+#pragma unroll
+    for (int i = 0; i < NUM; i++) {
+      shared[i][wid] = val[i];
+    }
+  }
+
+  __syncthreads();
+
+  bool is_mask = threadIdx.x < (blockDim.x / 32.f);
+#pragma unroll
+  for (int i = 0; i < NUM; i++) {
+    val[i] = is_mask ? shared[i][lane] : (T)(0.0f);
+  }
+  warpReduceSumV2<T, NUM>(val);
+  return (T)0.0f;
+}
+
+// for float4 version
+template <uint32_t kNumThreads, typename T, int ArraySize = 4>
+__device__ __forceinline__ void local_warp_reduce_sum_array(
+    T* value_ptr, uint32_t active_mask = 0xffffffffu) {
+  static_assert(kNumThreads >= 1 &&
+                kNumThreads <= MINIMAX_REDUCE_RMS_WARP_SIZE);
+#pragma unroll
+  for (int i = 0; i < ArraySize; ++i) {
+#pragma unroll
+    for (int mask = kNumThreads / 2; mask > 0; mask >>= 1) {
+      value_ptr[i] += __shfl_xor_sync(active_mask, value_ptr[i], mask,
+                                      MINIMAX_REDUCE_RMS_WARP_SIZE);
+    }
+  }
+}
+
+constexpr int next_pow2(int val) {
+  int result = 1;
+  while (result < val) {
+    result <<= 1;
+  }
+  return result;
+}
+
+// ---------------------------------------------------------------------------
+
+template <typename DType>
+class IndexHelper {
+ public:
+  __device__ __forceinline__ IndexHelper(MiniMaxReduceRMSParams const& params) {
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
+    namespace cg = cooperative_groups;
+    cg::cluster_group cluster = cg::this_cluster();
+    cg::grid_group grid = cg::this_grid();
+    token_id = grid.cluster_rank();
+    access_id_in_token = cluster.thread_rank();
+    token_stride = grid.num_clusters();
+#else
+    token_id = blockIdx.x;
+    access_id_in_token = threadIdx.x;
+    token_stride = gridDim.x;
+#endif
+    access_id = token_id * params.hidden_dim / kElemsPerAccess<DType> +
+                access_id_in_token;
+    access_stride = token_stride * params.hidden_dim / kElemsPerAccess<DType>;
+    tot_access = params.size_q / kElemsPerAccess<DType>;
+  }
+
+  int token_id;
+  int access_id_in_token;
+  int token_stride;
+  int access_id;
+  int access_stride;
+  int tot_access;
+};
+
+/**
+* this kernel is used to for minimax attention module
+* input tensor [total_tokens, hidden_dim / tp_size], fp32
+* rms weight [hidden_dim / tp_size], bf16
+step 1: reduce from single rank to get the variance sum (reduce(input^2,
+dim=-1)) step 2: reduce from all ranks to get the variance sum
+(all_reduce(variance_sum)) step 3: calculate the rms norm (input *
+rsqrt(variance + eps)) in this case, max hidden_dim is 6144 (float data), for
+each token, we only need 6144 / 4 / tp_size = (1536 / tp_size) threads so we can
+assume cluster size is 1 (tp_size >= 2)
+ */
+template <typename DType, int NRanks>
+__global__ void __launch_bounds__(1024)
+    minimax_reduce_rms_kernel_lamport(MiniMaxReduceRMSParams params) {
+  IndexHelper<DType> index_helper(params);
+  int token_id = index_helper.token_id;
+  int access_id_in_token = index_helper.access_id_in_token;
+  int token_stride = index_helper.token_stride;
+  int access_id = index_helper.access_id;
+  int access_stride = index_helper.access_stride;
+  int tot_access = index_helper.tot_access;
+  int tot_tokens = params.size_q / params.hidden_dim;
+  float4 clear_vec = get_neg_zero();
+
+  LamportComm<NRanks> comm(params.workspace, params.rank);
+  int clear_access = comm.clear_size / kElemsPerAccess<DType>;
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
+  asm volatile("griddepcontrol.wait;");
+#endif
+  for (int idx = access_id; idx < tot_access;
+       idx += access_stride, token_id += token_stride) {
+    alignas(16) DType vals[kElemsPerAccess<DType>];
+    float sum_variance = 0.F;
+    *reinterpret_cast<float4*>(vals) =
+        reinterpret_cast<float4*>(params.allreduce_in)[idx];
+#pragma unroll
+    for (int i = 0; i < kElemsPerAccess<DType>; ++i) {
+      sum_variance += static_cast<float>(vals[i]) * static_cast<float>(vals[i]);
+    }
+    blockReduceSumV2<float, 1>(&sum_variance);
+    if (is_neg_zero(sum_variance)) {
+      sum_variance = 0.F;
+    }
+    if (threadIdx.x == 0) {
+      for (int r = 0; r < NRanks; ++r) {
+        reinterpret_cast<float*>(
+            comm.data_bufs[r])[(params.rank * tot_tokens) + token_id] =
+            (sum_variance);
+      }
+    }
+
+    bool done = false;
+    float vars_all_ranks[NRanks];
+    while (!done) {
+      done = true;
+#pragma unroll
+      for (int r = 0; r < NRanks; ++r) {
+        vars_all_ranks[r] = ld_global_volatile(&reinterpret_cast<float*>(
+            comm.data_bufs[params.rank])[(r * tot_tokens) + token_id]);
+        done &= !is_neg_zero(vars_all_ranks[r]);
+      }
+    }
+    sum_variance = 0.F;
+#pragma unroll
+    for (int r = 0; r < NRanks; ++r) {
+      sum_variance += vars_all_ranks[r];
+    }
+
+    DType norm_weight[kElemsPerAccess<DType>];
+    *reinterpret_cast<typename ElemsPerAccess<DType>::vec_type*>(norm_weight) =
+        reinterpret_cast<typename ElemsPerAccess<DType>::vec_type*>(
+            params.rms_gamma)[access_id_in_token];
+
+#pragma unroll
+    for (int i = 0; i < kElemsPerAccess<DType>; ++i) {
+      vals[i] = static_cast<DType>(
+          static_cast<float>(vals[i]) *
+          rsqrtf(
+              (sum_variance / static_cast<float>(params.hidden_dim) / NRanks) +
+              params.rms_eps) *
+          static_cast<float>(norm_weight[i]));
+    }
+
+    reinterpret_cast<float4*>(params.rms_norm_out)[idx] =
+        *reinterpret_cast<float4*>(vals);
+  }
+  for (int idx = access_id; idx < clear_access; idx += access_stride) {
+    reinterpret_cast<float4*>(comm.clear_buf)[idx] = clear_vec;
+  }
+  comm.update(params.size_q * NRanks);
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
+  asm volatile("griddepcontrol.launch_dependents;");
+#endif
+}
+
+/**
+ * Float4 variant: process 4 rows at once, allreduce variance sums as float4 for
+ * better memory coalescing. sum_variance is always float; applies to all DTypes
+ * (half, bf16, float). When tot_tokens % 4 != 0, the last group pads rows with
+ * zeros; padded rows are not written to rms_norm_out. IsQK: when true, process
+ * Q+K in one loop with doubled comm buffer; when false, single-matrix (Q only).
+ */
+template <typename DType, int NRanks, int OriginQDim, int OriginKDim>
+__global__ void __launch_bounds__(1024)
+    minimax_reduce_qk_rms_kernel_lamport_float4(MiniMaxReduceRMSParams params) {
+  // Compile-time per-rank dimensions
+  constexpr int RankQDim = OriginQDim / NRanks;
+  constexpr int RankKDim = OriginKDim / NRanks;
+  // Threads needed to cover one row of Q / K with float4 accesses
+  constexpr int ThreadsPerRowQ = RankQDim / kElemsPerAccess<DType>;
+  constexpr int ThreadsPerRowK = RankKDim / kElemsPerAccess<DType>;
+  // Number of warps dedicated to Q / K
+  constexpr int NumWarpQ = (ThreadsPerRowQ + MINIMAX_REDUCE_RMS_WARP_SIZE - 1) /
+                           MINIMAX_REDUCE_RMS_WARP_SIZE;
+  constexpr int NumWarpK = (ThreadsPerRowK + MINIMAX_REDUCE_RMS_WARP_SIZE - 1) /
+                           MINIMAX_REDUCE_RMS_WARP_SIZE;
+
+  int tot_tokens = params.size_q / RankQDim;
+  int tot_groups = (tot_tokens + 3) / 4;  // ceiling; last group may be partial
+
+  // Memory strides for strided qkv tensors (elements -> float4-access units)
+  int access_stride_q = (params.stride_q > 0 ? params.stride_q : RankQDim) /
+                        kElemsPerAccess<DType>;
+  int access_stride_k = (params.stride_k > 0 ? params.stride_k : RankKDim) /
+                        kElemsPerAccess<DType>;
+  // Output strides: default to contiguous (hidden_dim / hidden_dim_k)
+  int access_stride_q_out =
+      (params.stride_q_out > 0 ? params.stride_q_out : params.hidden_dim) /
+      kElemsPerAccess<DType>;
+  int access_stride_k_out =
+      (params.stride_k_out > 0 ? params.stride_k_out : params.hidden_dim_k) /
+      kElemsPerAccess<DType>;
+
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
+  namespace cg = cooperative_groups;
+  cg::cluster_group cluster = cg::this_cluster();
+  cg::grid_group grid = cg::this_grid();
+  int group_id = grid.cluster_rank();
+  int access_id_in_token = cluster.thread_rank();
+  int group_stride = grid.num_clusters();
+#else
+  int group_id = blockIdx.x;
+  int access_id_in_token = threadIdx.x;
+  int group_stride = gridDim.x;
+#endif
+
+  bool is_q = (access_id_in_token < NumWarpQ * MINIMAX_REDUCE_RMS_WARP_SIZE);
+  int k_thread_idx =
+      access_id_in_token - (NumWarpQ * MINIMAX_REDUCE_RMS_WARP_SIZE);
+  bool is_valid_q = (access_id_in_token < ThreadsPerRowQ);
+  bool is_valid_k = (k_thread_idx >= 0 && k_thread_idx < ThreadsPerRowK);
+  float4 clear_vec = get_neg_zero();
+
+  // Shared memory for two-level block reduction and scale broadcast
+  __shared__ float block_reduce_sum[4][MINIMAX_REDUCE_RMS_WARP_SIZE + 1];
+  __shared__ float global_scale_q[4];
+  __shared__ float global_scale_k[4];
+
+  LamportComm<NRanks> comm(params.workspace, params.rank);
+
+  DType norm_weight[kElemsPerAccess<DType>]{};
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
+  asm volatile("griddepcontrol.wait;");
+#endif
+  if (is_q) {
+    if (is_valid_q) {
+      *reinterpret_cast<typename ElemsPerAccess<DType>::vec_type*>(
+          norm_weight) =
+          reinterpret_cast<typename ElemsPerAccess<DType>::vec_type const*>(
+              params.rms_gamma)[access_id_in_token];
+    }
+  } else {
+    if (is_valid_k) {
+      *reinterpret_cast<typename ElemsPerAccess<DType>::vec_type*>(
+          norm_weight) =
+          reinterpret_cast<typename ElemsPerAccess<DType>::vec_type const*>(
+              params.rms_gamma_k)[k_thread_idx];
+    }
+  }
+
+  // Main loop: process one group of 4 tokens per iteration.
+  for (int g = group_id; g < tot_groups; g += group_stride) {
+    alignas(16) DType vals[4][kElemsPerAccess<DType>]{};
+    float warp_sum_variance[4]{0.F, 0.F, 0.F, 0.F};
+
+    if (is_q) {
+#pragma unroll
+      for (int row = 0; row < 4; ++row) {
+        int token_r = g * 4 + row;
+        if (token_r >= tot_tokens || !is_valid_q) {
+          continue;
+        }
+        int idx_r = token_r * access_stride_q + access_id_in_token;
+        *reinterpret_cast<float4*>(&vals[row][0]) =
+            reinterpret_cast<float4 const*>(params.allreduce_in)[idx_r];
+#pragma unroll
+        for (int i = 0; i < kElemsPerAccess<DType>; ++i) {
+          float x = static_cast<float>(vals[row][i]);
+          warp_sum_variance[row] += x * x;
+        }
+      }
+    } else {
+#pragma unroll
+      for (int row = 0; row < 4; ++row) {
+        int token_r = g * 4 + row;
+        if (token_r >= tot_tokens || !is_valid_k) {
+          continue;
+        }
+        int idx_r = token_r * access_stride_k + k_thread_idx;
+        *reinterpret_cast<float4*>(&vals[row][0]) =
+            reinterpret_cast<float4 const*>(params.allreduce_in_k)[idx_r];
+#pragma unroll
+        for (int i = 0; i < kElemsPerAccess<DType>; ++i) {
+          float x = static_cast<float>(vals[row][i]);
+          warp_sum_variance[row] += x * x;
+        }
+      }
+    }
+
+    local_warp_reduce_sum_array<MINIMAX_REDUCE_RMS_WARP_SIZE, float, 4>(
+        warp_sum_variance);
+    // Warp lane 0 writes its warp's partial sum to shared memory
+    int lane = threadIdx.x & (MINIMAX_REDUCE_RMS_WARP_SIZE - 1);
+    if (lane == 0) {
+#pragma unroll
+      for (int t = 0; t < 4; ++t) {
+        block_reduce_sum[t][threadIdx.x / MINIMAX_REDUCE_RMS_WARP_SIZE] =
+            warp_sum_variance[t];
+      }
+    }
+    __syncthreads();
+
+    int tid = threadIdx.x;
+
+    if (tid < MINIMAX_REDUCE_RMS_WARP_SIZE) {
+      constexpr int kNumWarpQPow2 =
+          (next_pow2(NumWarpQ) > NRanks) ? next_pow2(NumWarpQ) : NRanks;
+      float local_sum[4];
+#pragma unroll
+      for (int t = 0; t < 4; ++t) {
+        local_sum[t] = (tid < NumWarpQ) ? block_reduce_sum[t][tid] : 0.F;
+      }
+      // After this, all kNumWarpQPow2 lanes (including tid 0..NRanks-1) have
+      // the total Q sum-of-squares for all 4 tokens.
+      local_warp_reduce_sum_array<kNumWarpQPow2, float, 4>(local_sum);
+
+      if (tid < NRanks) {
+#pragma unroll
+        for (int t = 0; t < 4; ++t) {
+          if (is_neg_zero(local_sum[t])) {
+            local_sum[t] = 0.F;
+          }
+        }
+        // Parallel push: thread tid writes this rank's Q sum to rank tid's buf
+        reinterpret_cast<float4*>(
+            comm.data_bufs[tid])[(params.rank * tot_groups * 2) + (2 * g)] =
+            *reinterpret_cast<float4*>(local_sum);
+
+        // Parallel pull: thread tid reads rank tid's contribution from
+        // this rank's (params.rank's) buffer
+        bool done = false;
+        float4 var_all_ranks;
+        while (!done) {
+          done = true;
+          var_all_ranks = ld_global_volatile(&reinterpret_cast<float4*>(
+              comm.data_bufs[params.rank])[(tid * tot_groups * 2) + (2 * g)]);
+          done &= !is_neg_zero(var_all_ranks);
+        }
+
+        // Warp-level allreduce: each of the NRanks threads holds one rank's
+        // partial sum; after this all NRanks threads have the global total.
+        constexpr uint32_t kQActiveMask = (1u << NRanks) - 1u;
+        local_warp_reduce_sum_array<NRanks, float, 4>(
+            reinterpret_cast<float*>(&var_all_ranks), kQActiveMask);
+
+        // Thread 0 computes rsqrt with compile-time Dim and writes to smem
+        if (tid == 0) {
+          *reinterpret_cast<float4*>(global_scale_q) =
+              rms_rsqrt<OriginQDim>(var_all_ranks, params.rms_eps);
+        }
+      }
+    } else if (tid >= MINIMAX_REDUCE_RMS_WARP_SIZE * NumWarpQ &&
+               tid < MINIMAX_REDUCE_RMS_WARP_SIZE * (NumWarpQ + 1)) {
+      // --- K leader warp ---
+      constexpr int kNumWarpKPow2 =
+          (next_pow2(NumWarpK) > NRanks) ? next_pow2(NumWarpK) : NRanks;
+      float local_sum[4];
+#pragma unroll
+      for (int t = 0; t < 4; ++t) {
+        local_sum[t] = (k_thread_idx < NumWarpK)
+                           ? block_reduce_sum[t][NumWarpQ + k_thread_idx]
+                           : 0.F;
+      }
+      local_warp_reduce_sum_array<kNumWarpKPow2, float, 4>(local_sum);
+
+      if (k_thread_idx < NRanks) {
+#pragma unroll
+        for (int t = 0; t < 4; ++t) {
+          if (is_neg_zero(local_sum[t])) {
+            local_sum[t] = 0.F;
+          }
+        }
+        reinterpret_cast<float4*>(
+            comm.data_bufs[k_thread_idx])[(params.rank * tot_groups * 2) +
+                                          (2 * g + 1)] =
+            *reinterpret_cast<float4*>(local_sum);
+
+        bool done = false;
+        float4 var_all_ranks;
+        while (!done) {
+          done = true;
+          var_all_ranks = ld_global_volatile(&reinterpret_cast<float4*>(
+              comm.data_bufs[params.rank])[(k_thread_idx * tot_groups * 2) +
+                                           (2 * g + 1)]);
+          done &= !is_neg_zero(var_all_ranks);
+        }
+
+        constexpr uint32_t kKActiveMask = (1u << NRanks) - 1u;
+        local_warp_reduce_sum_array<NRanks, float, 4>(
+            reinterpret_cast<float*>(&var_all_ranks), kKActiveMask);
+
+        if (k_thread_idx == 0) {
+          *reinterpret_cast<float4*>(global_scale_k) =
+              rms_rsqrt<OriginKDim>(var_all_ranks, params.rms_eps);
+        }
+      }
+    }
+    __syncthreads();
+
+    if (is_q) {
+#pragma unroll
+      for (int t = 0; t < 4; ++t) {
+        warp_sum_variance[t] = global_scale_q[t];
+      }
+#pragma unroll
+      for (int r = 0; r < 4; ++r) {
+#pragma unroll
+        for (int i = 0; i < kElemsPerAccess<DType>; ++i) {
+          vals[r][i] = static_cast<DType>(static_cast<float>(vals[r][i]) *
+                                          warp_sum_variance[r] *
+                                          static_cast<float>(norm_weight[i]));
+        }
+        int token_r = g * 4 + r;
+        if (token_r >= tot_tokens || !is_valid_q) {
+          continue;
+        }
+        int idx_out = token_r * access_stride_q_out + access_id_in_token;
+        reinterpret_cast<float4*>(params.rms_norm_out)[idx_out] =
+            *reinterpret_cast<float4*>(&vals[r][0]);
+      }
+    } else {
+#pragma unroll
+      for (int t = 0; t < 4; ++t) {
+        warp_sum_variance[t] = global_scale_k[t];
+      }
+#pragma unroll
+      for (int r = 0; r < 4; ++r) {
+#pragma unroll
+        for (int i = 0; i < kElemsPerAccess<DType>; ++i) {
+          vals[r][i] = static_cast<DType>(static_cast<float>(vals[r][i]) *
+                                          warp_sum_variance[r] *
+                                          static_cast<float>(norm_weight[i]));
+        }
+        int token_r = g * 4 + r;
+        if (token_r >= tot_tokens || !is_valid_k) {
+          continue;
+        }
+        int idx_out = token_r * access_stride_k_out + k_thread_idx;
+        reinterpret_cast<float4*>(params.rms_norm_out_k)[idx_out] =
+            *reinterpret_cast<float4*>(&vals[r][0]);
+      }
+    }
+  }  // end group loop
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
+  asm volatile("griddepcontrol.launch_dependents;");
+#endif
+
+  int clear_access = static_cast<int>(comm.clear_size / kElemsPerAccess<DType>);
+  int clear_stride = group_stride * blockDim.x;
+  for (int idx = group_id * blockDim.x + threadIdx.x; idx < clear_access;
+       idx += clear_stride) {
+    reinterpret_cast<float4*>(comm.clear_buf)[idx] = clear_vec;
+  }
+
+  comm.update(static_cast<int64_t>(2) * tot_groups * kElemsPerAccess<DType> *
+              NRanks);
+}
+
+int get_sm_count() {
+  static int sm_count = 0;
+  if (sm_count == 0) {
+    int device_id;
+    CUDA_CHECK(cudaGetDevice(&device_id));
+    cudaDeviceProp device_prop;
+    cudaGetDeviceProperties(&device_prop, device_id);
+    sm_count = device_prop.multiProcessorCount;
+  }
+  return sm_count;
+}
+
+inline int getSMVersion(bool queryRealSmArch = false) {
+  int device{-1};
+  CUDA_CHECK(cudaGetDevice(&device));
+  int sm_major = 0;
+  int sm_minor = 0;
+  CUDA_CHECK(cudaDeviceGetAttribute(&sm_major,
+                                    cudaDevAttrComputeCapabilityMajor, device));
+  CUDA_CHECK(cudaDeviceGetAttribute(&sm_minor,
+                                    cudaDevAttrComputeCapabilityMinor, device));
+  int sm = sm_major * 10 + sm_minor;
+  if (sm == 121 && !queryRealSmArch) {
+    return 120;
+  }
+  return sm;
+}
+
+template <typename KernelFunc>
+int get_max_active_blocks(KernelFunc kernel, int block_size,
+                          int dynamic_smem = 0) {
+  int max_active = 0;
+  CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
+      &max_active, kernel, block_size, dynamic_smem));
+  return std::max(max_active, 1);
+}
+
+template <typename DType, int NRanks>
+void minimax_reduce_rms_kernel_launcher(MiniMaxReduceRMSParams const& params) {
+  static int SM = getSMVersion();
+  int token_num = params.size_q / params.hidden_dim;
+  int sm_count = get_sm_count();
+  int cluster_size = 1;
+  int cluster_num = token_num;
+  int threads_per_token = params.hidden_dim / kElemsPerAccess<DType>;
+  int block_size = threads_per_token;
+
+  int max_blocks_per_sm = get_max_active_blocks(
+      minimax_reduce_rms_kernel_lamport<DType, NRanks>, block_size);
+  int max_grid = max_blocks_per_sm * sm_count;
+
+  int grid_size =
+      (std::min(max_grid, cluster_num * cluster_size) / cluster_size) *
+      cluster_size;
+
+  cudaLaunchConfig_t cfg;
+  cfg.gridDim = grid_size;
+  cfg.blockDim = block_size;
+  cfg.dynamicSmemBytes = 0;
+  cfg.stream = params.stream;
+
+  cudaLaunchAttribute attribute[2];
+  attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+  attribute[0].val.programmaticStreamSerializationAllowed = 1;
+  attribute[1].id = cudaLaunchAttributeClusterDimension;
+  attribute[1].val.clusterDim.x = cluster_size;
+  attribute[1].val.clusterDim.y = 1;
+  attribute[1].val.clusterDim.z = 1;
+  cfg.attrs = attribute;
+  cfg.numAttrs = SM >= 90 ? 2 : 0;
+
+  CUDA_CHECK(cudaLaunchKernelEx(
+      &cfg, minimax_reduce_rms_kernel_lamport<DType, NRanks>, params));
+}
+
+template <typename DType, int NRanks, int OriginQDim, int OriginKDim>
+void minimax_reduce_rms_kernel_launcher_float4(
+    MiniMaxReduceRMSParams const& params) {
+  TORCH_CHECK(params.size_q % params.hidden_dim == 0);
+  TORCH_CHECK(params.hidden_dim % kElemsPerAccess<DType> == 0);
+  if (params.stride_q > 0) {
+    TORCH_CHECK(params.stride_q % kElemsPerAccess<DType> == 0);
+  }
+  TORCH_CHECK(params.allreduce_in_k != nullptr,
+              "float4 QK kernel requires K input");
+  TORCH_CHECK(params.hidden_dim >= params.hidden_dim_k);
+  TORCH_CHECK(params.size_k % params.hidden_dim_k == 0);
+  TORCH_CHECK(params.hidden_dim_k % kElemsPerAccess<DType> == 0);
+  TORCH_CHECK(params.size_q / params.hidden_dim ==
+              params.size_k / params.hidden_dim_k);
+  if (params.stride_k > 0) {
+    TORCH_CHECK(params.stride_k % kElemsPerAccess<DType> == 0);
+  }
+
+  int token_num = params.size_q / params.hidden_dim;
+  int tot_groups = (token_num + 3) / 4;
+  if (tot_groups == 0) {
+    return;
+  }
+
+  static int SM = getSMVersion();
+  int sm_count = get_sm_count();
+  int cluster_size = 1;
+  int cluster_num = tot_groups;
+
+  int access_per_row_q = params.hidden_dim / kElemsPerAccess<DType>;
+  int access_per_row_k = params.hidden_dim_k / kElemsPerAccess<DType>;
+
+  // Round each section up to a warp boundary
+  auto divUp = [](int a, int b) { return (a + b - 1) / b * b; };
+  int block_size = divUp(access_per_row_q, MINIMAX_REDUCE_RMS_WARP_SIZE) +
+                   divUp(access_per_row_k, MINIMAX_REDUCE_RMS_WARP_SIZE);
+
+  auto kfn =
+      minimax_reduce_qk_rms_kernel_lamport_float4<DType, NRanks, OriginQDim,
+                                                  OriginKDim>;
+
+  int max_blocks_per_sm = get_max_active_blocks(kfn, block_size);
+  int max_grid = max_blocks_per_sm * sm_count;
+  int grid_size =
+      (std::min(max_grid, cluster_num * cluster_size) / cluster_size) *
+      cluster_size;
+
+  cudaLaunchConfig_t cfg;
+  cfg.gridDim = grid_size;
+  cfg.blockDim = block_size;
+  cfg.dynamicSmemBytes = 0;
+  cfg.stream = params.stream;
+
+  cudaLaunchAttribute attribute[2];
+  attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+  attribute[0].val.programmaticStreamSerializationAllowed = 1;
+  attribute[1].id = cudaLaunchAttributeClusterDimension;
+  attribute[1].val.clusterDim.x = cluster_size;
+  attribute[1].val.clusterDim.y = 1;
+  attribute[1].val.clusterDim.z = 1;
+  cfg.attrs = attribute;
+  cfg.numAttrs = SM >= 90 ? 2 : 0;
+
+  CUDA_CHECK(cudaLaunchKernelEx(&cfg, kfn, params));
+}
+
+template <int NRanks>
+void dispatch_dtype(MiniMaxReduceRMSParams const& params) {
+  // Use the optimized QK float4 kernel when:
+  //  - K input is present, AND
+  //  - the full (NRanks * per-rank) dimensions match the MiniMax M2 shape.
+  // Otherwise fall back to the scalar kernel.
+  bool use_float4 = (params.allreduce_in_k != nullptr) &&
+                    (params.hidden_dim * params.nranks == 6144) &&
+                    (params.hidden_dim_k * params.nranks == 1024);
+
+  if (params.dtype == at::ScalarType::Half) {
+    if (use_float4) {
+      minimax_reduce_rms_kernel_launcher_float4<half, NRanks, 6144, 1024>(
+          params);
+    } else {
+      minimax_reduce_rms_kernel_launcher<half, NRanks>(params);
+    }
+  } else if (params.dtype == at::ScalarType::BFloat16) {
+    if (use_float4) {
+      minimax_reduce_rms_kernel_launcher_float4<__nv_bfloat16, NRanks, 6144,
+                                                1024>(params);
+    } else {
+      minimax_reduce_rms_kernel_launcher<__nv_bfloat16, NRanks>(params);
+    }
+  } else if (params.dtype == at::ScalarType::Float) {
+    if (use_float4) {
+      minimax_reduce_rms_kernel_launcher_float4<float, NRanks, 6144, 1024>(
+          params);
+    } else {
+      minimax_reduce_rms_kernel_launcher<float, NRanks>(params);
+    }
+  } else {
+    TORCH_CHECK(false, "Unsupported data type for minimax_reduce_rms_op");
+  }
+}
+
+void minimax_reduce_rms_op(MiniMaxReduceRMSParams const& params) {
+  if (params.nranks == 2) {
+    dispatch_dtype<2>(params);
+  } else if (params.nranks == 4) {
+    dispatch_dtype<4>(params);
+  } else if (params.nranks == 8) {
+    dispatch_dtype<8>(params);
+  } else if (params.nranks == 16) {
+    dispatch_dtype<16>(params);
+  } else {
+    TORCH_CHECK(false, "minimax_reduce_rms_op: unsupported ranks number!");
+  }
+}
+}  // namespace tensorrt_llm
+}  // namespace vllm
+
+torch::Tensor minimax_allreduce_rms(torch::Tensor const& input,
+                                    torch::Tensor const& norm_weight,
+                                    torch::Tensor workspace, int64_t const rank,
+                                    int64_t const nranks, double const eps) {
+  auto allreduce_params = vllm::tensorrt_llm::MiniMaxReduceRMSParams();
+
+  allreduce_params.nranks = static_cast<int>(nranks);
+  allreduce_params.rank = static_cast<int>(rank);
+  allreduce_params.dtype = input.scalar_type();
+  allreduce_params.size_q = static_cast<int>(input.numel());
+  allreduce_params.hidden_dim = static_cast<int>(input.size(-1));
+  allreduce_params.stride_q = allreduce_params.hidden_dim;
+  allreduce_params.workspace =
+      reinterpret_cast<void**>(workspace.mutable_data_ptr());
+  allreduce_params.allreduce_in = input.data_ptr();
+  allreduce_params.rms_gamma = norm_weight.data_ptr();
+  allreduce_params.rms_eps = static_cast<float>(eps);
+  allreduce_params.stream = at::cuda::getCurrentCUDAStream(input.get_device());
+
+  torch::Tensor rms_norm_out = torch::empty_like(input);
+  allreduce_params.rms_norm_out = rms_norm_out.mutable_data_ptr();
+
+  vllm::tensorrt_llm::minimax_reduce_rms_op(allreduce_params);
+
+  return rms_norm_out;
+}
+
+std::tuple<torch::Tensor, torch::Tensor> minimax_allreduce_rms_qk(
+    torch::Tensor qkv, torch::Tensor const& norm_weight_q,
+    torch::Tensor const& norm_weight_k, torch::Tensor workspace,
+    int64_t const q_size, int64_t const kv_size, int64_t const rank,
+    int64_t const nranks, double const eps) {
+  TORCH_CHECK(qkv.dim() == 2, "minimax_allreduce_rms_qk: qkv must be 2D");
+  TORCH_CHECK(qkv.is_contiguous(),
+              "minimax_allreduce_rms_qk: qkv must be contiguous");
+  int64_t qkv_dim = qkv.size(-1);
+  TORCH_CHECK(qkv_dim == q_size + 2 * kv_size,
+              "minimax_allreduce_rms_qk: qkv last dim must equal "
+              "q_size + 2 * kv_size");
+  TORCH_CHECK(rank < nranks,
+              "minimax_allreduce_rms_qk: rank must be less than nranks");
+
+  int64_t num_tokens = qkv.size(0);
+  int elem_bytes = qkv.element_size();
+
+  torch::Tensor q_out = torch::empty({num_tokens, q_size}, qkv.options());
+  torch::Tensor k_out = torch::empty({num_tokens, kv_size}, qkv.options());
+
+  auto params = vllm::tensorrt_llm::MiniMaxReduceRMSParams();
+  params.nranks = static_cast<int>(nranks);
+  params.rank = static_cast<int>(rank);
+  params.dtype = qkv.scalar_type();
+  params.size_q = static_cast<int>(num_tokens * q_size);
+  params.hidden_dim = static_cast<int>(q_size);
+  params.size_k = static_cast<int>(num_tokens * kv_size);
+  params.hidden_dim_k = static_cast<int>(kv_size);
+  params.stride_q = static_cast<int>(qkv_dim);
+  params.stride_k = static_cast<int>(qkv_dim);
+  params.stride_q_out = 0;  // q_out is contiguous; kernel uses hidden_dim
+  params.stride_k_out = 0;  // k_out is contiguous; kernel uses hidden_dim_k
+  params.workspace = reinterpret_cast<void**>(workspace.mutable_data_ptr());
+
+  uint8_t* base = static_cast<uint8_t*>(qkv.data_ptr());
+  params.allreduce_in = base;
+  params.allreduce_in_k = base + q_size * elem_bytes;
+  params.rms_gamma = norm_weight_q.data_ptr();
+  params.rms_gamma_k = norm_weight_k.data_ptr();
+  params.rms_eps = static_cast<float>(eps);
+  params.stream = at::cuda::getCurrentCUDAStream(qkv.get_device());
+
+  params.rms_norm_out = q_out.mutable_data_ptr();
+  params.rms_norm_out_k = k_out.mutable_data_ptr();
+
+  vllm::tensorrt_llm::minimax_reduce_rms_op(params);
+  return {q_out, k_out};
+}

csrc/minimax_reduce_rms_kernel.h

@@ -0,0 +1,79 @@
+/*
+ * Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+
+#include <torch/types.h>
+
+namespace vllm {
+namespace tensorrt_llm {
+
+template <typename DType>
+struct ElemsPerAccess;
+
+template <>
+struct ElemsPerAccess<half> {
+  static constexpr int value = 8;
+  using vec_type = float4;
+};
+
+template <>
+struct ElemsPerAccess<nv_bfloat16> {
+  static constexpr int value = 8;
+  using vec_type = float4;
+};
+
+template <>
+struct ElemsPerAccess<float> {
+  static constexpr int value = 4;
+  using vec_type = float4;
+};
+
+template <typename DType>
+static constexpr int kElemsPerAccess = ElemsPerAccess<DType>::value;
+
+struct MiniMaxReduceRMSParams {
+  int nranks{};
+  int rank{};
+  at::ScalarType dtype{at::ScalarType::Undefined};
+  int size_q{};
+  int hidden_dim{};
+  int size_k{};
+  int hidden_dim_k{};
+  int stride_q{};  // row stride for q input (elements); when > hidden_dim,
+                   // q is part of a wider qkv tensor
+  int stride_k{};  // row stride for k input (elements); when > hidden_dim_k,
+                   // k is part of a wider qkv tensor
+  int stride_q_out{};  // row stride for q output (elements); 0 = contiguous
+  int stride_k_out{};  // row stride for k output (elements); 0 = contiguous
+  void** workspace{};
+  void* allreduce_in{};
+  void* rms_norm_out{};
+  void* rms_gamma{};
+  void* allreduce_in_k{};
+  void* rms_norm_out_k{};
+  void* rms_gamma_k{};
+  float rms_eps{};
+  cudaStream_t stream{};
+};
+
+void minimax_reduce_rms_op(MiniMaxReduceRMSParams const& params);
+
+}  // namespace tensorrt_llm
+}  // namespace vllm

csrc/ops.h

@@ -391,4 +391,16 @@ int64_t qr_max_size();
 #ifndef USE_ROCM
 void dsv3_fused_a_gemm(torch::Tensor& output, torch::Tensor const& mat_a,
                        torch::Tensor const& mat_b);
+#endif
+
+#ifndef USE_ROCM
+torch::Tensor minimax_allreduce_rms(torch::Tensor const& input,
+                                    torch::Tensor const& norm_weight,
+                                    torch::Tensor workspace, int64_t const rank,
+                                    int64_t const nranks, double const eps);
+std::tuple<torch::Tensor, torch::Tensor> minimax_allreduce_rms_qk(
+    torch::Tensor qkv, torch::Tensor const& norm_weight_q,
+    torch::Tensor const& norm_weight_k, torch::Tensor workspace,
+    int64_t const q_size, int64_t const kv_size, int64_t const rank,
+    int64_t const nranks, double const eps);
 #endif

csrc/torch_bindings.cpp

@@ -668,6 +668,29 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "Tensor? b_qzeros, "
       "SymInt n, SymInt group_size, SymInt sm_count, SymInt sm_version, SymInt "
       "CUBLAS_M_THRESHOLD, bool has_zp, bool n32k16_reorder) -> Tensor");
+
+  ops.def(
+      "minimax_allreduce_rms("
+      "Tensor input,"
+      "Tensor norm_weight,"
+      "Tensor workspace,"
+      "int rank,"
+      "int nranks,"
+      "float eps) -> Tensor");
+  ops.impl("minimax_allreduce_rms", torch::kCUDA, &minimax_allreduce_rms);
+  ops.def(
+      "minimax_allreduce_rms_qk("
+      "Tensor qkv,"
+      "Tensor norm_weight_q,"
+      "Tensor norm_weight_k,"
+      "Tensor workspace,"
+      "int q_size,"
+      "int kv_size,"
+      "int rank,"
+      "int nranks,"
+      "float eps) -> (Tensor, Tensor)");
+  ops.impl("minimax_allreduce_rms_qk", torch::kCUDA, &minimax_allreduce_rms_qk);
+
   //  conditionally compiled so impl in source file
 #endif
 }

docker/Dockerfile

@@ -649,7 +649,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     else \
         BITSANDBYTES_VERSION="${BITSANDBYTES_VERSION_X86}"; \
     fi; \
-    uv pip install --system accelerate hf_transfer modelscope \
+    uv pip install --system accelerate modelscope \
         "bitsandbytes>=${BITSANDBYTES_VERSION}" "timm${TIMM_VERSION}" "runai-model-streamer[s3,gcs,azure]${RUNAI_MODEL_STREAMER_VERSION}"
 
 # ============================================================
@@ -772,9 +772,10 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install --system -e tests/vllm_test_utils
 
 # enable fast downloads from hf (for testing)
-RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system hf_transfer
-ENV HF_HUB_ENABLE_HF_TRANSFER 1
+ENV HF_XET_HIGH_PERFORMANCE 1
+
+# increase timeout for hf downloads (for testing)
+ENV HF_HUB_DOWNLOAD_TIMEOUT 60
 
 # Copy in the v1 package for testing (it isn't distributed yet)
 COPY vllm/v1 /usr/local/lib/python${PYTHON_VERSION}/dist-packages/vllm/v1

docker/Dockerfile.cpu

@@ -140,9 +140,11 @@ RUN \
     esac; \
     }; \
     remove_packages_not_supported_on_aarch64 && \
-    sed -i 's/^torch==.*/torch==2.10.0/g' requirements/cpu-test.in && \
+    sed -i 's/^torch==.*/torch==2.11.0/g' requirements/cpu-test.in && \
     sed -i 's/torchaudio.*/torchaudio/g' requirements/cpu-test.in && \
     sed -i 's/torchvision.*/torchvision/g' requirements/cpu-test.in && \
+    # Related issue: https://github.com/vllm-project/vllm/pull/38800#issuecomment-4228314305
+    sed -i 's/^sentence-transformers.*/sentence-transformers==5.3.0/g' requirements/cpu-test.in && \
     uv pip compile requirements/cpu-test.in -o requirements/cpu-test.txt --index-strategy unsafe-best-match --torch-backend cpu
 
 RUN --mount=type=cache,target=/root/.cache/uv \
@@ -195,6 +197,12 @@ ADD ./.buildkite/ ./.buildkite/
 RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install -e tests/vllm_test_utils
 
+# enable fast downloads from hf (for testing)
+ENV HF_XET_HIGH_PERFORMANCE 1
+
+# increase timeout for hf downloads (for testing)
+ENV HF_HUB_DOWNLOAD_TIMEOUT 60
+
 ######################### RELEASE IMAGE #########################
 FROM base AS vllm-openai
 

docker/Dockerfile.nightly_torch

@@ -269,9 +269,10 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install --system -e tests/vllm_test_utils
 
 # enable fast downloads from hf (for testing)
-RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system hf_transfer
-ENV HF_HUB_ENABLE_HF_TRANSFER 1
+ENV HF_XET_HIGH_PERFORMANCE 1
+
+# increase timeout for hf downloads (for testing)
+ENV HF_HUB_DOWNLOAD_TIMEOUT 60
 
 RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install --system -r requirements/nightly_torch_test.txt

docker/Dockerfile.rocm

@@ -364,9 +364,10 @@ RUN cd /vllm-workspace \
     && python3 -m pip install pytest-shard
 
 # enable fast downloads from hf (for testing)
-RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system hf_transfer
-ENV HF_HUB_ENABLE_HF_TRANSFER=1
+ENV HF_XET_HIGH_PERFORMANCE=1
+
+# increase timeout for hf downloads (for testing)
+ENV HF_HUB_DOWNLOAD_TIMEOUT 60
 
 # install audio decode package `torchcodec` from source (required due to 
 # ROCm and torch version mismatch) for tests with datasets package

docs/features/reasoning_outputs.md

@@ -244,12 +244,12 @@ response = client.chat.completions.create(
 
 Some models, such as [Qwen3](https://qwen.readthedocs.io/en/latest/getting_started/quickstart.html#thinking-budget), [DeepSeek](https://www.alibabacloud.com/help/en/model-studio/deep-thinking), and [Nemotron3](https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16), support a thinking budget that limits the maximum number of tokens used for reasoning.
 
-Token counting starts from `think_start_str`. Once the reasoning token count reaches the configured `thinking_token_budget`, vLLM forces the model to produce `think_end_str`, effectively terminating the reasoning block.
+Token counting starts from `reasoning_start_str`. Once the reasoning token count reaches the configured `thinking_token_budget`, vLLM forces the model to produce `reasoning_end_str`, effectively terminating the reasoning block.
 
 To use this feature:
 
 - `--reasoning-parser` enables reasoning extraction.
-- `--reasoning-config` defines the reasoning boundary tokens (e.g., `think_start_str`, `think_end_str`).
+- `--reasoning-config` defines the reasoning boundary tokens (e.g., `reasoning_start_str`, `reasoning_end_str`).
 - `thinking_token_budget` (a sampling parameter) sets the per-request reasoning token limit.
 
 If `thinking_token_budget` is not specified, no explicit reasoning limit is applied beyond normal generation constraints such as `max_tokens`.
@@ -257,20 +257,20 @@ If `thinking_token_budget` is not specified, no explicit reasoning limit is appl
 `--reasoning-config` accepts a JSON object corresponding to  
 [ReasoningConfig][vllm.config.ReasoningConfig] with the following fields:
 
-| Field             | Type           | Description                                      |
-|-------------------|----------------|--------------------------------------------------|
-| `think_start_str` | `str \| null`  | String that marks the start of reasoning content |
-| `think_end_str`   | `str \| null`  | String that marks the end of reasoning content   |
+| Field                 | Type           | Description                                      |
+|-----------------------|----------------|--------------------------------------------------|
+| `reasoning_start_str` | `str \| null`  | String that marks the start of reasoning content |
+| `reasoning_end_str`   | `str \| null`  | String that marks the end of reasoning content   |
 
 !!! note
-    `think_end_str` can include a transition phrase before the think end token. For example, setting `think_end_str` to `"I have to give the solution based on the thinking directly now.</think>"` instructs the model to emit that phrase when the budget is exhausted, making the reasoning termination more natural.
+    `reasoning_end_str` can include a transition phrase before the reasoning end token. For example, setting `reasoning_end_str` to `"I have to give the solution based on the reasoning directly now.</think>"` instructs the model to emit that phrase when the budget is exhausted, making the reasoning termination more natural.
 
 ### Online Serving
 
 ```bash
 vllm serve Qwen/Qwen3-0.6B \
     --reasoning-parser qwen3 \
-    --reasoning-config '{"think_start_str": "<think>", "think_end_str": "I have to give the solution based on the thinking directly now.</think>"}'
+    --reasoning-config '{"reasoning_start_str": "<think>", "reasoning_end_str": "I have to give the solution based on the reasoning directly now.</think>"}'
 ```
 
 Then make a request with `thinking_token_budget` to limit the reasoning tokens:
@@ -298,8 +298,8 @@ from vllm.config import ReasoningConfig
 llm = LLM(
     model="Qwen/Qwen3-0.6B",
     reasoning_config=ReasoningConfig(
-        think_start_str="<think>",
-        think_end_str="I have to give the solution based on the thinking directly now.</think>",
+        reasoning_start_str="<think>",
+        reasoning_end_str="I have to give the solution based on the thinking directly now.</think>",
     ),
 )
 

docs/getting_started/installation/gpu.rocm.inc.md

@@ -147,7 +147,7 @@ uv pip install vllm --extra-index-url https://wheels.vllm.ai/rocm/0.15.0/rocm700
         # Install dependencies
         pip install --upgrade numba \
             scipy \
-            huggingface-hub[cli,hf_transfer] \
+            huggingface-hub[cli] \
             setuptools_scm
         pip install -r requirements/rocm.txt
 

examples/tool_chat_template_gemma4.jinja

@@ -0,0 +1,331 @@
+{%- macro format_parameters(properties, required) -%}
+    {%- set standard_keys = ['description', 'type', 'properties', 'required', 'nullable'] -%}
+    {%- set ns = namespace(found_first=false) -%}
+    {%- for key, value in properties | dictsort -%}
+        {%- set add_comma = false -%}
+        {%- if key not in standard_keys -%}
+            {%- if ns.found_first %},{% endif -%}
+            {%- set ns.found_first = true -%}
+            {{ key }}:{
+            {%- if value['description'] -%}
+                description:<|"|>{{ value['description'] }}<|"|>
+                {%- set add_comma = true -%}
+            {%- endif -%}
+            {%- if value['nullable'] %}
+                {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                nullable:true
+            {%- endif -%}
+            {%- if value['type'] | upper == 'STRING' -%}
+                {%- if value['enum'] -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    enum:{{ format_argument(value['enum']) }}
+                {%- endif -%}
+            {%- elif value['type'] | upper == 'OBJECT' -%}
+                ,properties:{
+                {%- if value['properties'] is defined and value['properties'] is mapping -%}
+                    {{- format_parameters(value['properties'], value['required'] | default([])) -}}
+                {%- elif value is mapping -%}
+                    {{- format_parameters(value, value['required'] | default([])) -}}
+                {%- endif -%}
+                }
+                {%- if value['required'] -%}
+                    ,required:[
+                    {%- for item in value['required'] | default([]) -%}
+                        <|"|>{{- item -}}<|"|>
+                        {%- if not loop.last %},{% endif -%}
+                    {%- endfor -%}
+                    ]
+                {%- endif -%}
+            {%- elif value['type'] | upper == 'ARRAY' -%}
+                {%- if value['items'] is mapping and value['items'] -%}
+                    ,items:{
+                    {%- set ns_items = namespace(found_first=false) -%}
+                    {%- for item_key, item_value in value['items'] | dictsort -%}
+                        {%- if item_value is not none -%}
+                            {%- if ns_items.found_first %},{% endif -%}
+                            {%- set ns_items.found_first = true -%}
+                            {%- if item_key == 'properties' -%}
+                                properties:{
+                                {%- if item_value is mapping -%}
+                                    {{- format_parameters(item_value, value['items']['required'] | default([])) -}}
+                                {%- endif -%}
+                                }
+                            {%- elif item_key == 'required' -%}
+                                required:[
+                                {%- for req_item in item_value -%}
+                                    <|"|>{{- req_item -}}<|"|>
+                                    {%- if not loop.last %},{% endif -%}
+                                {%- endfor -%}
+                                ]
+                            {%- elif item_key == 'type' -%}
+                                {%- if item_value is string -%}
+                                    type:{{ format_argument(item_value | upper) }}
+                                {%- else -%}
+                                    type:{{ format_argument(item_value | map('upper') | list) }}
+                                {%- endif -%}
+                            {%- else -%}
+                                {{ item_key }}:{{ format_argument(item_value) }}
+                            {%- endif -%}
+                        {%- endif -%}
+                    {%- endfor -%}
+                    }
+                {%- endif -%}
+            {%- endif -%}
+            {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+            type:<|"|>{{ value['type'] | upper }}<|"|>}
+        {%- endif -%}
+    {%- endfor -%}
+{%- endmacro -%}
+{%- macro format_function_declaration(tool_data) -%}
+    declaration:{{- tool_data['function']['name'] -}}{description:<|"|>{{- tool_data['function']['description'] -}}<|"|>
+    {%- set params = tool_data['function']['parameters'] -%}
+    {%- if params -%}
+        ,parameters:{
+        {%- if params['properties'] -%}
+            properties:{ {{- format_parameters(params['properties'], params['required']) -}} },
+        {%- endif -%}
+        {%- if params['required'] -%}
+            required:[
+            {%- for item in params['required'] -%}
+                <|"|>{{- item -}}<|"|>
+                {{- ',' if not loop.last -}}
+            {%- endfor -%}
+            ],
+        {%- endif -%}
+        {%- if params['type'] -%}
+            type:<|"|>{{- params['type'] | upper -}}<|"|>}
+        {%- endif -%}
+    {%- endif -%}
+    {%- if 'response' in tool_data['function'] -%}
+        {%- set response_declaration = tool_data['function']['response'] -%}
+        ,response:{
+        {%- if response_declaration['description'] -%}
+            description:<|"|>{{- response_declaration['description'] -}}<|"|>,
+        {%- endif -%}
+        {%- if response_declaration['type'] | upper == 'OBJECT' -%}
+            type:<|"|>{{- response_declaration['type'] | upper -}}<|"|>}
+        {%- endif -%}
+    {%- endif -%}
+    }
+{%- endmacro -%}
+{%- macro format_argument(argument, escape_keys=True) -%}
+    {%- if argument is string -%}
+        {{- '<|"|>' + argument + '<|"|>' -}}
+    {%- elif argument is boolean -%}
+        {{- 'true' if argument else 'false' -}}
+    {%- elif argument is mapping -%}
+        {{- '{' -}}
+        {%- set ns = namespace(found_first=false) -%}
+        {%- for key, value in argument | dictsort -%}
+            {%- if ns.found_first %},{% endif -%}
+            {%- set ns.found_first = true -%}
+            {%- if escape_keys -%}
+                {{- '<|"|>' + key + '<|"|>' -}}
+            {%- else -%}
+                {{- key -}}
+            {%- endif -%}
+            :{{- format_argument(value, escape_keys=escape_keys) -}}
+        {%- endfor -%}
+        {{- '}' -}}
+    {%- elif argument is sequence -%}
+        {{- '[' -}}
+        {%- for item in argument -%}
+            {{- format_argument(item, escape_keys=escape_keys) -}}
+            {%- if not loop.last %},{% endif -%}
+        {%- endfor -%}
+        {{- ']' -}}
+    {%- else -%}
+        {{- argument -}}
+    {%- endif -%}
+{%- endmacro -%}
+{%- macro strip_thinking(text) -%}
+    {%- set ns = namespace(result='') -%}
+    {%- for part in text.split('<channel|>') -%}
+        {%- if '<|channel>' in part -%}
+            {%- set ns.result = ns.result + part.split('<|channel>')[0] -%}
+        {%- else -%}
+            {%- set ns.result = ns.result + part -%}
+        {%- endif -%}
+    {%- endfor -%}
+    {{- ns.result | trim -}}
+{%- endmacro -%}
+
+{%- macro format_tool_response_block(tool_name, response) -%}
+    {{- '<|tool_response>' -}}
+    {%- if response is mapping -%}
+        {{- 'response:' + tool_name + '{' -}}
+        {%- for key, value in response | dictsort -%}
+            {{- key -}}:{{- format_argument(value, escape_keys=False) -}}
+            {%- if not loop.last %},{% endif -%}
+        {%- endfor -%}
+        {{- '}' -}}
+    {%- else -%}
+        {{- 'response:' + tool_name + '{value:' + format_argument(response, escape_keys=False) + '}' -}}
+    {%- endif -%}
+    {{- '<tool_response|>' -}}
+{%- endmacro -%}
+
+{%- set ns = namespace(prev_message_type=None) -%}
+{%- set loop_messages = messages -%}
+{{ bos_token }}
+{%- if (enable_thinking is defined and enable_thinking) or tools or messages[0]['role'] in ['system', 'developer'] -%}
+    {{- '<|turn>system\n' -}}
+
+    {%- if enable_thinking is defined and enable_thinking -%}
+        {{- '<|think|>' -}}
+        {%- set ns.prev_message_type = 'think' -%}
+    {%- endif -%}
+
+    {%- if messages[0]['role'] in ['system', 'developer'] -%}
+        {{- messages[0]['content'] | trim -}}
+        {%- set loop_messages = messages[1:] -%}
+    {%- endif -%}
+
+    {%- if tools -%}
+        {%- for tool in tools %}
+            {{- '<|tool>' -}}
+            {{- format_function_declaration(tool) | trim -}}
+            {{- '<tool|>' -}}
+        {%- endfor %}
+        {%- set ns.prev_message_type = 'tool' -%}
+    {%- endif -%}
+
+    {{- '<turn|>\n' -}}
+{%- endif %}
+
+{%- set ns_turn = namespace(last_user_idx=-1) -%}
+{%- for i in range(loop_messages | length) -%}
+    {%- if loop_messages[i]['role'] == 'user' -%}
+        {%- set ns_turn.last_user_idx = i -%}
+    {%- endif -%}
+{%- endfor -%}
+
+{%- for message in loop_messages -%}
+    {%- if message['role'] != 'tool' -%}
+    {%- set ns.prev_message_type = None -%}
+    {%- set role = 'model' if message['role'] == 'assistant' else message['role'] -%}
+    {#- OpenAI may emit multiple assistant messages in one tool loop (user → asst → tool → asst → tool).
+        Only the first of those should open <|turn>model; later ones continue the same model turn. -#}
+    {%- set prev_nt = namespace(role=None, found=false) -%}
+    {%- if loop.index0 > 0 -%}
+        {%- for j in range(loop.index0 - 1, -1, -1) -%}
+            {%- if not prev_nt.found -%}
+                {%- if loop_messages[j]['role'] != 'tool' -%}
+                    {%- set prev_nt.role = loop_messages[j]['role'] -%}
+                    {%- set prev_nt.found = true -%}
+                {%- endif -%}
+            {%- endif -%}
+        {%- endfor -%}
+    {%- endif -%}
+    {%- set continue_same_model_turn = (role == 'model' and prev_nt.role == 'assistant') -%}
+    {%- if not continue_same_model_turn -%}
+        {{- '<|turn>' + role + '\n' }}
+    {%- endif -%}
+
+    {%- if message.get('reasoning') and loop.index0 > ns_turn.last_user_idx and message.get('tool_calls') -%}
+        {{- '<|channel>thought\n' + message['reasoning'] + '\n<channel|>'}}
+    {%- endif -%}
+
+            {%- if message['tool_calls'] -%}
+                {%- for tool_call in message['tool_calls'] -%}
+                    {%- set function = tool_call['function'] -%}
+                    {{- '<|tool_call>call:' + function['name'] + '{' -}}
+                    {%- if function['arguments'] is mapping -%}
+                        {%- set ns_args = namespace(found_first=false) -%}
+                        {%- for key, value in function['arguments'] | dictsort -%}
+                            {%- if ns_args.found_first %},{% endif -%}
+                            {%- set ns_args.found_first = true -%}
+                            {{- key -}}:{{- format_argument(value, escape_keys=False) -}}
+                        {%- endfor -%}
+                    {%- elif function['arguments'] is string -%}
+                        {{- function['arguments'] -}}
+                    {%- endif -%}
+                    {{- '}<tool_call|>' -}}
+                {%- endfor -%}
+                {%- set ns.prev_message_type = 'tool_call' -%}
+            {%- endif -%}
+
+            {%- set ns_tr_out = namespace(flag=false) -%}
+            {%- if message.get('tool_responses') -%}
+                {#- Legacy: tool_responses embedded on the assistant message -#}
+                {%- for tool_response in message['tool_responses'] -%}
+                    {{- format_tool_response_block(tool_response['name'] | default('unknown'), tool_response['response']) -}}
+                    {%- set ns_tr_out.flag = true -%}
+                    {%- set ns.prev_message_type = 'tool_response' -%}
+                {%- endfor -%}
+            {%- elif message.get('tool_calls') -%}
+                {#- OpenAI Chat Completions: consecutive following messages with role "tool" (no break/continue; range scan) -#}
+                {%- set ns_tool_scan = namespace(stopped=false) -%}
+                {%- for k in range(loop.index0 + 1, loop_messages | length) -%}
+                    {%- if ns_tool_scan.stopped -%}
+                    {%- elif loop_messages[k]['role'] != 'tool' -%}
+                        {%- set ns_tool_scan.stopped = true -%}
+                    {%- else -%}
+                        {%- set follow = loop_messages[k] -%}
+                        {%- set ns_tname = namespace(name=follow.get('name') | default('unknown')) -%}
+                        {%- for tc in message['tool_calls'] -%}
+                            {%- if tc.get('id') == follow.get('tool_call_id') -%}
+                                {%- set ns_tname.name = tc['function']['name'] -%}
+                            {%- endif -%}
+                        {%- endfor -%}
+                        {%- set tool_body = follow.get('content') -%}
+                        {%- if tool_body is string -%}
+                            {{- format_tool_response_block(ns_tname.name, tool_body) -}}
+                        {%- elif tool_body is sequence and tool_body is not string -%}
+                            {%- set ns_txt = namespace(s='') -%}
+                            {%- for part in tool_body -%}
+                                {%- if part.get('type') == 'text' -%}
+                                    {%- set ns_txt.s = ns_txt.s + (part.get('text') | default('')) -%}
+                                {%- endif -%}
+                            {%- endfor -%}
+                            {{- format_tool_response_block(ns_tname.name, ns_txt.s) -}}
+                        {%- else -%}
+                            {{- format_tool_response_block(ns_tname.name, tool_body) -}}
+                        {%- endif -%}
+                        {%- set ns_tr_out.flag = true -%}
+                        {%- set ns.prev_message_type = 'tool_response' -%}
+                    {%- endif -%}
+                {%- endfor -%}
+            {%- endif -%}
+
+            {%- if message['content'] is string -%}
+                {%- if role == 'model' -%}
+                    {{- strip_thinking(message['content']) -}}
+                {%- else -%}
+                    {{- message['content'] | trim -}}
+                {%- endif -%}
+            {%- elif message['content'] is sequence -%}
+                {%- for item in message['content'] -%}
+                    {%- if item['type'] == 'text' -%}
+                        {%- if role == 'model' -%}
+                            {{- strip_thinking(item['text']) -}}
+                        {%- else -%}
+                            {{- item['text'] | trim -}}
+                        {%- endif -%}
+                    {%- elif item['type'] == 'image' -%}
+                        {{- '\n\n<|image|>\n\n' -}}
+                        {%- set ns.prev_message_type = 'image' -%}
+                    {%- elif item['type'] == 'audio' -%}
+                        {{- '<|audio|>' -}}
+                        {%- set ns.prev_message_type = 'audio' -%}
+                    {%- elif item['type'] == 'video' -%}
+                        {{- '\n\n<|video|>\n\n' -}}
+                        {%- set ns.prev_message_type = 'video' -%}
+                    {%- endif -%}
+                {%- endfor -%}
+            {%- endif -%}
+
+        {%- if not (ns_tr_out.flag and not message.get('content')) -%}
+            {{- '<turn|>\n' -}}
+        {%- endif -%}
+    {%- endif -%}
+{%- endfor -%}
+
+{%- if add_generation_prompt -%}
+    {%- if ns.prev_message_type != 'tool_response' -%}
+        {{- '<|turn>model\n' -}}
+    {%- endif -%}
+    {%- if not enable_thinking | default(false) -%}
+        {{- '<|channel>thought\n<channel|>' -}}
+    {%- endif -%}
+{%- endif -%}

requirements/common.txt

@@ -7,7 +7,7 @@ requests >= 2.26.0
 tqdm
 blake3
 py-cpuinfo
-transformers >= 4.56.0, < 5
+transformers >= 4.56.0, != 5.0.*, != 5.1.*, != 5.2.*, != 5.3.*, != 5.4.*, != 5.5.0
 tokenizers >= 0.21.1  # Required for fast incremental detokenization.
 protobuf >= 5.29.6, !=6.30.*, !=6.31.*, !=6.32.*, !=6.33.0.*, !=6.33.1.*, !=6.33.2.*, !=6.33.3.*, !=6.33.4.* # Required by LlamaTokenizer, gRPC. CVE-2026-0994
 fastapi[standard] >= 0.115.0 # Required by FastAPI's form models in the OpenAI API server's audio transcriptions endpoint.
@@ -37,7 +37,7 @@ pyyaml
 six>=1.16.0; python_version > '3.11' # transitive dependency of pandas that needs to be the latest version for python 3.12
 setuptools>=77.0.3,<81.0.0; python_version > '3.11' # Setuptools is used by triton, we need to ensure a modern version is installed for 3.12+ so that it does not try to import distutils, which was removed in 3.12
 einops # Required for Qwen2-VL.
-compressed-tensors == 0.14.0.1 # required for compressed-tensors
+compressed-tensors == 0.15.0.1 # required for compressed-tensors
 depyf==0.20.0 # required for profiling and debugging with compilation config
 cloudpickle # allows pickling lambda functions in model_executor/models/registry.py
 watchfiles # required for http server to monitor the updates of TLS files

requirements/kv_connectors.txt

@@ -1,3 +1,5 @@
 lmcache >= 0.3.9
 nixl >= 0.7.1, < 0.10.0 # Required for disaggregated prefill
+nixl-cu12 >= 0.7.1, < 0.10.0
+nixl-cu13 >= 0.7.1, < 0.10.0
 mooncake-transfer-engine >= 0.3.8

requirements/nightly_torch_test.txt

@@ -29,8 +29,8 @@ 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[bm25s]>=2, <3 # required for mteb test
-transformers==4.57.5
-tokenizers==0.22.0
+transformers==5.5.3
+tokenizers==0.22.2
 schemathesis>=3.39.15 # Required for openai schema test.
 # quantization
 bitsandbytes>=0.49.2

requirements/rocm-test.in

@@ -36,8 +36,8 @@ 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[bm25s]>=2, <3 # required for mteb test
-transformers==4.57.5
-tokenizers==0.22.0
+transformers==5.5.3
+tokenizers==0.22.2
 schemathesis>=3.39.15 # Required for openai schema test
 # quantization
 bitsandbytes==0.49.2
@@ -80,4 +80,3 @@ plotly # required for perf comparison html report
 rapidfuzz
 torchgeo==0.7.0
 multiprocess==0.70.16
-huggingface-hub==0.36.2

requirements/rocm-test.txt

@@ -232,7 +232,6 @@ filelock==3.25.2
     #   python-discovery
     #   ray
     #   torch
-    #   transformers
     #   virtualenv
 fiona==1.10.1
     # via torchgeo
@@ -318,7 +317,7 @@ h5py==3.16.0
     # via terratorch
 harfile==0.4.0
     # via schemathesis
-hf-xet==1.4.2
+hf-xet==1.4.3
     # via huggingface-hub
 hiredis==3.3.1
     # via tensorizer
@@ -332,11 +331,11 @@ httpx==0.27.2
     # via
     #   -r requirements/rocm-test.in
     #   diffusers
+    #   huggingface-hub
     #   perceptron
     #   schemathesis
-huggingface-hub==0.36.2
+huggingface-hub==1.10.2
     # via
-    #   -r requirements/rocm-test.in
     #   accelerate
     #   datasets
     #   diffusers
@@ -970,7 +969,6 @@ requests==2.32.5
     #   google-api-core
     #   google-cloud-storage
     #   gpt-oss
-    #   huggingface-hub
     #   lightly
     #   lm-eval
     #   mistral-common
@@ -983,7 +981,6 @@ requests==2.32.5
     #   starlette-testclient
     #   tacoreader
     #   tiktoken
-    #   transformers
     #   wandb
 resampy==0.4.3
     # via -r requirements/rocm-test.in
@@ -1191,7 +1188,7 @@ timm==1.0.17
     #   segmentation-models-pytorch
     #   terratorch
     #   torchgeo
-tokenizers==0.22.0
+tokenizers==0.22.2
     # via
     #   -c requirements/common.txt
     #   -r requirements/rocm-test.in
@@ -1230,7 +1227,7 @@ tqdm==4.67.3
     #   tacoreader
     #   terratorch
     #   transformers
-transformers==4.57.5
+transformers==5.5.3
     # via
     #   -c requirements/common.txt
     #   -r requirements/rocm-test.in
@@ -1252,7 +1249,9 @@ typepy==1.3.4
 typer==0.24.1
     # via
     #   fastsafetensors
+    #   huggingface-hub
     #   perceptron
+    #   transformers
 typeshed-client==2.9.0
     # via jsonargparse
 typing-extensions==4.15.0

requirements/test.in

@@ -18,7 +18,7 @@ httpx
 librosa # required for audio tests
 vector_quantize_pytorch # required for minicpmo_26 test
 vocos # required for minicpmo_26 test
-peft>=0.15.0 # required for phi-4-mm test
+peft>=0.18.1 # required for phi-4-mm test
 pqdm
 ray[cgraph,default]>=2.48.0 # Ray Compiled Graph, required by pipeline parallelism tests
 resampy # required for audio tests
@@ -39,8 +39,8 @@ 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[bm25s]>=2, <3 # required for mteb test
-transformers==4.57.5
-tokenizers==0.22.0
+transformers==5.5.3
+tokenizers==0.22.2
 schemathesis>=3.39.15 # Required for openai schema test.
 # quantization
 bitsandbytes==0.49.2

requirements/test.txt

@@ -4,7 +4,7 @@ absl-py==2.1.0
     # via
     #   rouge-score
     #   tensorboard
-accelerate==1.0.1
+accelerate==1.13.0
     # via peft
 aenum==3.1.16
     # via lightly
@@ -240,7 +240,6 @@ filelock==3.16.1
     #   huggingface-hub
     #   ray
     #   torch
-    #   transformers
     #   virtualenv
 fiona==1.10.1
     # via torchgeo
@@ -323,7 +322,7 @@ h5py==3.13.0
     # via terratorch
 harfile==0.3.0
     # via schemathesis
-hf-xet==1.1.7
+hf-xet==1.4.3
     # via huggingface-hub
 hiredis==3.0.0
     # via tensorizer
@@ -337,9 +336,10 @@ httpx==0.27.2
     # via
     #   -r requirements/test.in
     #   diffusers
+    #   huggingface-hub
     #   perceptron
     #   schemathesis
-huggingface-hub==0.36.2
+huggingface-hub==1.10.2
     # via
     #   accelerate
     #   datasets
@@ -740,7 +740,7 @@ pathvalidate==3.2.1
     # via pytablewriter
 patsy==1.0.1
     # via statsmodels
-peft==0.16.0
+peft==0.18.1
     # via -r requirements/test.in
 perceptron==0.1.4
     # via -r requirements/test.in
@@ -963,7 +963,7 @@ referencing==0.35.1
     # via
     #   jsonschema
     #   jsonschema-specifications
-regex==2024.9.11
+regex==2026.2.28
     # via
     #   diffusers
     #   nltk
@@ -982,7 +982,6 @@ requests==2.32.3
     #   google-api-core
     #   google-cloud-storage
     #   gpt-oss
-    #   huggingface-hub
     #   lightly
     #   lm-eval
     #   mistral-common
@@ -995,7 +994,6 @@ requests==2.32.3
     #   starlette-testclient
     #   tacoreader
     #   tiktoken
-    #   transformers
     #   wandb
 resampy==0.4.3
     # via -r requirements/test.in
@@ -1193,7 +1191,7 @@ timm==1.0.17
     #   segmentation-models-pytorch
     #   terratorch
     #   torchgeo
-tokenizers==0.22.0
+tokenizers==0.22.2
     # via
     #   -r requirements/test.in
     #   transformers
@@ -1269,7 +1267,7 @@ tqdm==4.67.3
     #   tacoreader
     #   terratorch
     #   transformers
-transformers==4.57.5
+transformers==5.5.3
     # via
     #   -r requirements/test.in
     #   genai-perf
@@ -1290,7 +1288,9 @@ typepy==1.3.2
 typer==0.15.2
     # via
     #   fastsafetensors
+    #   huggingface-hub
     #   perceptron
+    #   transformers
 types-python-dateutil==2.9.0.20241206
     # via arrow
 typeshed-client==2.8.2

requirements/test/xpu.txt

@@ -0,0 +1,736 @@
+# This file was autogenerated by uv via the following command:
+#    uv pip compile requirements/test/xpu.in -c requirements/xpu.txt -o requirements/test/xpu.txt --index-strategy unsafe-best-match --torch-backend xpu --python-platform x86_64-manylinux_2_39 --python-version 3.12
+absl-py==2.4.0
+    # via
+    #   -r requirements/test/xpu.in
+    #   rouge-score
+accelerate==1.13.0
+    # via -r requirements/test/xpu.in
+aiohappyeyeballs==2.6.1
+    # via aiohttp
+aiohttp==3.13.4
+    # via
+    #   -c requirements/common.txt
+    #   fsspec
+    #   gpt-oss
+    #   lm-eval
+aiosignal==1.4.0
+    # via aiohttp
+albumentations==1.4.6
+    # via -r requirements/test/xpu.in
+annotated-doc==0.0.4
+    # via
+    #   fastapi
+    #   typer
+annotated-types==0.7.0
+    # via pydantic
+anyio==4.13.0
+    # via
+    #   httpx
+    #   starlette
+arctic-inference==0.1.1
+    # via -r requirements/test/xpu.in
+attrs==26.1.0
+    # via
+    #   aiohttp
+    #   jsonlines
+    #   jsonschema
+    #   referencing
+audioread==3.0.1
+    # via
+    #   -r requirements/test/xpu.in
+    #   librosa
+blobfile==3.0.0
+    # via -r requirements/test/xpu.in
+bm25s==0.2.13
+    # via
+    #   -r requirements/test/xpu.in
+    #   mteb
+bounded-pool-executor==0.0.3
+    # via pqdm
+certifi==2026.2.25
+    # via
+    #   httpcore
+    #   httpx
+    #   requests
+cffi==2.0.0
+    # via soundfile
+chardet==5.2.0
+    # via mbstrdecoder
+charset-normalizer==3.4.6
+    # via requests
+chz==0.4.0
+    # via gpt-oss
+click==8.3.1
+    # via
+    #   jiwer
+    #   nltk
+    #   schemathesis
+    #   typer
+    #   uvicorn
+colorama==0.4.6
+    # via sacrebleu
+coverage==7.13.5
+    # via pytest-cov
+dataproperty==1.1.0
+    # via
+    #   pytablewriter
+    #   tabledata
+datasets==4.8.4
+    # via
+    #   evaluate
+    #   lm-eval
+    #   mteb
+decorator==5.2.1
+    # via librosa
+dill==0.4.1
+    # via
+    #   datasets
+    #   evaluate
+    #   lm-eval
+    #   multiprocess
+docker==7.1.0
+    # via gpt-oss
+docopt==0.6.2
+    # via num2words
+dpcpp-cpp-rt==2025.3.1
+    # via
+    #   onemkl-sycl-blas
+    #   onemkl-sycl-dft
+    #   onemkl-sycl-lapack
+    #   onemkl-sycl-rng
+    #   onemkl-sycl-sparse
+    #   torch
+evaluate==0.4.6
+    # via lm-eval
+fastapi==0.135.2
+    # via
+    #   -c requirements/common.txt
+    #   gpt-oss
+filelock==3.25.2
+    # via
+    #   -c requirements/common.txt
+    #   blobfile
+    #   datasets
+    #   huggingface-hub
+    #   modelscope
+    #   torch
+frozenlist==1.8.0
+    # via
+    #   aiohttp
+    #   aiosignal
+fsspec==2026.2.0
+    # via
+    #   datasets
+    #   evaluate
+    #   huggingface-hub
+    #   torch
+gpt-oss==0.0.8
+    # via -r requirements/test/xpu.in
+graphql-core==3.2.8
+    # via hypothesis-graphql
+h11==0.16.0
+    # via
+    #   httpcore
+    #   uvicorn
+harfile==0.4.0
+    # via schemathesis
+hf-xet==1.4.3
+    # via huggingface-hub
+html2text==2025.4.15
+    # via gpt-oss
+httpcore==1.0.9
+    # via httpx
+httpx==0.28.1
+    # via
+    #   datasets
+    #   huggingface-hub
+    #   schemathesis
+huggingface-hub==1.10.2
+    # via
+    #   accelerate
+    #   datasets
+    #   evaluate
+    #   sentence-transformers
+    #   timm
+    #   tokenizers
+    #   transformers
+hypothesis==6.151.10
+    # via
+    #   hypothesis-graphql
+    #   hypothesis-jsonschema
+    #   schemathesis
+hypothesis-graphql==0.12.0
+    # via schemathesis
+hypothesis-jsonschema==0.23.1
+    # via schemathesis
+idna==3.11
+    # via
+    #   anyio
+    #   httpx
+    #   requests
+    #   yarl
+imageio==2.37.3
+    # via scikit-image
+impi-rt==2021.17.0
+    # via
+    #   oneccl
+    #   torch
+iniconfig==2.3.0
+    # via pytest
+intel-cmplr-lib-rt==2025.3.1
+    # via
+    #   intel-sycl-rt
+    #   torch
+intel-cmplr-lib-ur==2025.3.1
+    # via
+    #   intel-openmp
+    #   intel-sycl-rt
+    #   torch
+intel-cmplr-lic-rt==2025.3.1
+    # via
+    #   intel-opencl-rt
+    #   intel-sycl-rt
+    #   torch
+intel-opencl-rt==2025.3.1
+    # via
+    #   dpcpp-cpp-rt
+    #   onemkl-sycl-blas
+    #   onemkl-sycl-dft
+    #   onemkl-sycl-lapack
+    #   onemkl-sycl-rng
+    #   onemkl-sycl-sparse
+    #   torch
+intel-openmp==2025.3.1
+    # via
+    #   dpcpp-cpp-rt
+    #   mkl
+    #   torch
+intel-pti==0.15.0
+    # via torch
+intel-sycl-rt==2025.3.1
+    # via
+    #   dpcpp-cpp-rt
+    #   oneccl
+    #   torch
+jinja2==3.1.6
+    # via
+    #   -c requirements/xpu.txt
+    #   lm-eval
+    #   torch
+jiwer==4.0.0
+    # via -r requirements/test/xpu.in
+joblib==1.5.3
+    # via
+    #   librosa
+    #   nltk
+    #   scikit-learn
+jsonlines==4.0.0
+    # via lm-eval
+jsonschema==4.26.0
+    # via
+    #   hypothesis-jsonschema
+    #   mistral-common
+    #   schemathesis
+jsonschema-rs==0.45.0
+    # via schemathesis
+jsonschema-specifications==2025.9.1
+    # via jsonschema
+junit-xml==1.9
+    # via schemathesis
+lazy-loader==0.5
+    # via
+    #   librosa
+    #   scikit-image
+librosa==0.10.2.post1
+    # via -r requirements/test/xpu.in
+llvmlite==0.44.0
+    # via numba
+lm-eval==0.4.11
+    # via -r requirements/test/xpu.in
+lxml==6.0.2
+    # via
+    #   blobfile
+    #   gpt-oss
+    #   sacrebleu
+markdown-it-py==4.0.0
+    # via rich
+markupsafe==3.0.3
+    # via
+    #   jinja2
+    #   werkzeug
+mbstrdecoder==1.1.4
+    # via
+    #   dataproperty
+    #   pytablewriter
+    #   typepy
+mdurl==0.1.2
+    # via markdown-it-py
+mistral-common==1.11.0
+    # via
+    #   -c requirements/common.txt
+    #   -r requirements/test/xpu.in
+mkl==2025.3.0
+    # via
+    #   onemkl-sycl-blas
+    #   onemkl-sycl-dft
+    #   onemkl-sycl-lapack
+    #   onemkl-sycl-rng
+    #   onemkl-sycl-sparse
+    #   torch
+modelscope==1.35.3
+    # via -r requirements/test/xpu.in
+more-itertools==10.8.0
+    # via lm-eval
+mpmath==1.3.0
+    # via sympy
+msgpack==1.1.2
+    # via librosa
+mteb==2.12.7
+    # via -r requirements/test/xpu.in
+multidict==6.7.1
+    # via
+    #   aiohttp
+    #   yarl
+multiprocess==0.70.19
+    # via
+    #   datasets
+    #   evaluate
+networkx==3.6.1
+    # via
+    #   scikit-image
+    #   torch
+nltk==3.9.4
+    # via rouge-score
+num2words==0.5.14
+    # via -r requirements/test/xpu.in
+numba==0.61.2
+    # via
+    #   -c requirements/xpu.txt
+    #   librosa
+numpy==2.2.6
+    # via
+    #   accelerate
+    #   albumentations
+    #   bm25s
+    #   datasets
+    #   evaluate
+    #   imageio
+    #   librosa
+    #   lm-eval
+    #   mistral-common
+    #   mteb
+    #   numba
+    #   opencv-python-headless
+    #   pandas
+    #   pytrec-eval-terrier
+    #   rouge-score
+    #   sacrebleu
+    #   scikit-image
+    #   scikit-learn
+    #   scipy
+    #   sentence-transformers
+    #   soundfile
+    #   soxr
+    #   tifffile
+    #   torchvision
+    #   transformers
+oneccl==2021.17.1
+    # via
+    #   oneccl-devel
+    #   torch
+oneccl-devel==2021.17.1
+    # via torch
+onemkl-license==2025.3.0
+    # via
+    #   mkl
+    #   torch
+onemkl-sycl-blas==2025.3.0
+    # via
+    #   onemkl-sycl-lapack
+    #   onemkl-sycl-sparse
+    #   torch
+onemkl-sycl-dft==2025.3.0
+    # via torch
+onemkl-sycl-lapack==2025.3.0
+    # via torch
+onemkl-sycl-rng==2025.3.0
+    # via torch
+onemkl-sycl-sparse==2025.3.0
+    # via torch
+openai-harmony==0.0.8
+    # via
+    #   -c requirements/common.txt
+    #   gpt-oss
+opencv-python-headless==4.13.0.92
+    # via
+    #   -c requirements/common.txt
+    #   albumentations
+    #   mistral-common
+packaging==26.0
+    # via
+    #   -c requirements/xpu.txt
+    #   accelerate
+    #   datasets
+    #   evaluate
+    #   huggingface-hub
+    #   lazy-loader
+    #   modelscope
+    #   pooch
+    #   pytest
+    #   pytest-rerunfailures
+    #   scikit-image
+    #   transformers
+    #   typepy
+pandas==3.0.1
+    # via
+    #   datasets
+    #   evaluate
+pathvalidate==3.3.1
+    # via pytablewriter
+pillow==12.1.1
+    # via
+    #   imageio
+    #   mistral-common
+    #   scikit-image
+    #   torchvision
+platformdirs==4.9.4
+    # via pooch
+pluggy==1.6.0
+    # via
+    #   pytest
+    #   pytest-cov
+polars==1.39.3
+    # via mteb
+polars-runtime-32==1.39.3
+    # via polars
+pooch==1.8.2
+    # via
+    #   -r requirements/test/xpu.in
+    #   librosa
+portalocker==3.2.0
+    # via sacrebleu
+pqdm==0.2.0
+    # via -r requirements/test/xpu.in
+propcache==0.4.1
+    # via
+    #   aiohttp
+    #   yarl
+psutil==7.2.2
+    # via accelerate
+py==1.11.0
+    # via pytest-forked
+pyarrow==23.0.1
+    # via datasets
+pycountry==26.2.16
+    # via pydantic-extra-types
+pycparser==3.0
+    # via cffi
+pycryptodomex==3.23.0
+    # via blobfile
+pydantic==2.12.5
+    # via
+    #   -c requirements/common.txt
+    #   albumentations
+    #   fastapi
+    #   gpt-oss
+    #   mistral-common
+    #   mteb
+    #   openai-harmony
+    #   pydantic-extra-types
+pydantic-core==2.41.5
+    # via pydantic
+pydantic-extra-types==2.11.1
+    # via mistral-common
+pyelftools==0.32
+    # via triton-xpu
+pygments==2.20.0
+    # via
+    #   pytest
+    #   rich
+pyrate-limiter==4.1.0
+    # via schemathesis
+pystemmer==3.0.0
+    # via
+    #   -r requirements/test/xpu.in
+    #   mteb
+pytablewriter==1.2.1
+    # via lm-eval
+pytest==9.0.2
+    # via
+    #   -r requirements/test/xpu.in
+    #   pytest-asyncio
+    #   pytest-cov
+    #   pytest-forked
+    #   pytest-rerunfailures
+    #   pytest-shard
+    #   pytest-timeout
+    #   schemathesis
+pytest-asyncio==1.3.0
+    # via -r requirements/test/xpu.in
+pytest-cov==6.3.0
+    # via -r requirements/test/xpu.in
+pytest-forked==1.6.0
+    # via -r requirements/test/xpu.in
+pytest-rerunfailures==14.0
+    # via -r requirements/test/xpu.in
+pytest-shard==0.1.2
+    # via -r requirements/test/xpu.in
+pytest-timeout==2.3.1
+    # via -r requirements/test/xpu.in
+python-dateutil==2.9.0.post0
+    # via
+    #   pandas
+    #   typepy
+pytrec-eval-terrier==0.5.10
+    # via mteb
+pytz==2026.1.post1
+    # via typepy
+pyyaml==6.0.3
+    # via
+    #   accelerate
+    #   albumentations
+    #   datasets
+    #   huggingface-hub
+    #   schemathesis
+    #   timm
+    #   transformers
+rapidfuzz==3.12.1
+    # via
+    #   -r requirements/test/xpu.in
+    #   jiwer
+referencing==0.37.0
+    # via
+    #   jsonschema
+    #   jsonschema-specifications
+regex==2026.3.32
+    # via
+    #   nltk
+    #   sacrebleu
+    #   tiktoken
+    #   transformers
+requests==2.33.1
+    # via
+    #   -c requirements/common.txt
+    #   datasets
+    #   docker
+    #   evaluate
+    #   gpt-oss
+    #   lm-eval
+    #   mistral-common
+    #   modelscope
+    #   mteb
+    #   pooch
+    #   schemathesis
+    #   starlette-testclient
+    #   tiktoken
+rich==14.3.3
+    # via
+    #   mteb
+    #   schemathesis
+    #   typer
+rouge-score==0.1.2
+    # via lm-eval
+rpds-py==0.30.0
+    # via
+    #   jsonschema
+    #   referencing
+sacrebleu==2.6.0
+    # via lm-eval
+safetensors==0.7.0
+    # via
+    #   accelerate
+    #   timm
+    #   transformers
+schemathesis==4.14.2
+    # via -r requirements/test/xpu.in
+scikit-image==0.26.0
+    # via albumentations
+scikit-learn==1.8.0
+    # via
+    #   albumentations
+    #   librosa
+    #   lm-eval
+    #   mteb
+    #   sentence-transformers
+scipy==1.17.1
+    # via
+    #   albumentations
+    #   bm25s
+    #   librosa
+    #   mteb
+    #   pytrec-eval-terrier
+    #   scikit-image
+    #   scikit-learn
+    #   sentence-transformers
+sentence-transformers==5.3.0
+    # via mteb
+setuptools==80.10.2
+    # via
+    #   -c requirements/common.txt
+    #   -c requirements/xpu.txt
+    #   modelscope
+    #   pytablewriter
+    #   torch
+shellingham==1.5.4
+    # via typer
+six==1.17.0
+    # via
+    #   -c requirements/common.txt
+    #   junit-xml
+    #   python-dateutil
+    #   rouge-score
+sortedcontainers==2.4.0
+    # via hypothesis
+soundfile==0.13.1
+    # via
+    #   -r requirements/test/xpu.in
+    #   librosa
+    #   mistral-common
+soxr==0.5.0.post1
+    # via
+    #   -r requirements/test/xpu.in
+    #   librosa
+    #   mistral-common
+sqlitedict==2.1.0
+    # via lm-eval
+starlette==1.0.0
+    # via
+    #   fastapi
+    #   starlette-testclient
+starlette-testclient==0.4.1
+    # via schemathesis
+structlog==25.5.0
+    # via gpt-oss
+sympy==1.14.0
+    # via torch
+tabledata==1.3.4
+    # via pytablewriter
+tabulate==0.10.0
+    # via sacrebleu
+tbb==2022.3.0
+    # via
+    #   intel-opencl-rt
+    #   mkl
+    #   torch
+tblib==3.1.0
+    # via -r requirements/test/xpu.in
+tcmlib==1.4.1
+    # via
+    #   tbb
+    #   torch
+    #   umf
+tcolorpy==0.1.7
+    # via pytablewriter
+tenacity==9.1.4
+    # via
+    #   gpt-oss
+    #   lm-eval
+    #   schemathesis
+termcolor==3.3.0
+    # via gpt-oss
+threadpoolctl==3.6.0
+    # via scikit-learn
+tifffile==2026.3.3
+    # via scikit-image
+tiktoken==0.12.0
+    # via
+    #   -c requirements/common.txt
+    #   gpt-oss
+    #   lm-eval
+    #   mistral-common
+timm==1.0.17
+    # via -r requirements/test/xpu.in
+tokenizers==0.22.2
+    # via
+    #   -c requirements/common.txt
+    #   transformers
+torch==2.10.0+xpu
+    # via
+    #   -c requirements/xpu.txt
+    #   accelerate
+    #   mteb
+    #   sentence-transformers
+    #   timm
+    #   torchvision
+torchvision==0.25.0+xpu
+    # via timm
+tqdm==4.67.3
+    # via
+    #   datasets
+    #   evaluate
+    #   huggingface-hub
+    #   lm-eval
+    #   modelscope
+    #   mteb
+    #   nltk
+    #   pqdm
+    #   sentence-transformers
+    #   transformers
+transformers==5.5.3
+    # via
+    #   -c requirements/common.txt
+    #   sentence-transformers
+triton-xpu==3.6.0
+    # via torch
+typepy==1.3.4
+    # via
+    #   dataproperty
+    #   pytablewriter
+    #   tabledata
+typer==0.24.1
+    # via
+    #   huggingface-hub
+    #   transformers
+typing-extensions==4.15.0
+    # via
+    #   -c requirements/common.txt
+    #   aiosignal
+    #   albumentations
+    #   anyio
+    #   chz
+    #   fastapi
+    #   huggingface-hub
+    #   librosa
+    #   lm-eval
+    #   mistral-common
+    #   mteb
+    #   pqdm
+    #   pydantic
+    #   pydantic-core
+    #   pydantic-extra-types
+    #   pytest-asyncio
+    #   referencing
+    #   schemathesis
+    #   sentence-transformers
+    #   starlette
+    #   torch
+    #   typing-inspection
+typing-inspection==0.4.2
+    # via
+    #   fastapi
+    #   pydantic
+umf==1.0.2
+    # via
+    #   intel-cmplr-lib-ur
+    #   torch
+urllib3==2.6.3
+    # via
+    #   blobfile
+    #   docker
+    #   modelscope
+    #   requests
+uvicorn==0.42.0
+    # via gpt-oss
+werkzeug==3.1.7
+    # via schemathesis
+word2number==1.1
+    # via lm-eval
+xxhash==3.6.0
+    # via
+    #   datasets
+    #   evaluate
+yarl==1.23.0
+    # via aiohttp
+zstandard==0.25.0
+    # via lm-eval

requirements/xpu-test.in

@@ -9,6 +9,8 @@ pytest-shard
 # --- Core Tools & Bindings ---
 absl-py
 arctic-inference
+lm_eval[api]
+modelscope
 
 # --- Audio Processing ---
 librosa

tests/conftest.py

@@ -409,6 +409,15 @@ class HfRunner:
             model_name,
             trust_remote_code=trust_remote_code,
         )
+        # HF runner should use the HF config so that it's consistent with the HF model
+        if self.config.__module__.startswith("vllm.transformers_utils.configs"):
+            from transformers.models.auto.configuration_auto import CONFIG_MAPPING
+
+            del CONFIG_MAPPING._extra_content[self.config.model_type]
+            self.config = AutoConfig.from_pretrained(
+                model_name,
+                trust_remote_code=trust_remote_code,
+            )
         self.device = self.get_default_device()
         self.dtype = dtype = _get_and_verify_dtype(
             self.model_name,

tests/evals/gsm8k/configs/Qwen3.5-397B-A17B-NVFP4-DEP2.yaml

@@ -7,3 +7,4 @@ server_args: >-
   --max-model-len 4096
   --data-parallel-size 2
   --enable-expert-parallel
+  --max-num-seqs 512

tests/kernels/core/test_minimax_reduce_rms.py

@@ -0,0 +1,152 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for MiniMax QK RMS-norm: NCCL reference vs Lamport fused kernel."""
+
+import pytest
+import torch
+import torch.nn as nn
+from torch.multiprocessing import spawn
+
+from tests.kernels.utils import opcheck
+from tests.utils import ensure_current_vllm_config, init_test_distributed_environment
+from vllm.distributed import cleanup_dist_env_and_memory
+from vllm.model_executor.layers.mamba.linear_attn import MiniMaxText01RMSNormTP
+from vllm.platforms import current_platform
+from vllm.utils.network_utils import get_open_port
+from vllm.utils.torch_utils import set_random_seed
+
+
+@ensure_current_vllm_config()
+def _worker_forward_qk(
+    local_rank,
+    world_size,
+    port,
+    num_tokens,
+    hidden_q_full,
+    hidden_k_full,
+    dtype,
+    seed,
+    eps,
+):
+    """Per-rank worker: compare NCCL allreduce path vs Lamport fused kernel."""
+
+    if not hasattr(torch.ops._C, "minimax_allreduce_rms_qk"):
+        cleanup_dist_env_and_memory()
+        return
+    device = torch.device(f"cuda:{local_rank}")
+    torch.accelerator.set_device_index(device)
+    init_test_distributed_environment(
+        world_size, 1, local_rank, port, local_rank=local_rank
+    )
+
+    hq = hidden_q_full // world_size
+    hk = hidden_k_full // world_size
+
+    q_norm = MiniMaxText01RMSNormTP(hidden_q_full, eps=eps).cuda()
+    k_norm = MiniMaxText01RMSNormTP(hidden_k_full, eps=eps).cuda()
+
+    set_random_seed(seed)
+    qw = torch.randn(hidden_q_full, dtype=dtype, device="cuda")
+    kw = torch.randn(hidden_k_full, dtype=dtype, device="cuda")
+    q_norm.weight = nn.Parameter(qw[local_rank * hq : (local_rank + 1) * hq])
+    k_norm.weight = nn.Parameter(kw[local_rank * hk : (local_rank + 1) * hk])
+
+    torch.manual_seed(seed + 1000 + local_rank)
+    qkv = torch.randn(num_tokens, hq + hk + hk, dtype=dtype, device="cuda")
+
+    q_ref, k_ref, v_ref = qkv.clone().split([hq, hk, hk], dim=-1)
+    ref_q, ref_k = MiniMaxText01RMSNormTP.forward_qk(q_norm, k_norm, q_ref, k_ref)
+
+    # Set up Lamport workspace.
+    from vllm.distributed.parallel_state import get_tp_group
+    from vllm.model_executor.layers.mamba.lamport_workspace import (
+        get_allreduce_workspace,
+    )
+
+    workspace = get_allreduce_workspace(
+        rank=local_rank,
+        world_size=world_size,
+        max_tokens=num_tokens,
+        process_group=get_tp_group().cpu_group,
+    )
+
+    opcheck(
+        torch.ops._C.minimax_allreduce_rms_qk,
+        (
+            qkv.clone(),
+            q_norm.weight,
+            k_norm.weight,
+            workspace,
+            hq,
+            hk,
+            local_rank,
+            world_size,
+            eps,
+        ),
+    )
+    fused_q, fused_k = torch.ops._C.minimax_allreduce_rms_qk(
+        qkv.clone(),
+        q_norm.weight,
+        k_norm.weight,
+        workspace,
+        hq,
+        hk,
+        local_rank,
+        world_size,
+        eps,
+    )
+    _, _, fused_v = qkv.split([hq, hk, hk], dim=-1)
+    torch.accelerator.synchronize()
+
+    torch.testing.assert_close(
+        fused_q,
+        ref_q,
+        atol=3e-2,
+        rtol=3e-2,
+    )
+    torch.testing.assert_close(fused_k, ref_k, atol=3e-2, rtol=3e-2)
+
+    cleanup_dist_env_and_memory()
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda(),
+    reason="CUDA required",
+)
+@pytest.mark.parametrize("world_size", [2, 4, 8])
+@pytest.mark.parametrize("num_tokens", [1, 128, 333])
+@pytest.mark.parametrize(
+    "hidden_dims",
+    [(6144, 1024)],
+)
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@pytest.mark.parametrize("eps", [1e-6])
+@pytest.mark.parametrize("seed", [42])
+def test_minimax_reduce_rms_qk(
+    world_size,
+    num_tokens,
+    hidden_dims,
+    dtype,
+    eps,
+    seed,
+):
+    num_gpus = current_platform.device_count()
+    if num_gpus < world_size:
+        pytest.skip(f"Need >= {world_size} GPUs, have {num_gpus}")
+    hidden_q_full, hidden_k_full = hidden_dims
+    port = str(get_open_port())
+    spawn(
+        _worker_forward_qk,
+        args=(
+            world_size,
+            port,
+            num_tokens,
+            hidden_q_full,
+            hidden_k_full,
+            dtype,
+            seed,
+            eps,
+        ),
+        nprocs=world_size,
+        join=True,
+    )

tests/lora/conftest.py

@@ -3,6 +3,7 @@
 
 import tempfile
 from collections import OrderedDict
+from importlib import reload
 from unittest.mock import MagicMock
 
 import pytest
@@ -43,6 +44,18 @@ def cleanup_fixture(should_do_global_cleanup_after_test: bool):
         cleanup_dist_env_and_memory(shutdown_ray=True)
 
 
+@pytest.fixture
+def maybe_enable_lora_dual_stream(monkeypatch: pytest.MonkeyPatch):
+    if current_platform.is_cuda():
+        monkeypatch.setenv("VLLM_LORA_ENABLE_DUAL_STREAM", "1")
+        import vllm.lora.layers.base_linear
+
+        if not hasattr(vllm.lora.layers.base_linear, "lora_linear_async"):
+            # Reload the module to ensure the environment variable takes effect.
+            reload(vllm.lora.layers.base_linear)
+    yield
+
+
 @pytest.fixture
 def dist_init():
     from tests.utils import ensure_current_vllm_config

tests/lora/test_lora_checkpoints.py

@@ -5,7 +5,9 @@ import pytest
 
 from vllm.lora.lora_model import LoRAModel
 from vllm.lora.peft_helper import PEFTHelper
+from vllm.lora.utils import parse_fine_tuned_lora_name
 from vllm.model_executor.models.baichuan import BaiChuanBaseForCausalLM
+from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM
 from vllm.model_executor.models.utils import WeightsMapper
 
 lora_lst = ["baichuan7B", "baichuan7B-zero", "baichuan7B-zero-regex", "chatglm3-6b"]
@@ -128,3 +130,24 @@ def test_lora_weights_mapping(baichuan_lora_files):
     for name in lora_model.loras:
         assert name.startswith(hf_to_vllm_mapper.orig_to_new_prefix["model."])
         assert ".baichuan_layers." in name
+
+
+def test_gemma4_lora_weights_mapping():
+    mapper = Gemma4ForCausalLM.hf_to_vllm_mapper
+    name = "base_model.model.model.language_model.layers.9.mlp.down_proj.lora_A.weight"
+    assert parse_fine_tuned_lora_name(name, mapper) == (
+        "model.layers.9.mlp.down_proj",
+        True,
+    )
+
+
+def test_gemma4_moe_lora_weights_mapping():
+    mapper = Gemma4ForCausalLM.hf_to_vllm_mapper
+    name = (
+        "base_model.model.model.language_model.layers.9.moe.experts."
+        "gate_up_proj.lora_B.weight"
+    )
+    assert parse_fine_tuned_lora_name(name, mapper) == (
+        "model.layers.9.moe.gate_up_proj",
+        False,
+    )

tests/lora/test_minicpmv_tp.py

@@ -1,7 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from importlib.metadata import version
+
 import pytest
+from packaging.version import Version
 
 import vllm
 from vllm.assets.image import ImageAsset
@@ -10,6 +13,14 @@ from vllm.platforms import current_platform
 
 from ..utils import multi_gpu_test
 
+pytestmark = pytest.mark.skipif(
+    Version("5.0") <= Version(version("transformers")),
+    reason=(
+        "MiniCPMV custom processor uses tokenizer.im_start_id which is not "
+        "available on TokenizersBackend in transformers v5.0+"
+    ),
+)
+
 MODEL_PATH = "openbmb/MiniCPM-Llama3-V-2_5"
 
 PROMPT_TEMPLATE = (

tests/model_executor/test_weight_utils.py

@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-import os
 import tempfile
 
 import huggingface_hub.constants
@@ -10,26 +9,10 @@ from huggingface_hub.utils import LocalEntryNotFoundError
 
 from vllm.model_executor.model_loader.weight_utils import (
     download_weights_from_hf,
-    enable_hf_transfer,
     maybe_remap_kv_scale_name,
 )
 
 
-def test_hf_transfer_auto_activation():
-    if "HF_HUB_ENABLE_HF_TRANSFER" in os.environ:
-        # in case it is already set, we can't test the auto activation
-        pytest.skip("HF_HUB_ENABLE_HF_TRANSFER is set, can't test auto activation")
-    enable_hf_transfer()
-    try:
-        # enable hf hub transfer if available
-        import hf_transfer  # type: ignore # noqa
-
-        HF_TRANSFER_ACTIVE = True
-    except ImportError:
-        HF_TRANSFER_ACTIVE = False
-    assert huggingface_hub.constants.HF_HUB_ENABLE_HF_TRANSFER == HF_TRANSFER_ACTIVE
-
-
 def test_download_weights_from_hf():
     with tempfile.TemporaryDirectory() as tmpdir:
         # assert LocalEntryNotFoundError error is thrown
@@ -178,5 +161,4 @@ class TestMaybeRemapKvScaleName:
 
 
 if __name__ == "__main__":
-    test_hf_transfer_auto_activation()
     test_download_weights_from_hf()

tests/models/language/generation/test_common.py

@@ -143,6 +143,11 @@ def test_models(
         # in parts of the operators
         pytest.skip(f"Skipping '{model}' model test with AITER kernel.")
 
+    if current_platform.is_cpu() and model == "TitanML/tiny-mixtral":
+        # This untrained model is sensitive to the rounding error
+        # Fuse ops to reduce bfloat16 rounding
+        monkeypatch.setenv("VLLM_CPU_CI_ENV", "0")
+
     with hf_runner(model) as hf_model:
         hf_outputs = hf_model.generate_greedy_logprobs_limit(
             example_prompts, max_tokens, num_logprobs

tests/models/language/pooling/test_colbert.py

@@ -109,6 +109,14 @@ def _load_hf_model(model_name: str, hf_spec: dict, device: torch.device):
         **extra,
     ).to(device)
     model.eval()
+
+    # Transformers 5.0 weight materialization can clear non-persistent
+    # buffers (e.g. rotary inv_freq) that were registered with
+    # persistent=False.  Re-compute them so the model produces valid output.
+    for mod in model.modules():
+        if hasattr(mod, "_compute_inv_freq") and hasattr(mod, "inv_freq"):
+            mod.inv_freq = mod._compute_inv_freq(device=device)
+
     return model
 
 

tests/models/language/pooling/test_nomic_max_model_len.py

@@ -8,7 +8,13 @@ import pytest
 from ...utils import EmbedModelInfo
 
 MODELS = [
-    EmbedModelInfo("nomic-ai/nomic-embed-text-v1"),
+    EmbedModelInfo(
+        "nomic-ai/nomic-embed-text-v1",
+        # Fixme:
+        #  Update nomic-embed code to support the latest
+        #  HF version and remove revision set.
+        revision="720244025c1a7e15661a174c63cce63c8218e52b",
+    ),
     # EmbedModelInfo("nomic-ai/nomic-embed-text-v1.5"),
     # EmbedModelInfo("nomic-ai/CodeRankEmbed"),
     EmbedModelInfo("nomic-ai/nomic-embed-text-v2-moe"),
@@ -24,7 +30,10 @@ max_model_len = int(original_max_position_embeddings * factor)
 @pytest.mark.parametrize("model_info", MODELS)
 def test_default(model_info, vllm_runner):
     with vllm_runner(
-        model_info.name, runner="pooling", max_model_len=None
+        model_info.name,
+        revision=model_info.revision,
+        runner="pooling",
+        max_model_len=None,
     ) as vllm_model:
         model_config = vllm_model.llm.llm_engine.model_config
         if model_info.name == "nomic-ai/nomic-embed-text-v2-moe":
@@ -39,7 +48,10 @@ def test_default(model_info, vllm_runner):
 def test_set_max_model_len_legal(model_info, vllm_runner):
     # set max_model_len <= 512
     with vllm_runner(
-        model_info.name, runner="pooling", max_model_len=256
+        model_info.name,
+        revision=model_info.revision,
+        runner="pooling",
+        max_model_len=256,
     ) as vllm_model:
         model_config = vllm_model.llm.llm_engine.model_config
         assert model_config.max_model_len == 256
@@ -49,11 +61,19 @@ def test_set_max_model_len_legal(model_info, vllm_runner):
         # For nomic-embed-text-v2-moe the length is set to 512
         # by sentence_bert_config.json.
         with pytest.raises(ValueError):
-            with vllm_runner(model_info.name, runner="pooling", max_model_len=1024):
+            with vllm_runner(
+                model_info.name,
+                revision=model_info.revision,
+                runner="pooling",
+                max_model_len=1024,
+            ):
                 pass
     else:
         with vllm_runner(
-            model_info.name, runner="pooling", max_model_len=1024
+            model_info.name,
+            revision=model_info.revision,
+            runner="pooling",
+            max_model_len=1024,
         ) as vllm_model:
             model_config = vllm_model.llm.llm_engine.model_config
             assert model_config.max_model_len == 1024
@@ -63,7 +83,12 @@ def test_set_max_model_len_legal(model_info, vllm_runner):
 def test_set_max_model_len_illegal(model_info, vllm_runner):
     # set max_model_len > 2048
     with pytest.raises(ValueError):
-        with vllm_runner(model_info.name, runner="pooling", max_model_len=4096):
+        with vllm_runner(
+            model_info.name,
+            revision=model_info.revision,
+            runner="pooling",
+            max_model_len=4096,
+        ):
             pass
 
     # set max_model_len > 2048 by hf_overrides
@@ -71,6 +96,7 @@ def test_set_max_model_len_illegal(model_info, vllm_runner):
     with pytest.raises(ValueError):
         with vllm_runner(
             model_info.name,
+            revision=model_info.revision,
             runner="pooling",
             max_model_len=None,
             hf_overrides=hf_overrides,
@@ -91,7 +117,11 @@ def test_use_rope_scaling_legal(model_info, vllm_runner):
     }
 
     with vllm_runner(
-        model_info.name, runner="pooling", max_model_len=None, hf_overrides=hf_overrides
+        model_info.name,
+        revision=model_info.revision,
+        runner="pooling",
+        max_model_len=None,
+        hf_overrides=hf_overrides,
     ):
         pass
 
@@ -110,6 +140,7 @@ def test_use_rope_scaling_illegal(model_info, vllm_runner):
     with pytest.raises(ValueError):
         with vllm_runner(
             model_info.name,
+            revision=model_info.revision,
             runner="pooling",
             max_model_len=max_model_len + 1,
             hf_overrides=hf_overrides,
@@ -129,6 +160,7 @@ def test_use_rope_scaling_illegal(model_info, vllm_runner):
     with pytest.raises(ValueError):
         with vllm_runner(
             model_info.name,
+            revision=model_info.revision,
             runner="pooling",
             max_model_len=None,
             hf_overrides=hf_overrides,

tests/models/language/pooling/test_splade_sparse_pooler.py

@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-import types
-
 import pytest
 import torch
 import torch.nn as nn
@@ -11,6 +9,8 @@ from vllm.model_executor.models.bert import (
     BertMLMHead,
     SPLADESparsePooler,
 )
+from vllm.pooling_params import PoolingParams
+from vllm.v1.pool.metadata import PoolingMetadata, PoolingStates
 
 # ---------------------------------------------------------------------
 # Functional test: SPLADE formula correctness (no HF download needed)
@@ -38,8 +38,12 @@ def test_splade_pooler_matches_reference_formula(B, T, H, V):
         ],
         dtype=torch.long,
     )
-    meta = types.SimpleNamespace(
-        prompt_lens=prompt_lens_tenser, prompt_token_ids=token_ids
+    meta = PoolingMetadata(
+        prompt_lens=prompt_lens_tenser,
+        prompt_token_ids=token_ids,
+        prompt_token_ids_cpu=token_ids,
+        pooling_params=[PoolingParams(task="embed")] * B,
+        pooling_states=[PoolingStates() for _ in range(B)],
     )
 
     # MLM head (prefer BertMLMHead, fallback to Linear if unavailable)

tests/models/language/pooling_mteb_test/mteb_embed_utils.py

@@ -151,6 +151,7 @@ def mteb_test_embed_models(
 
     with vllm_runner(
         model_info.name,
+        revision=model_info.revision,
         runner="pooling",
         max_model_len=model_info.max_model_len,
         **vllm_extra_kwargs,
@@ -201,6 +202,7 @@ def mteb_test_embed_models(
     if model_info.mteb_score is None:
         with hf_runner(
             model_info.name,
+            revision=model_info.revision,
             is_sentence_transformer=True,
             dtype=ci_envs.VLLM_CI_HF_DTYPE or model_info.hf_dtype,
         ) as hf_model:

tests/models/language/pooling_mteb_test/mteb_score_utils.py

@@ -241,6 +241,7 @@ def mteb_test_rerank_models(
 
     with vllm_runner(
         model_info.name,
+        revision=model_info.revision,
         runner="pooling",
         max_model_len=None,
         max_num_seqs=8,
@@ -286,7 +287,9 @@ def mteb_test_rerank_models(
     # Accelerate mteb test by setting
     # SentenceTransformers mteb score to a constant
     if model_info.mteb_score is None:
-        with hf_runner(model_info.name, dtype=model_info.hf_dtype) as hf_model:
+        with hf_runner(
+            model_info.name, revision=model_info.revision, dtype=model_info.hf_dtype
+        ) as hf_model:
             hf_model.chat_template = chat_template
             st_main_score = run_mteb_rerank(
                 hf_model,

tests/models/language/pooling_mteb_test/test_baai.py

@@ -69,7 +69,10 @@ MODELS = [
         attn_type="decoder",
         is_prefix_caching_supported=True,
         is_chunked_prefill_supported=True,
-        enable_test=True,
+        # Skip: model's custom tokenizer on HF hub is incompatible with
+        # transformers v5 (sets attrs before super().__init__, triggering
+        # AttributeError on 'verbose' in __getattr__).
+        enable_test=False,
     ),
 ]
 

tests/models/language/pooling_mteb_test/test_gte.py

@@ -72,7 +72,8 @@ MODELS = [
         attn_type="encoder_only",
         is_prefix_caching_supported=False,
         is_chunked_prefill_supported=False,
-        enable_test=True,
+        # Skip: numerical regression with transformers v5.
+        enable_test=False,
     ),
     ########## ModernBertModel
     EmbedModelInfo(

tests/models/language/pooling_mteb_test/test_jina.py

@@ -75,6 +75,10 @@ def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:
     mteb_test_rerank_models(vllm_runner, model_info)
 
 
+@pytest.mark.skip(
+    reason="jinaai/jina-embeddings-v3 custom XLMRobertaLoRA model on HF hub "
+    "is incompatible with transformers v5 (missing all_tied_weights_keys)"
+)
 @pytest.mark.parametrize("model_info", EMBEDDING_MODELS)
 @pytest.mark.parametrize("dtype", ["half"])
 @pytest.mark.parametrize("dimensions", [16, 32])

tests/models/language/pooling_mteb_test/test_nomic.py

@@ -12,6 +12,10 @@ MODELS = [
     EmbedModelInfo(
         "nomic-ai/nomic-embed-text-v1",
         architecture="NomicBertModel",
+        # Fixme:
+        #  Update nomic-embed code to support the latest
+        #  HF version and remove revision set.
+        revision="720244025c1a7e15661a174c63cce63c8218e52b",
         mteb_score=0.737568559,
         enable_test=True,
         seq_pooling_type="MEAN",

tests/models/multimodal/generation/test_common.py

@@ -186,7 +186,14 @@ VLM_TEST_SETTINGS = {
         max_num_seqs=2,
         auto_cls=AutoModel,
         hf_output_post_proc=model_utils.ultravox_trunc_hf_output,
-        marks=[pytest.mark.core_model, pytest.mark.cpu_model],
+        marks=[
+            pytest.mark.core_model,
+            pytest.mark.cpu_model,
+            # TODO: Remove skip once model has been upstreamed to Transformers
+            pytest.mark.skip(
+                reason="Custom model code is not compatible with Transformers v5"
+            ),
+        ],
     ),
     #### Transformers fallback to test
     ## To reduce test burden, we only test batching arbitrary image size
@@ -394,6 +401,22 @@ VLM_TEST_SETTINGS = {
         vllm_runner_kwargs={"mm_processor_kwargs": {"do_pan_and_scan": True}},
         patch_hf_runner=model_utils.gemma3_patch_hf_runner,
     ),
+    "gemma4": VLMTestInfo(
+        models=["google/gemma-4-E2B-it"],
+        test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE),
+        prompt_formatter=lambda img_prompt: f"<bos><|turn>user\n{img_prompt}<turn|>\n<|turn>model\n",  # noqa: E501
+        single_image_prompts=IMAGE_ASSETS.prompts(
+            {
+                "stop_sign": "<|image|>What's the content in the center of the image?",  # noqa: E501
+                "cherry_blossom": "<|image|>What is the season?",
+            }
+        ),
+        multi_image_prompt="<|image|><|image|>Describe the two images in detail.",  # noqa: E501
+        max_model_len=4096,
+        max_num_seqs=2,
+        auto_cls=AutoModelForImageTextToText,
+        vllm_runner_kwargs={"limit_mm_per_prompt": {"image": 4}},
+    ),
     "granite_vision": VLMTestInfo(
         models=["ibm-granite/granite-vision-3.3-2b"],
         test_type=(VLMTestType.IMAGE),
@@ -517,6 +540,12 @@ VLM_TEST_SETTINGS = {
         max_model_len=4096,
         use_tokenizer_eos=True,
         patch_hf_runner=model_utils.internvl_patch_hf_runner,
+        # TODO: Remove skip once model has been upstreamed to Transformers
+        marks=[
+            pytest.mark.skip(
+                reason="Custom model code tries to access data from meta-tensor"
+            )
+        ],
     ),
     "intern_vl-video": VLMTestInfo(
         models=[
@@ -529,6 +558,12 @@ VLM_TEST_SETTINGS = {
         use_tokenizer_eos=True,
         patch_hf_runner=model_utils.internvl_patch_hf_runner,
         num_logprobs=10 if current_platform.is_rocm() else 5,
+        # TODO: Remove skip once model has been upstreamed to Transformers
+        marks=[
+            pytest.mark.skip(
+                reason="Custom model code tries to access data from meta-tensor"
+            )
+        ],
     ),
     "intern_vl-hf": VLMTestInfo(
         models=["OpenGVLab/InternVL3-1B-hf"],
@@ -575,6 +610,8 @@ VLM_TEST_SETTINGS = {
         hf_model_kwargs={"device_map": "auto"},
         patch_hf_runner=model_utils.isaac_patch_hf_runner,
         image_size_factors=[(0.25,), (0.25, 0.25, 0.25), (0.25, 0.2, 0.15)],
+        # TODO: Remove skip once model has been upstreamed to Transformers
+        marks=[pytest.mark.skip(reason="Custom model imports deleted object")],  # noqa: E501
     ),
     "kimi_vl": VLMTestInfo(
         models=["moonshotai/Kimi-VL-A3B-Instruct"],
@@ -790,7 +827,12 @@ VLM_TEST_SETTINGS = {
             pytest.mark.skipif(
                 Version(TRANSFORMERS_VERSION) == Version("4.57.3"),
                 reason="This model is broken in Transformers v4.57.3",
-            )
+            ),
+            pytest.mark.skipif(
+                Version(TRANSFORMERS_VERSION) >= Version("5.0.0"),
+                reason="Model's custom code uses ROPE_INIT_FUNCTIONS"
+                "['default'] which was removed in transformers v5",
+            ),
         ],
     ),
     "phi3v": VLMTestInfo(
@@ -944,6 +986,12 @@ VLM_TEST_SETTINGS = {
             )
             for inp in custom_inputs.different_patch_input_cases_internvl()
         ],
+        # TODO: Remove skip once model has been upstreamed to Transformers
+        marks=[
+            pytest.mark.skip(
+                reason="Custom model code tries to access data from meta-tensor"
+            )
+        ],
     ),
     "llava_onevision-multiple-images": VLMTestInfo(
         models=["llava-hf/llava-onevision-qwen2-0.5b-ov-hf"],

tests/models/multimodal/generation/test_nemotron_parse.py

@@ -103,6 +103,10 @@ def run_test(
         )
 
 
+@pytest.mark.skip(
+    reason="Model's custom MBart decoder has head count mismatch with "
+    "transformers v5's GQA-aware cross-attention (8 vs 16 heads)"
+)
 @pytest.mark.parametrize("model", ["nvidia/NVIDIA-Nemotron-Parse-v1.1"])
 @pytest.mark.parametrize("dtype", ["bfloat16"])
 @pytest.mark.parametrize("num_logprobs", [5])

tests/models/multimodal/generation/test_phi4siglip.py

@@ -0,0 +1,198 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections.abc import Sequence
+from importlib.metadata import version
+
+import pytest
+import regex as re
+from packaging.version import Version
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from vllm.logprobs import SampleLogprobs
+from vllm.multimodal.image import rescale_image_size
+
+from ....conftest import (
+    IMAGE_ASSETS,
+    HfRunner,
+    PromptImageInput,
+    VllmRunner,
+)
+from ....utils import multi_gpu_test
+from ...utils import check_logprobs_close
+
+pytestmark = pytest.mark.skipif(
+    Version("5.0") <= Version(version("transformers")),
+    reason=(
+        "vllm upgraded transformers above v5.4 where HF model custom code uses siglip2 "
+        "internals (filter_out_non_signature_kwargs) removed by "
+        "huggingface/transformers#43514"
+    ),
+)
+
+MODEL_ID = "microsoft/Phi-4-reasoning-vision-15B"
+
+HF_IMAGE_PROMPTS = IMAGE_ASSETS.prompts(
+    {
+        "stop_sign": "<|user|>\n<image>\nWhat's the content of the image?<|end|>\n<|assistant|>\n",  # noqa: E501
+        "cherry_blossom": "<|user|>\n<image>\nPlease infer the season with reason in details.<|end|>\n<|assistant|>\n",  # noqa: E501
+    }
+)
+HF_MULTIIMAGE_IMAGE_PROMPT = (
+    "<|user|>\n<image>\n<image>\nDescribe these images.<|end|>\n<|assistant|>\n"  # noqa: E501
+)
+
+DTYPE = "half"
+MAX_TOKENS = 128
+NUM_LOGPROBS = 10
+
+
+def vllm_to_hf_output(
+    vllm_output: tuple[list[int], str, SampleLogprobs | None], model: str
+):
+    """Sanitize vllm output to be comparable with hf output."""
+    _, output_str, out_logprobs = vllm_output
+
+    output_str_without_image = re.sub(r"(<image>)+", "", output_str)
+    if output_str_without_image and output_str_without_image[0] == " ":
+        output_str_without_image = output_str_without_image[1:]
+
+    hf_output_str = output_str_without_image + "<|end|><|endoftext|>"
+
+    tokenizer = AutoTokenizer.from_pretrained(model, trust_remote_code=True)
+    hf_output_ids = tokenizer.encode(output_str_without_image)
+    if hf_output_ids and hf_output_ids[0] == tokenizer.bos_token_id:
+        hf_output_ids = hf_output_ids[1:]
+
+    return hf_output_ids, hf_output_str, out_logprobs
+
+
+def _build_single_image_inputs(
+    image_assets,
+) -> list[tuple[list[str], PromptImageInput]]:
+    """Build single-image inputs for all size_factors at once."""
+    images = [asset.pil_image for asset in image_assets]
+    all_inputs: list[tuple[list[str], PromptImageInput]] = []
+    for size_factors in [[1.0], [0.25, 0.5, 1.0]]:
+        for image, prompt in zip(images, HF_IMAGE_PROMPTS):
+            all_inputs.append(
+                (
+                    [prompt for _ in size_factors],
+                    [rescale_image_size(image, f) for f in size_factors],
+                )
+            )
+    return all_inputs
+
+
+def _build_multi_image_inputs(
+    image_assets,
+) -> list[tuple[list[str], PromptImageInput]]:
+    """Build multi-image inputs for all size_factors at once."""
+    images = [asset.pil_image for asset in image_assets]
+    all_inputs: list[tuple[list[str], PromptImageInput]] = []
+    for size_factors in [[0.5], [0.15, 0.30]]:
+        all_inputs.append(
+            (
+                [HF_MULTIIMAGE_IMAGE_PROMPT for _ in size_factors],
+                [
+                    [rescale_image_size(image, factor) for image in images]
+                    for factor in size_factors
+                ],
+            )
+        )
+    return all_inputs
+
+
+def _run_and_compare(
+    hf_runner: type[HfRunner],
+    vllm_runner: type[VllmRunner],
+    all_inputs: Sequence[tuple[list[str], PromptImageInput]],
+    model: str,
+    max_model_len: int,
+    max_num_seqs: int,
+    mm_limit: int,
+    gpu_memory_utilization: float,
+):
+    """Load each runner once, run all inputs, then compare."""
+    # NOTE: run vLLM first, then HF.  vLLM needs a fresh process without
+    # cuda initialization; running HF first would break the multiprocessing
+    # backend with fork method.
+    with vllm_runner(
+        model,
+        runner="generate",
+        max_model_len=max_model_len,
+        max_num_seqs=max_num_seqs,
+        gpu_memory_utilization=gpu_memory_utilization,
+        dtype=DTYPE,
+        limit_mm_per_prompt={"image": mm_limit},
+        tensor_parallel_size=2,
+        trust_remote_code=True,
+        enforce_eager=True,
+    ) as vllm_model:
+        vllm_outputs_per_case = [
+            vllm_model.generate_greedy_logprobs(
+                prompts,
+                MAX_TOKENS,
+                num_logprobs=NUM_LOGPROBS,
+                images=images,
+            )
+            for prompts, images in all_inputs
+        ]
+
+    hf_model_kwargs = {"_attn_implementation": "sdpa", "device_map": "auto"}
+    with hf_runner(
+        model,
+        dtype=DTYPE,
+        model_kwargs=hf_model_kwargs,
+        auto_cls=AutoModelForCausalLM,
+        trust_remote_code=True,
+    ) as hf_model:
+        hf_outputs_per_case = [
+            hf_model.generate_greedy_logprobs_limit(
+                prompts,
+                MAX_TOKENS,
+                num_logprobs=NUM_LOGPROBS,
+                images=images,
+            )
+            for prompts, images in all_inputs
+        ]
+
+    for hf_outputs, vllm_outputs in zip(hf_outputs_per_case, vllm_outputs_per_case):
+        check_logprobs_close(
+            outputs_0_lst=hf_outputs,
+            outputs_1_lst=vllm_outputs,
+            name_0="hf",
+            name_1="vllm",
+        )
+
+
+@multi_gpu_test(num_gpus=2)
+@pytest.mark.parametrize("model", [MODEL_ID])
+def test_models(hf_runner, vllm_runner, image_assets, model) -> None:
+    all_inputs = _build_single_image_inputs(image_assets)
+    _run_and_compare(
+        hf_runner,
+        vllm_runner,
+        all_inputs,
+        model,
+        max_model_len=8192,
+        max_num_seqs=2,
+        mm_limit=1,
+        gpu_memory_utilization=0.80,
+    )
+
+
+@multi_gpu_test(num_gpus=2)
+@pytest.mark.parametrize("model", [MODEL_ID])
+def test_multi_images_models(hf_runner, vllm_runner, image_assets, model) -> None:
+    all_inputs = _build_multi_image_inputs(image_assets)
+    _run_and_compare(
+        hf_runner,
+        vllm_runner,
+        all_inputs,
+        model,
+        max_model_len=8192,
+        max_num_seqs=2,
+        mm_limit=2,
+        gpu_memory_utilization=0.80,
+    )

tests/models/multimodal/generation/test_voxtral.py

@@ -149,6 +149,10 @@ def test_online_serving(vllm_runner, audio_assets: AudioTestAssets):
     )
 
 
+@pytest.mark.skip(
+    reason="VoxtralProcessor.apply_chat_template() in transformers v5 "
+    "doesn't resolve chat_template=None to the default template"
+)
 def test_hf_reference(hf_runner, vllm_runner, audio_assets: AudioTestAssets):
     """Compare vLLM Mistral-format output against HF Transformers reference.
 

tests/models/multimodal/generation/vlm_utils/core.py

@@ -80,6 +80,11 @@ def run_test(
     if vllm_runner_kwargs:
         vllm_runner_kwargs_.update(vllm_runner_kwargs)
 
+    # Avoid passing limit_mm_per_prompt twice when vllm_runner_kwargs
+    # already contains it (e.g. gemma4 sets it via vllm_runner_kwargs).
+    if "limit_mm_per_prompt" in vllm_runner_kwargs_:
+        limit_mm_per_prompt = vllm_runner_kwargs_.pop("limit_mm_per_prompt")
+
     with vllm_runner(
         model,
         max_model_len=max_model_len,

tests/models/multimodal/pooling/test_colmodernvbert.py

@@ -15,6 +15,10 @@ from vllm.entrypoints.pooling.score.utils import compute_maxsim_score
 MODEL_NAME = "ModernVBERT/colmodernvbert-merged"
 COLBERT_DIM = 128
 DTYPE = "half"
+# Fixme:
+#  Update colmodernvbert code to support the latest HF version
+#  and remove revision set.
+REVISION = "4a0a9f3ac7a7992fec410bfa8e3d080ac9a5bcee"
 
 
 # -----------------------------------------------------------------------
@@ -26,6 +30,7 @@ def test_colmodernvbert_text_token_embed(vllm_runner):
     """Text query produces per-token embeddings with shape (seq_len, 128)."""
     with vllm_runner(
         MODEL_NAME,
+        revision=REVISION,
         runner="pooling",
         dtype=DTYPE,
         enforce_eager=True,
@@ -49,6 +54,7 @@ def test_colmodernvbert_text_relevance_ordering(vllm_runner):
 
     with vllm_runner(
         MODEL_NAME,
+        revision=REVISION,
         runner="pooling",
         dtype=DTYPE,
         enforce_eager=True,
@@ -66,6 +72,7 @@ def test_colmodernvbert_text_late_interaction(vllm_runner):
 
     with vllm_runner(
         MODEL_NAME,
+        revision=REVISION,
         runner="pooling",
         dtype=DTYPE,
         enforce_eager=True,
@@ -92,6 +99,7 @@ def test_colmodernvbert_image_token_embed(vllm_runner, image_assets):
     """Image input produces per-token embeddings including vision tokens."""
     with vllm_runner(
         MODEL_NAME,
+        revision=REVISION,
         runner="pooling",
         dtype=DTYPE,
         enforce_eager=True,

tests/models/multimodal/pooling/test_colqwen3.py

@@ -22,6 +22,11 @@ from vllm.entrypoints.pooling.score.utils import ScoreMultiModalParam
 
 from ....conftest import VllmRunner
 
+pytestmark = pytest.mark.skip(
+    reason="ColQwen3 model's weight tying is incompatible with "
+    "transformers v5 (missing all_tied_weights_keys)"
+)
+
 MODELS = [
     "TomoroAI/tomoro-colqwen3-embed-4b",
     "OpenSearch-AI/Ops-Colqwen3-4B",

tests/models/multimodal/pooling/test_intern_vit.py

@@ -11,6 +11,11 @@ from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
 
 from ....conftest import ImageTestAssets
 
+pytestmark = pytest.mark.skip(
+    reason="InternVisionModel's custom code is incompatible with "
+    "transformers v5 (missing all_tied_weights_keys)"
+)
+
 # we use snapshot_download to prevent conflicts between
 # dynamic_module and trust_remote_code for hf_runner
 DOWNLOAD_PATTERN = ["*.json", "*.py", "*.safetensors", "*.txt", "*.model"]

tests/models/multimodal/pooling/test_jinavl_reranker.py

@@ -15,6 +15,11 @@ from vllm.entrypoints.pooling.score.utils import ScoreMultiModalParam
 
 from ....conftest import HfRunner, VllmRunner
 
+pytestmark = pytest.mark.skip(
+    reason="jinaai/jina-reranker-m0 custom code is incompatible with "
+    "transformers v5 (missing all_tied_weights_keys)"
+)
+
 MODELS = ["jinaai/jina-reranker-m0"]
 
 MM_PROCESSOR_KWARGS = {

tests/models/multimodal/processing/test_gemma4.py

@@ -0,0 +1,44 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+from vllm.multimodal import MULTIMODAL_REGISTRY
+
+from ....conftest import ImageTestAssets
+from ...utils import build_model_context
+
+# TODO: to be updated to "google/gemma-4-e2b-it" once the models are available
+GEMMA4_MODEL_ID = "google/gemma-4-E2B-it"
+
+
+@pytest.mark.parametrize("model_id", [GEMMA4_MODEL_ID])
+def test_limit_mm_per_prompt(
+    image_assets: ImageTestAssets,
+    model_id: str,
+):
+    """Test that limit_mm_per_prompt accurately restricts multiple images."""
+    # We only allow 1 image
+    ctx = build_model_context(
+        model_id,
+        mm_processor_kwargs={},
+        limit_mm_per_prompt={"image": 1},
+    )
+    processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
+
+    # Provide 2 images in the prompt
+    prompt = "<image><image>"
+    # image_assets usually has multiple images
+    images = [asset.pil_image for asset in image_assets][:2]
+    if len(images) < 2:
+        images = [images[0], images[0]]
+
+    mm_data = {"image": images}
+
+    # Expect ValueError when exceeding limit
+    with pytest.raises(ValueError, match="At most 1 image"):
+        processor(
+            prompt,
+            mm_items=processor.info.parse_mm_data(mm_data),
+            hf_processor_mm_kwargs={},
+        )

tests/models/multimodal/processing/test_musicflamingo.py

@@ -17,11 +17,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from importlib.metadata import version
 from unittest.mock import MagicMock
 
 import numpy as np
 import pytest
 import torch
+from packaging.version import Version
 from transformers import PretrainedConfig
 
 from tests.models.registry import HF_EXAMPLE_MODELS
@@ -122,6 +124,11 @@ def test_musicflamingo_dummy_text_uses_plain_audio_tokens(mock_ctx):
     assert builder.get_dummy_text({"audio": 2}) == "<sound><sound>"
 
 
+@pytest.mark.skipif(
+    Version(version("transformers")) >= Version("5.5"),
+    reason="transformers v5.5 added native MusicFlamingoForConditionalGeneration "
+    "with a different get_audio_features signature (requires input_ids)",
+)
 def test_musicflamingo_audio_feature_pipeline_matches_hf_small_config():
     from transformers.models.musicflamingo import (
         modeling_musicflamingo as hf_musicflamingo_modeling,

tests/models/registry.py

@@ -277,6 +277,10 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
         "google/gemma-2-9b", extras={"tiny": "google/gemma-2-2b-it"}
     ),
     "Gemma3ForCausalLM": _HfExamplesInfo("google/gemma-3-1b-it"),
+    "Gemma4ForCausalLM": _HfExamplesInfo(
+        "google/gemma-4-E2B-it",
+        min_transformers_version="5.0.0",
+    ),
     "Gemma3nForCausalLM": _HfExamplesInfo("google/gemma-3n-E2B-it"),
     "GlmForCausalLM": _HfExamplesInfo("zai-org/glm-4-9b-chat-hf"),
     "Glm4ForCausalLM": _HfExamplesInfo("zai-org/GLM-4-9B-0414"),
@@ -330,7 +334,15 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
         "internlm/internlm2-chat-7b", trust_remote_code=True
     ),
     "InternLM2VEForCausalLM": _HfExamplesInfo(
-        "OpenGVLab/Mono-InternVL-2B", trust_remote_code=True
+        "OpenGVLab/Mono-InternVL-2B",
+        trust_remote_code=True,
+        max_transformers_version="4.57",
+        transformers_version_reason={
+            "vllm": (
+                "Custom config cannot be loaded with Transformers "
+                "v5 because `vision_config` is not always set"
+            )
+        },
     ),
     "InternLM3ForCausalLM": _HfExamplesInfo(
         "internlm/internlm3-8b-instruct", trust_remote_code=True
@@ -465,6 +477,13 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
     "Plamo2ForCausalLM": _HfExamplesInfo(
         "pfnet/plamo-2-1b",
         trust_remote_code=True,
+        max_transformers_version="4.57",
+        transformers_version_reason={
+            "hf": (
+                "Custom model code uses `_tied_weight_keys: list[str]` but "
+                "Transformers v5 now expects `_tied_weight_keys: dict[str, str]`"
+            )
+        },
     ),
     "Plamo3ForCausalLM": _HfExamplesInfo(
         "pfnet/plamo-3-nict-2b-base",
@@ -505,6 +524,13 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
         trust_remote_code=True,
         max_model_len=4096,
         is_available_online=True,
+        max_transformers_version="5.3",
+        transformers_version_reason={
+            "vllm": (
+                "vllm upgraded transformers above v5.4 where "
+                "validate_rope() no longer accepts ignore_keys param"
+            )
+        },
     ),
     "SeedOssForCausalLM": _HfExamplesInfo(
         "ByteDance-Seed/Seed-OSS-36B-Instruct",
@@ -540,6 +566,11 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
         "xverse/XVERSE-7B-Chat",
         tokenizer="meta-llama/Llama-2-7b",
         trust_remote_code=True,
+        max_transformers_version="4.57",
+        transformers_version_reason={
+            "vllm": "XVERSE tokenizer is incompatible with transformers v5 "
+            "(add_prefix_space / prepend_scheme mismatch).",
+        },
     ),
     "Zamba2ForCausalLM": _HfExamplesInfo("Zyphra/Zamba2-7B-instruct"),
     "MiMoForCausalLM": _HfExamplesInfo("XiaomiMiMo/MiMo-7B-RL", trust_remote_code=True),
@@ -636,6 +667,7 @@ _LATE_INTERACTION_EXAMPLE_MODELS = {
     # [Multimodal]
     "ColModernVBertForRetrieval": _HfExamplesInfo(
         "ModernVBERT/colmodernvbert-merged",
+        revision="4a0a9f3ac7a7992fec410bfa8e3d080ac9a5bcee",
     ),
     "ColPaliForRetrieval": _HfExamplesInfo("vidore/colpali-v1.3-hf"),
     "ColQwen3": _HfExamplesInfo(
@@ -749,10 +781,18 @@ _MULTIMODAL_EXAMPLE_MODELS = {
     # [Decoder-only]
     "AriaForConditionalGeneration": _HfExamplesInfo("rhymes-ai/Aria"),
     "AudioFlamingo3ForConditionalGeneration": _HfExamplesInfo(
-        "nvidia/audio-flamingo-3-hf", min_transformers_version="5.0.0"
+        "nvidia/audio-flamingo-3-hf",
+        min_transformers_version="5.3.0",
+        transformers_version_reason={
+            "vllm": "Needs https://github.com/huggingface/transformers/pull/43538"
+        },
     ),
     "MusicFlamingoForConditionalGeneration": _HfExamplesInfo(
-        "nvidia/music-flamingo-2601-hf", min_transformers_version="5.3.0"
+        "nvidia/music-flamingo-2601-hf",
+        min_transformers_version="5.3.0",
+        transformers_version_reason={
+            "vllm": "Needs https://github.com/huggingface/transformers/pull/43538"
+        },
     ),
     "AyaVisionForConditionalGeneration": _HfExamplesInfo("CohereLabs/aya-vision-8b"),
     "BagelForConditionalGeneration": _HfExamplesInfo("ByteDance-Seed/BAGEL-7B-MoT"),
@@ -795,15 +835,40 @@ _MULTIMODAL_EXAMPLE_MODELS = {
     ),
     "FireRedASR2ForConditionalGeneration": _HfExamplesInfo(
         "allendou/FireRedASR2-LLM-vllm",
+        trust_remote_code=True,
+        max_transformers_version="5.1",
+        transformers_version_reason={
+            "vllm": "Incompatible with transformers v5.2+ "
+            "(dict object has no attribute '__name__').",
+        },
+    ),
+    "FireRedLIDForConditionalGeneration": _HfExamplesInfo(
+        "PatchyTisa/FireRedLID-vllm",
+        trust_remote_code=True,
+        max_transformers_version="5.1",
+        transformers_version_reason={
+            "vllm": "Incompatible with transformers v5.2+ "
+            "(dict object has no attribute '__name__').",
+        },
     ),
     "FunASRForConditionalGeneration": _HfExamplesInfo(
         "allendou/Fun-ASR-Nano-2512-vllm",
+        trust_remote_code=True,
+        max_transformers_version="5.1",
+        transformers_version_reason={
+            "vllm": "Incompatible with transformers v5.2+ "
+            "(dict object has no attribute '__name__').",
+        },
     ),
     "FunAudioChatForConditionalGeneration": _HfExamplesInfo(
         "funaudiochat", is_available_online=False
     ),
     "FuyuForCausalLM": _HfExamplesInfo("adept/fuyu-8b"),
     "Gemma3ForConditionalGeneration": _HfExamplesInfo("google/gemma-3-4b-it"),
+    "Gemma4ForConditionalGeneration": _HfExamplesInfo(
+        "google/gemma-4-E2B-it",
+        min_transformers_version="5.5.0",
+    ),
     "Gemma3nForConditionalGeneration": _HfExamplesInfo("google/gemma-3n-E2B-it"),
     "GlmAsrForConditionalGeneration": _HfExamplesInfo(
         "zai-org/GLM-ASR-Nano-2512",
@@ -835,6 +900,13 @@ _MULTIMODAL_EXAMPLE_MODELS = {
     "HCXVisionForCausalLM": _HfExamplesInfo(
         "naver-hyperclovax/HyperCLOVAX-SEED-Vision-Instruct-3B",
         trust_remote_code=True,
+        max_transformers_version="4.57",
+        transformers_version_reason={
+            "vllm": (
+                "Custom config cannot be loaded with Transformers "
+                "v5 because `text_config` is not always set"
+            )
+        },
     ),
     "HCXVisionV2ForCausalLM": _HfExamplesInfo(
         "naver-hyperclovax/HyperCLOVAX-SEED-Think-32B",
@@ -854,7 +926,12 @@ _MULTIMODAL_EXAMPLE_MODELS = {
         extras={"0.2-2B-Preview": "PerceptronAI/Isaac-0.2-2B-Preview"},
     ),
     "InternS1ForConditionalGeneration": _HfExamplesInfo(
-        "internlm/Intern-S1", trust_remote_code=True
+        "internlm/Intern-S1",
+        trust_remote_code=True,
+        max_transformers_version="4.57",
+        transformers_version_reason={
+            "vllm": "Custom tokenizer code is not compatible with Transformers v5."
+        },
     ),
     "InternS1ProForConditionalGeneration": _HfExamplesInfo(
         "internlm/Intern-S1-Pro",
@@ -943,7 +1020,14 @@ _MULTIMODAL_EXAMPLE_MODELS = {
     "MiDashengLMModel": _HfExamplesInfo(
         "mispeech/midashenglm-7b", trust_remote_code=True
     ),
-    "MiniCPMO": _HfExamplesInfo("openbmb/MiniCPM-o-2_6", trust_remote_code=True),
+    "MiniCPMO": _HfExamplesInfo(
+        "openbmb/MiniCPM-o-2_6",
+        trust_remote_code=True,
+        max_transformers_version="4.57",
+        transformers_version_reason={
+            "hf": "Custom processor code is not compatible with Transformers v5."
+        },
+    ),
     "MiniCPMV": _HfExamplesInfo(
         "openbmb/MiniCPM-Llama3-V-2_5",
         extras={
@@ -951,6 +1035,13 @@ _MULTIMODAL_EXAMPLE_MODELS = {
             "4.0": "openbmb/MiniCPM-V-4",
             "4.5": "openbmb/MiniCPM-V-4_5",
         },
+        max_transformers_version="4.57",
+        transformers_version_reason={
+            "vllm": (
+                "MiniCPMVBatchFeature is incompatible with its base class in "
+                "Transformers v5. See https://huggingface.co/openbmb/MiniCPM-Llama3-V-2_5/discussions/78"
+            )
+        },
         trust_remote_code=True,
     ),
     "MiniMaxVL01ForConditionalGeneration": _HfExamplesInfo(
@@ -987,13 +1078,25 @@ _MULTIMODAL_EXAMPLE_MODELS = {
         "nano_vl_dummy", is_available_online=False, trust_remote_code=True
     ),
     "OpenCUAForConditionalGeneration": _HfExamplesInfo(
-        "xlangai/OpenCUA-7B", trust_remote_code=True
+        "xlangai/OpenCUA-7B",
+        trust_remote_code=True,
+        max_transformers_version="4.57",
+        transformers_version_reason={
+            "vllm": "Tokenizer cannot be initialised in Transformers v5."
+        },
     ),
     "OpenPanguVLForConditionalGeneration": _HfExamplesInfo(
         "FreedomIntelligence/openPangu-VL-7B",
         trust_remote_code=True,
         max_model_len=4096,
         enforce_eager=True,
+        max_transformers_version="4.57",
+        transformers_version_reason={
+            "vllm": (
+                "OpenPanguVLVideoProcessorInitKwargs does not specify total=False, "
+                "making all kwargs required. See https://huggingface.co/FreedomIntelligence/openPangu-VL-7B/discussions/2"
+            )
+        },
     ),
     "Ovis": _HfExamplesInfo(
         "AIDC-AI/Ovis2-1B",
@@ -1005,12 +1108,24 @@ _MULTIMODAL_EXAMPLE_MODELS = {
             "1.6-gemma": "AIDC-AI/Ovis1.6-Gemma2-9B",
         },
     ),
-    "Ovis2_5": _HfExamplesInfo("AIDC-AI/Ovis2.5-2B", trust_remote_code=True),
+    "Ovis2_5": _HfExamplesInfo(
+        "AIDC-AI/Ovis2.5-2B",
+        trust_remote_code=True,
+        max_transformers_version="4.57",
+        transformers_version_reason={
+            "vllm": "Custom processor code is not compatible with Transformers v5."
+        },
+    ),
     "Ovis2_6ForCausalLM": _HfExamplesInfo(
         "AIDC-AI/Ovis2.6-2B", is_available_online=False, trust_remote_code=True
     ),
     "Ovis2_6_MoeForCausalLM": _HfExamplesInfo(
-        "AIDC-AI/Ovis2.6-30B-A3B", trust_remote_code=True
+        "AIDC-AI/Ovis2.6-30B-A3B",
+        trust_remote_code=True,
+        max_transformers_version="4.57",
+        transformers_version_reason={
+            "vllm": "Custom processor code is not compatible with Transformers v5."
+        },
     ),
     "PaddleOCRVLForConditionalGeneration": _HfExamplesInfo(
         "PaddlePaddle/PaddleOCR-VL",
@@ -1029,6 +1144,19 @@ _MULTIMODAL_EXAMPLE_MODELS = {
         },  # noqa: E501
         extras={"phi3.5": "microsoft/Phi-3.5-vision-instruct"},
     ),
+    "Phi4ForCausalLMV": _HfExamplesInfo(
+        "microsoft/Phi-4-reasoning-vision-15B",
+        trust_remote_code=True,
+        max_transformers_version="5.3",
+        transformers_version_reason={
+            "vllm": (
+                "vllm upgraded transformers above v5.4 where HF model "
+                "custom code uses siglip2 internals "
+                "(filter_out_non_signature_kwargs) removed "
+                "by huggingface/transformers#43514"
+            )
+        },
+    ),
     "Phi4MMForCausalLM": _HfExamplesInfo(
         "microsoft/Phi-4-multimodal-instruct", trust_remote_code=True
     ),
@@ -1124,6 +1252,14 @@ _MULTIMODAL_EXAMPLE_MODELS = {
             "architectures": ["Tarsier2ForConditionalGeneration"],
             "model_type": "tarsier2",
         },
+        max_transformers_version="5.3",
+        transformers_version_reason={
+            "vllm": (
+                "Qwen2VLConfig was split into Qwen2VLConfig + "
+                "Qwen2VLTextConfig in transformers v5, breaking "
+                "attribute access (num_attention_heads, hidden_size, etc.)"
+            )
+        },
     ),
     "VoxtralForConditionalGeneration": _HfExamplesInfo(
         "mistralai/Voxtral-Mini-3B-2507",

tests/models/utils.py

@@ -375,6 +375,7 @@ def softmax(data):
 @dataclass
 class ModelInfo:
     name: str
+    revision: str | None = None
     architecture: str = ""
     dtype: str = "auto"
     max_model_len: int | None = None
@@ -468,7 +469,16 @@ def dummy_hf_overrides(
     else:
         # Use minimal layers for testing
         num_layers = 1
-        num_hidden_layers = 3 if model_arch == "Gemma3nForConditionalGeneration" else 1
+        num_hidden_layers = (
+            3
+            if model_arch
+            in (
+                "Gemma3nForConditionalGeneration",
+                "Gemma4ForCausalLM",
+                "Gemma4ForConditionalGeneration",
+            )
+            else 1
+        )
 
     update_dict = {
         "num_layers": num_layers,

tests/multimodal/media/test_video.py

@@ -239,6 +239,17 @@ def test_video_media_io_backend_env_var_fallback(monkeypatch: pytest.MonkeyPatch
         assert metadata_missing["video_backend"] == "test_video_backend_override_2"
 
 
+def _make_jpeg_b64_frames(n: int, width: int = 8, height: int = 8) -> list[str]:
+    """Return *n* tiny base64-encoded JPEG frames."""
+    frames: list[str] = []
+    for i in range(n):
+        img = Image.new("RGB", (width, height), color=(i % 256, 0, 0))
+        buf = io.BytesIO()
+        img.save(buf, format="JPEG")
+        frames.append(pybase64.b64encode(buf.getvalue()).decode("ascii"))
+    return frames
+
+
 def test_load_base64_jpeg_returns_metadata():
     """Regression test: load_base64 with video/jpeg must return metadata.
 
@@ -248,16 +259,8 @@ def test_load_base64_jpeg_returns_metadata():
     """
 
     num_test_frames = 3
-    frame_width, frame_height = 8, 8
-
-    # Build a few tiny JPEG frames and base64-encode them
-    b64_frames = []
-    for i in range(num_test_frames):
-        img = Image.new("RGB", (frame_width, frame_height), color=(i * 80, 0, 0))
-        buf = io.BytesIO()
-        img.save(buf, format="JPEG")
-        b64_frames.append(pybase64.b64encode(buf.getvalue()).decode("ascii"))
 
+    b64_frames = _make_jpeg_b64_frames(num_test_frames)
     data = ",".join(b64_frames)
 
     imageio = ImageMediaIO()
@@ -287,3 +290,52 @@ def test_load_base64_jpeg_returns_metadata():
     # Default fps=1 → duration == num_frames
     assert metadata["fps"] == 1.0
     assert metadata["duration"] == float(num_test_frames)
+
+
+def test_load_base64_jpeg_enforces_num_frames_limit():
+    """Frames beyond num_frames must be truncated in the video/jpeg path.
+
+    Without the limit an attacker can send thousands of base64 JPEG frames
+    in a single request and exhaust server memory (OOM).
+    """
+    num_frames_limit = 4
+    sent_frames = 20
+
+    b64_frames = _make_jpeg_b64_frames(sent_frames)
+    data = ",".join(b64_frames)
+
+    imageio = ImageMediaIO()
+    videoio = VideoMediaIO(imageio, num_frames=num_frames_limit)
+    frames, metadata = videoio.load_base64("video/jpeg", data)
+
+    assert frames.shape[0] == num_frames_limit
+    assert metadata["total_num_frames"] == num_frames_limit
+    assert metadata["frames_indices"] == list(range(num_frames_limit))
+
+
+def test_load_base64_jpeg_no_limit_when_num_frames_negative():
+    """When num_frames is -1, all frames should be loaded without truncation."""
+    sent_frames = 10
+
+    b64_frames = _make_jpeg_b64_frames(sent_frames)
+    data = ",".join(b64_frames)
+
+    imageio = ImageMediaIO()
+    videoio = VideoMediaIO(imageio, num_frames=-1)
+    frames, metadata = videoio.load_base64("video/jpeg", data)
+
+    assert frames.shape[0] == sent_frames
+    assert metadata["total_num_frames"] == sent_frames
+    assert metadata["frames_indices"] == list(range(sent_frames))
+
+
+def test_load_base64_jpeg_raises_on_zero_num_frames():
+    """num_frames=0 is invalid and should raise ValueError."""
+    b64_frames = _make_jpeg_b64_frames(3)
+    data = ",".join(b64_frames)
+
+    imageio = ImageMediaIO()
+    videoio = VideoMediaIO(imageio, num_frames=0)
+
+    with pytest.raises(ValueError, match="num_frames must be greater than 0 or -1"):
+        videoio.load_base64("video/jpeg", data)

tests/reasoning/test_gemma4_reasoning_parser.py

@@ -0,0 +1,275 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+from tests.reasoning.utils import run_reasoning_extraction
+from vllm.entrypoints.openai.chat_completion.protocol import (
+    ChatCompletionRequest,
+)
+from vllm.reasoning import ReasoningParser, ReasoningParserManager
+
+# Using mistral tokenizer as a generic mock since the actual model is not on HF
+from vllm.tokenizers.registry import get_tokenizer
+
+parser_name = "gemma4"
+
+
+@pytest.fixture(scope="module")
+def generic_tokenizer():
+    return get_tokenizer("google/gemma-4-E2B-it")
+
+
+INVALID_SIMPLE_NONSTREAMING = {
+    "output": "This is a reasoning section<channel|>This is the rest",
+    "reasoning": "This is a reasoning section",
+    "content": "This is the rest",
+    "is_reasoning_end": True,
+}
+INVALID_SIMPLE_STREAMING = {
+    "output": "This is a reasoning section<channel|>This is the rest",
+    "reasoning": None,
+    "content": "This is a reasoning sectionThis is the rest",
+    "is_reasoning_end": True,
+}
+INVALID_COMPLETE_NONSTREAMING = {
+    "output": "This is a reasoning section<channel|>",
+    "reasoning": "This is a reasoning section",
+    "content": None,
+    "is_reasoning_end": True,
+}
+INVALID_COMPLETE_STREAMING = {
+    "output": "This is a reasoning section<channel|>",
+    "reasoning": None,
+    "content": "This is a reasoning section",
+    "is_reasoning_end": True,
+}
+NO_CONTENT = {
+    "output": "<|channel>This is reasoning",
+    "reasoning": "This is reasoning",
+    "content": None,
+    "is_reasoning_end": False,
+}
+NO_REASONING = {
+    "output": "This is content",
+    "reasoning": None,
+    "content": "This is content",
+    "is_reasoning_end": False,
+}
+REASONING_WITH_CHANNEL = {
+    "output": "<|channel>This is a reasoning section<channel|>This is the rest",
+    "reasoning": "This is a reasoning section",
+    "content": "This is the rest",
+    "is_reasoning_end": True,
+}
+COMPLETE_REASONING_WITH_CHANNEL = {
+    "output": "<|channel>This is a reasoning section<channel|>",
+    "reasoning": "This is a reasoning section",
+    "content": None,
+    "is_reasoning_end": True,
+}
+MULTIPLE_LINES_WITH_CHANNEL = {
+    "output": "<|channel>This\nThat<channel|>This is the rest\nThat",
+    "reasoning": "This\nThat",
+    "content": "This is the rest\nThat",
+    "is_reasoning_end": True,
+}
+CHANNEL_NO_END = {
+    "output": "<|channel>This is a reasoning section",
+    "reasoning": "This is a reasoning section",
+    "content": None,
+    "is_reasoning_end": False,
+}
+EMPTY = {
+    "output": "",
+    "reasoning": None,
+    "content": "",
+    "is_reasoning_end": False,
+}
+NEW_LINE_NONSTREAMING = {
+    "output": (
+        "Before\n<|channel>This is a reasoning section<channel|>\nThis is the rest"
+    ),
+    "reasoning": "This is a reasoning section",
+    "content": "\nThis is the rest",
+    "is_reasoning_end": True,
+}
+NEW_LINE_STREAMING = {
+    "output": (
+        "Before\n<|channel>This is a reasoning section<channel|>\nThis is the rest"
+    ),
+    "reasoning": "This is a reasoning section",
+    "content": "Before\n\nThis is the rest",
+    "is_reasoning_end": True,
+}
+
+THOUGHT_PREFIX = {
+    "output": "<|channel>thought\nActual reasoning here<channel|>Final answer",
+    "reasoning": "Actual reasoning here",
+    "content": "Final answer",
+    "is_reasoning_end": True,
+}
+THOUGHT_PREFIX_ONLY = {
+    "output": "<|channel>thought\n<channel|>",
+    "reasoning": "",
+    "content": None,
+    "is_reasoning_end": True,
+}
+THOUGHT_PREFIX_MULTILINE = {
+    "output": "<|channel>thought\nLine1\nLine2<channel|>Answer",
+    "reasoning": "Line1\nLine2",
+    "content": "Answer",
+    "is_reasoning_end": True,
+}
+# "thousand" starts like "thought" but diverges — exercises Case 2→3 in streaming.
+THOUGHT_PREFIX_DIVERGE = {
+    "output": "<|channel>thousand reasons<channel|>Done",
+    "reasoning": "thousand reasons",
+    "content": "Done",
+    "is_reasoning_end": True,
+}
+# The model isn't reasoning if we're generating tool calls.
+TOOL_CALL_STARTED = {
+    "output": "<|tool_call>",
+    "reasoning": None,
+    "content": "<|tool_call>",
+    "is_reasoning_end": True,
+}
+
+TEST_CASES = [
+    pytest.param(False, INVALID_SIMPLE_NONSTREAMING, id="invalid_simple"),
+    pytest.param(True, INVALID_SIMPLE_STREAMING, id="invalid_simple_streaming"),
+    pytest.param(False, INVALID_COMPLETE_NONSTREAMING, id="invalid_complete"),
+    pytest.param(True, INVALID_COMPLETE_STREAMING, id="invalid_complete_streaming"),
+    pytest.param(False, NO_CONTENT, id="no_content"),
+    pytest.param(False, NO_REASONING, id="no_reasoning"),
+    pytest.param(False, REASONING_WITH_CHANNEL, id="reasoning"),
+    pytest.param(True, REASONING_WITH_CHANNEL, id="reasoning_streaming"),
+    pytest.param(False, COMPLETE_REASONING_WITH_CHANNEL, id="complete_reasoning"),
+    pytest.param(
+        True, COMPLETE_REASONING_WITH_CHANNEL, id="complete_reasoning_streaming"
+    ),
+    pytest.param(False, MULTIPLE_LINES_WITH_CHANNEL, id="multiple_lines"),
+    pytest.param(True, MULTIPLE_LINES_WITH_CHANNEL, id="multiple_lines_streaming"),
+    pytest.param(False, CHANNEL_NO_END, id="no_end"),
+    pytest.param(True, CHANNEL_NO_END, id="no_end_streaming"),
+    pytest.param(False, EMPTY, id="empty"),
+    pytest.param(False, NEW_LINE_NONSTREAMING, id="new_line"),
+    pytest.param(True, NEW_LINE_STREAMING, id="new_line_streaming"),
+    pytest.param(False, THOUGHT_PREFIX, id="thought_prefix"),
+    pytest.param(True, THOUGHT_PREFIX, id="thought_prefix_streaming"),
+    pytest.param(False, THOUGHT_PREFIX_ONLY, id="thought_prefix_only"),
+    pytest.param(True, THOUGHT_PREFIX_ONLY, id="thought_prefix_only_streaming"),
+    pytest.param(False, THOUGHT_PREFIX_MULTILINE, id="thought_prefix_multiline"),
+    pytest.param(
+        True, THOUGHT_PREFIX_MULTILINE, id="thought_prefix_multiline_streaming"
+    ),
+    pytest.param(False, THOUGHT_PREFIX_DIVERGE, id="thought_prefix_diverge"),
+    pytest.param(True, THOUGHT_PREFIX_DIVERGE, id="thought_prefix_diverge_streaming"),
+    pytest.param(False, TOOL_CALL_STARTED, id="tool_call_started"),
+    pytest.param(True, TOOL_CALL_STARTED, id="tool_call_started_streaming"),
+]
+
+
+def gemma4_encode_output(generic_tokenizer, output: str) -> list[int]:
+    # Resolve token IDs dynamically from the real tokenizer
+    vocab = generic_tokenizer.get_vocab()
+    start_token_id = vocab["<|channel>"]
+    end_token_id = vocab["<channel|>"]
+
+    index_start = output.find("<|channel>")
+    len_start = len("<|channel>")
+    index_end = output.find("<channel|>")
+    len_end = len("<channel|>")
+
+    output_tokens = []
+
+    def _encode(text: str) -> list[int]:
+        if not text:
+            return []
+        # Handle both raw transformers and vLLM wrappers
+        enc = getattr(generic_tokenizer, "tokenizer", generic_tokenizer)
+        try:
+            return enc.encode(text, add_special_tokens=False)
+        except TypeError:
+            return enc.encode(text)
+
+    if index_start != -1:
+        output_before = output[:index_start]
+        output_tokens += _encode(output_before)
+        output_tokens += [start_token_id]
+
+        if index_end != -1:
+            output_middle = output[index_start + len_start : index_end]
+            output_after = output[index_end + len_end :]
+            output_tokens += _encode(output_middle)
+            output_tokens += [end_token_id]
+            output_tokens += _encode(output_after)
+        else:
+            output_middle = output[index_start + len_start :]
+            output_tokens += _encode(output_middle)
+    elif index_end != -1:
+        output_before = output[:index_end]
+        output_after = output[index_end + len_end :]
+        output_tokens += _encode(output_before)
+        output_tokens += [end_token_id]
+        output_tokens += _encode(output_after)
+    else:
+        output_tokens += _encode(output)
+
+    return output_tokens
+
+
+@pytest.mark.parametrize("streaming, param_dict", TEST_CASES)
+def test_gemma4_reasoning(
+    streaming: bool,
+    param_dict: dict,
+    generic_tokenizer,
+):
+    output = param_dict["output"]
+    output_tokens = gemma4_encode_output(generic_tokenizer, output)
+
+    parser: ReasoningParser = ReasoningParserManager.get_reasoning_parser(parser_name)(
+        generic_tokenizer
+    )
+
+    # We use the generic run_reasoning_extraction from utils
+    # Use decode per token to get standard spaces instead of
+    # SentencePiece space characters
+    output_token_strings = [generic_tokenizer.decode([t]) for t in output_tokens]
+    reasoning, content = run_reasoning_extraction(
+        parser, output_token_strings, streaming=streaming
+    )
+
+    assert reasoning == param_dict["reasoning"]
+    assert content == param_dict["content"]
+
+    # Test is_reasoning_end
+    is_reasoning_end = parser.is_reasoning_end(output_tokens)
+    assert is_reasoning_end == param_dict["is_reasoning_end"]
+
+
+def test_gemma4_adjust_request(generic_tokenizer):
+    parser: ReasoningParser = ReasoningParserManager.get_reasoning_parser(parser_name)(
+        generic_tokenizer
+    )
+
+    request = ChatCompletionRequest(messages=[], model="test-model")
+    assert request.skip_special_tokens is True
+
+    result = parser.adjust_request(request)
+    assert result.skip_special_tokens is False
+    assert result is request
+
+
+def test_gemma4_previous_turn_reasoning_is_reasoning_end(generic_tokenizer):
+    output = (
+        "<|channel>thought\n1st thought<channel|>1st content<turn|>\n"
+        "<|turn>user\nThanks<|turn>model\n"
+    )
+    output_tokens = gemma4_encode_output(generic_tokenizer, output)
+    parser: ReasoningParser = ReasoningParserManager.get_reasoning_parser(parser_name)(
+        generic_tokenizer
+    )
+    is_reasoning_end = parser.is_reasoning_end(output_tokens)
+    assert not is_reasoning_end

tests/reasoning/test_step3p5_reasoning_parser.py

@@ -2,10 +2,10 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import pytest
-from transformers import AutoTokenizer
 
 from tests.reasoning.utils import run_reasoning_extraction
 from vllm.reasoning import ReasoningParser, ReasoningParserManager
+from vllm.tokenizers import get_tokenizer
 
 parser_name = "step3p5"
 start_token = "<think>"
@@ -16,7 +16,7 @@ REASONING_MODEL_NAME = "stepfun-ai/Step-3.5-Flash"
 
 @pytest.fixture(scope="module")
 def step3p5_tokenizer():
-    return AutoTokenizer.from_pretrained(REASONING_MODEL_NAME)
+    return get_tokenizer(tokenizer_name=REASONING_MODEL_NAME)
 
 
 SIMPLE_REASONING = {

tests/renderers/test_gemma4_chat_template.py

@@ -0,0 +1,345 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Tests for Gemma4 chat template rendering."""
+
+from pathlib import Path
+
+import jinja2.sandbox
+import pytest
+
+TEMPLATE_PATH = (
+    Path(__file__).resolve().parent.parent.parent
+    / "examples"
+    / "tool_chat_template_gemma4.jinja"
+)
+
+
+@pytest.fixture(scope="module")
+def gemma4_template():
+    """Load and compile the Gemma4 chat template."""
+    template_str = TEMPLATE_PATH.read_text()
+    env = jinja2.sandbox.ImmutableSandboxedEnvironment()
+    return env.from_string(template_str)
+
+
+def _render(template, messages, **kwargs):
+    """Render the template with sensible defaults."""
+    kwargs.setdefault("bos_token", "<bos>")
+    kwargs.setdefault("add_generation_prompt", False)
+    return template.render(messages=messages, **kwargs)
+
+
+class TestGemma4ChatTemplate:
+    def test_basic_multiturn_thinking_disabled(self, gemma4_template):
+        """With enable_thinking=False (default), generation prompt ends with
+        an empty thought channel to suppress thinking."""
+        messages = [
+            {"role": "user", "content": "Hello"},
+            {"role": "assistant", "content": "Hi there!"},
+            {"role": "user", "content": "How are you?"},
+        ]
+        result = _render(gemma4_template, messages, add_generation_prompt=True)
+        assert "<|turn>user\n" in result
+        assert "<|turn>model\n" in result
+        assert "Hello" in result
+        assert "Hi there!" in result
+        assert "How are you?" in result
+        assert result.rstrip("\n").endswith("<|channel>thought\n<channel|>")
+
+    def test_basic_multiturn_thinking_enabled(self, gemma4_template):
+        """With enable_thinking=True, generation prompt ends with model
+        turn opener (no thought suppression)."""
+        messages = [
+            {"role": "user", "content": "Hello"},
+            {"role": "assistant", "content": "Hi there!"},
+            {"role": "user", "content": "How are you?"},
+        ]
+        result = _render(
+            gemma4_template,
+            messages,
+            add_generation_prompt=True,
+            enable_thinking=True,
+        )
+        assert "<|turn>user\n" in result
+        assert "<|turn>model\n" in result
+        assert "Hello" in result
+        assert "Hi there!" in result
+        assert "How are you?" in result
+        assert result.rstrip("\n").endswith("<|turn>model")
+
+    def test_system_message(self, gemma4_template):
+        messages = [
+            {"role": "system", "content": "You are helpful."},
+            {"role": "user", "content": "Hi"},
+        ]
+        result = _render(gemma4_template, messages)
+        assert "<|turn>system\n" in result
+        assert "You are helpful." in result
+
+    def test_thinking_enabled(self, gemma4_template):
+        messages = [{"role": "user", "content": "Think about this"}]
+        result = _render(
+            gemma4_template,
+            messages,
+            add_generation_prompt=True,
+            enable_thinking=True,
+        )
+        assert "<|think|>" in result
+        assert "<|turn>system\n" in result
+
+    def test_tool_declarations(self, gemma4_template):
+        tools = [
+            {
+                "type": "function",
+                "function": {
+                    "name": "get_weather",
+                    "description": "Get weather for a city",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "city": {
+                                "type": "string",
+                                "description": "City name",
+                            }
+                        },
+                        "required": ["city"],
+                    },
+                },
+            }
+        ]
+        messages = [{"role": "user", "content": "What is the weather?"}]
+        result = _render(
+            gemma4_template,
+            messages,
+            tools=tools,
+            add_generation_prompt=True,
+        )
+        assert "<|tool>" in result
+        assert "declaration:get_weather" in result
+        assert "<tool|>" in result
+        assert '<|"|>City name<|"|>' in result
+
+    def test_tool_calls_in_assistant(self, gemma4_template):
+        messages = [
+            {"role": "user", "content": "Weather in London?"},
+            {
+                "role": "assistant",
+                "content": "",
+                "tool_calls": [
+                    {
+                        "id": "call_1",
+                        "function": {
+                            "name": "get_weather",
+                            "arguments": {"city": "London"},
+                        },
+                    }
+                ],
+            },
+        ]
+        result = _render(gemma4_template, messages)
+        assert "<|tool_call>call:get_weather{" in result
+        assert "}<tool_call|>" in result
+        assert '<|"|>London<|"|>' in result
+
+    def test_tool_responses_openai_style(self, gemma4_template):
+        """role='tool' messages are formatted as <|tool_response> blocks
+        with content dumped as-is."""
+        messages = [
+            {"role": "user", "content": "Weather?"},
+            {
+                "role": "assistant",
+                "content": "",
+                "tool_calls": [
+                    {
+                        "id": "call_1",
+                        "function": {
+                            "name": "get_weather",
+                            "arguments": {"city": "London"},
+                        },
+                    }
+                ],
+            },
+            {
+                "role": "tool",
+                "tool_call_id": "call_1",
+                "content": '{"temperature": 15, "condition": "sunny"}',
+            },
+        ]
+        result = _render(gemma4_template, messages, add_generation_prompt=True)
+        assert "<|tool_response>" in result
+        assert "response:get_weather{" in result
+        assert "<tool_response|>" in result
+        assert '"temperature": 15' in result
+
+    def test_tool_responses_legacy_style(self, gemma4_template):
+        """tool_responses embedded on the assistant message."""
+        messages = [
+            {"role": "user", "content": "Weather?"},
+            {
+                "role": "assistant",
+                "content": "",
+                "tool_calls": [
+                    {
+                        "function": {
+                            "name": "get_weather",
+                            "arguments": {"city": "London"},
+                        },
+                    }
+                ],
+                "tool_responses": [
+                    {
+                        "name": "get_weather",
+                        "response": {"temperature": 20},
+                    }
+                ],
+            },
+        ]
+        result = _render(gemma4_template, messages)
+        assert "<|tool_response>" in result
+        assert "response:get_weather{" in result
+        assert "temperature:" in result
+
+    def test_generation_prompt_not_after_tool_response(self, gemma4_template):
+        """add_generation_prompt=True should NOT add <|turn>model when the
+        last message type was tool_response (the model turn continues)."""
+        messages = [
+            {"role": "user", "content": "Weather?"},
+            {
+                "role": "assistant",
+                "content": "",
+                "tool_calls": [
+                    {
+                        "id": "call_1",
+                        "function": {
+                            "name": "get_weather",
+                            "arguments": {"city": "London"},
+                        },
+                    }
+                ],
+            },
+            {
+                "role": "tool",
+                "tool_call_id": "call_1",
+                "content": "sunny",
+            },
+        ]
+        result = _render(gemma4_template, messages, add_generation_prompt=True)
+        assert not result.strip().endswith("<|turn>model\n")
+
+    def test_reasoning_in_tool_chains(self, gemma4_template):
+        """reasoning field on assistant with tool_calls after last user
+        message emits <|channel>thought\\n...<channel|>."""
+        messages = [
+            {"role": "user", "content": "Calculate something"},
+            {
+                "role": "assistant",
+                "content": "",
+                "reasoning": "Let me think about this...",
+                "tool_calls": [
+                    {
+                        "function": {
+                            "name": "calculator",
+                            "arguments": {"expr": "2+2"},
+                        },
+                    }
+                ],
+            },
+        ]
+        result = _render(gemma4_template, messages)
+        assert "<|channel>thought\n" in result
+        assert "Let me think about this..." in result
+        assert "<channel|>" in result
+
+    def test_reasoning_not_before_last_user(self, gemma4_template):
+        """reasoning on assistant BEFORE the last user message is dropped."""
+        messages = [
+            {"role": "user", "content": "First"},
+            {
+                "role": "assistant",
+                "content": "Response",
+                "reasoning": "Old reasoning that should be dropped",
+                "tool_calls": [
+                    {
+                        "function": {
+                            "name": "fn",
+                            "arguments": {},
+                        },
+                    }
+                ],
+            },
+            {"role": "user", "content": "Second"},
+        ]
+        result = _render(gemma4_template, messages, add_generation_prompt=True)
+        assert "Old reasoning" not in result
+
+    def test_strip_thinking_in_model_content(self, gemma4_template):
+        """<|channel>...<channel|> in model content is stripped by the
+        strip_thinking macro."""
+        messages = [
+            {"role": "user", "content": "Hi"},
+            {
+                "role": "assistant",
+                "content": ("<|channel>internal thought<channel|>Visible answer"),
+            },
+        ]
+        result = _render(gemma4_template, messages)
+        assert "internal thought" not in result
+        assert "Visible answer" in result
+
+    def test_multi_turn_tool_chain(self, gemma4_template):
+        """assistant->tool->assistant->tool produces exactly one
+        <|turn>model (later assistants continue the same turn)."""
+        messages = [
+            {"role": "user", "content": "Do two things"},
+            {
+                "role": "assistant",
+                "content": "",
+                "tool_calls": [
+                    {
+                        "id": "c1",
+                        "function": {"name": "step1", "arguments": {}},
+                    },
+                ],
+            },
+            {"role": "tool", "tool_call_id": "c1", "content": "result1"},
+            {
+                "role": "assistant",
+                "content": "",
+                "tool_calls": [
+                    {
+                        "id": "c2",
+                        "function": {"name": "step2", "arguments": {}},
+                    },
+                ],
+            },
+            {"role": "tool", "tool_call_id": "c2", "content": "result2"},
+        ]
+        result = _render(gemma4_template, messages, add_generation_prompt=True)
+        assert result.count("<|turn>model\n") == 1
+
+    def test_format_argument_types(self, gemma4_template):
+        """Strings wrapped in <|"|>, booleans as true/false, numbers bare."""
+        messages = [
+            {"role": "user", "content": "Test"},
+            {
+                "role": "assistant",
+                "content": "",
+                "tool_calls": [
+                    {
+                        "function": {
+                            "name": "test_fn",
+                            "arguments": {
+                                "name": "Alice",
+                                "active": True,
+                                "count": 42,
+                            },
+                        },
+                    }
+                ],
+            },
+        ]
+        result = _render(gemma4_template, messages)
+        assert '<|"|>Alice<|"|>' in result
+        assert "active:true" in result
+        assert "count:42" in result

tests/tool_parsers/test_gemma4_tool_parser.py

@@ -0,0 +1,686 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import json
+from typing import Any
+from unittest.mock import MagicMock
+
+import pytest
+
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
+from vllm.tool_parsers.gemma4_tool_parser import (
+    TOOL_CALL_END,
+    TOOL_CALL_START,
+    Gemma4ToolParser,
+    _parse_gemma4_args,
+    _parse_gemma4_array,
+)
+
+# ---------------------------------------------------------------------------
+# Fixtures
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def mock_tokenizer():
+    tokenizer = MagicMock()
+    tokenizer.encode.return_value = [1, 2, 3]
+    # Include the tool call start token in the vocab for the parser
+    tokenizer.get_vocab.return_value = {TOOL_CALL_START: 48, TOOL_CALL_END: 49}
+    return tokenizer
+
+
+@pytest.fixture
+def parser(mock_tokenizer):
+    return Gemma4ToolParser(mock_tokenizer)
+
+
+@pytest.fixture
+def mock_request():
+    request = MagicMock(spec=ChatCompletionRequest)
+    request.tools = []
+    request.tool_choice = "auto"
+    return request
+
+
+# ---------------------------------------------------------------------------
+# Unit tests for _parse_gemma4_args (shared parser logic)
+# ---------------------------------------------------------------------------
+
+
+class TestParseGemma4Args:
+    def test_empty_string(self):
+        assert _parse_gemma4_args("") == {}
+
+    def test_whitespace_only(self):
+        assert _parse_gemma4_args("   ") == {}
+
+    def test_single_string_value(self):
+        result = _parse_gemma4_args('location:<|"|>Paris<|"|>')
+        assert result == {"location": "Paris"}
+
+    def test_string_value_with_comma(self):
+        result = _parse_gemma4_args('location:<|"|>Paris, France<|"|>')
+        assert result == {"location": "Paris, France"}
+
+    def test_multiple_string_values(self):
+        result = _parse_gemma4_args(
+            'location:<|"|>San Francisco<|"|>,unit:<|"|>celsius<|"|>'
+        )
+        assert result == {"location": "San Francisco", "unit": "celsius"}
+
+    def test_integer_value(self):
+        result = _parse_gemma4_args("count:42")
+        assert result == {"count": 42}
+
+    def test_float_value(self):
+        result = _parse_gemma4_args("score:3.14")
+        assert result == {"score": 3.14}
+
+    def test_boolean_true(self):
+        result = _parse_gemma4_args("flag:true")
+        assert result == {"flag": True}
+
+    def test_boolean_false(self):
+        result = _parse_gemma4_args("flag:false")
+        assert result == {"flag": False}
+
+    def test_null_value(self):
+        # Bare `null` must parse as None (Python), not the string "null".
+        # Without this, tool_choice=auto would emit `{"param": "null"}`
+        # instead of `{"param": null}` for nullable tool parameters.
+        result = _parse_gemma4_args("param:null")
+        assert result == {"param": None}
+        assert json.dumps(result) == '{"param": null}'
+
+    def test_mixed_types(self):
+        result = _parse_gemma4_args(
+            'name:<|"|>test<|"|>,count:42,active:true,score:3.14'
+        )
+        assert result == {
+            "name": "test",
+            "count": 42,
+            "active": True,
+            "score": 3.14,
+        }
+
+    def test_nested_object(self):
+        result = _parse_gemma4_args('nested:{inner:<|"|>value<|"|>}')
+        assert result == {"nested": {"inner": "value"}}
+
+    def test_array_of_strings(self):
+        result = _parse_gemma4_args('items:[<|"|>a<|"|>,<|"|>b<|"|>]')
+        assert result == {"items": ["a", "b"]}
+
+    def test_unterminated_string(self):
+        """Unterminated strings should take everything after the delimiter."""
+        result = _parse_gemma4_args('key:<|"|>unterminated')
+        assert result == {"key": "unterminated"}
+
+    def test_empty_value(self):
+        """Key with no value after colon."""
+        result = _parse_gemma4_args("key:")
+        assert result == {"key": ""}
+
+    def test_empty_value_partial_withheld(self):
+        """Key with no value is withheld in partial mode to avoid premature emission."""
+        result = _parse_gemma4_args("key:", partial=True)
+        assert result == {}
+        # also with a space after the colon
+        result = _parse_gemma4_args("key: ", partial=True)
+        assert result == {}
+
+    def test_empty_value_after_other_keys_partial_withheld(self):
+        """Trailing key with no value is withheld; earlier keys are kept."""
+        result = _parse_gemma4_args('name:<|"|>test<|"|>,flag:', partial=True)
+        assert result == {"name": "test"}
+
+
+class TestParseGemma4Array:
+    def test_string_array(self):
+        result = _parse_gemma4_array('<|"|>a<|"|>,<|"|>b<|"|>')
+        assert result == ["a", "b"]
+
+    def test_empty_array(self):
+        result = _parse_gemma4_array("")
+        assert result == []
+
+    def test_bare_values(self):
+        result = _parse_gemma4_array("42,true,3.14")
+        assert result == [42, True, 3.14]
+
+
+# ---------------------------------------------------------------------------
+# Non-streaming extraction tests
+# ---------------------------------------------------------------------------
+
+
+class TestExtractToolCalls:
+    def test_no_tool_calls(self, parser, mock_request):
+        model_output = "Hello, how can I help you today?"
+        result = parser.extract_tool_calls(model_output, mock_request)
+
+        assert result.tools_called is False
+        assert result.tool_calls == []
+        assert result.content == model_output
+
+    def test_single_tool_call(self, parser, mock_request):
+        model_output = (
+            '<|tool_call>call:get_weather{location:<|"|>London<|"|>}<tool_call|>'
+        )
+        result = parser.extract_tool_calls(model_output, mock_request)
+
+        assert result.tools_called is True
+        assert len(result.tool_calls) == 1
+        assert result.tool_calls[0].function.name == "get_weather"
+        args = json.loads(result.tool_calls[0].function.arguments)
+        assert args == {"location": "London"}
+
+    def test_multiple_arguments(self, parser, mock_request):
+        model_output = (
+            "<|tool_call>call:get_weather{"
+            'location:<|"|>San Francisco<|"|>,'
+            'unit:<|"|>celsius<|"|>}'
+            "<tool_call|>"
+        )
+        result = parser.extract_tool_calls(model_output, mock_request)
+
+        assert result.tools_called is True
+        assert len(result.tool_calls) == 1
+        assert result.tool_calls[0].function.name == "get_weather"
+        args = json.loads(result.tool_calls[0].function.arguments)
+        assert args == {"location": "San Francisco", "unit": "celsius"}
+
+    def test_text_before_tool_call(self, parser, mock_request):
+        model_output = (
+            "Let me check the weather for you. "
+            '<|tool_call>call:get_weather{location:<|"|>Paris<|"|>}'
+            "<tool_call|>"
+        )
+        result = parser.extract_tool_calls(model_output, mock_request)
+
+        assert result.tools_called is True
+        assert result.content == "Let me check the weather for you."
+        assert len(result.tool_calls) == 1
+        assert result.tool_calls[0].function.name == "get_weather"
+
+    def test_multiple_tool_calls(self, parser, mock_request):
+        model_output = (
+            '<|tool_call>call:get_weather{location:<|"|>London<|"|>}'
+            "<tool_call|>"
+            '<|tool_call>call:get_time{location:<|"|>London<|"|>}'
+            "<tool_call|>"
+        )
+        result = parser.extract_tool_calls(model_output, mock_request)
+
+        assert result.tools_called is True
+        assert len(result.tool_calls) == 2
+        assert result.tool_calls[0].function.name == "get_weather"
+        assert result.tool_calls[1].function.name == "get_time"
+
+    def test_nested_arguments(self, parser, mock_request):
+        model_output = (
+            "<|tool_call>call:complex_function{"
+            'nested:{inner:<|"|>value<|"|>},'
+            'list:[<|"|>a<|"|>,<|"|>b<|"|>]}'
+            "<tool_call|>"
+        )
+        result = parser.extract_tool_calls(model_output, mock_request)
+
+        assert result.tools_called is True
+        assert len(result.tool_calls) == 1
+        assert result.tool_calls[0].function.name == "complex_function"
+        args = json.loads(result.tool_calls[0].function.arguments)
+        assert args == {"nested": {"inner": "value"}, "list": ["a", "b"]}
+
+    def test_tool_call_with_number_and_boolean(self, parser, mock_request):
+        model_output = (
+            "<|tool_call>call:set_status{"
+            "is_active:true,"
+            "count:42,"
+            "score:3.14}"
+            "<tool_call|>"
+        )
+        result = parser.extract_tool_calls(model_output, mock_request)
+
+        assert result.tools_called is True
+        assert len(result.tool_calls) == 1
+        assert result.tool_calls[0].function.name == "set_status"
+        args = json.loads(result.tool_calls[0].function.arguments)
+        assert args == {"is_active": True, "count": 42, "score": 3.14}
+
+    def test_incomplete_tool_call(self, parser, mock_request):
+        model_output = '<|tool_call>call:get_weather{location:<|"|>London'
+        result = parser.extract_tool_calls(model_output, mock_request)
+
+        # Incomplete — no <tool_call|> end marker, regex won't match
+        assert result.tools_called is False
+        assert result.content == model_output
+
+    def test_hyphenated_function_name(self, parser, mock_request):
+        """Ensure function names with hyphens are parsed correctly."""
+        model_output = (
+            '<|tool_call>call:get-weather{location:<|"|>London<|"|>}<tool_call|>'
+        )
+        result = parser.extract_tool_calls(model_output, mock_request)
+
+        assert result.tools_called is True
+        assert result.tool_calls[0].function.name == "get-weather"
+
+    def test_dotted_function_name(self, parser, mock_request):
+        """Ensure function names with dots are parsed correctly."""
+        model_output = (
+            '<|tool_call>call:weather.get{location:<|"|>London<|"|>}<tool_call|>'
+        )
+        result = parser.extract_tool_calls(model_output, mock_request)
+
+        assert result.tools_called is True
+        assert result.tool_calls[0].function.name == "weather.get"
+
+    def test_no_arguments(self, parser, mock_request):
+        """Tool calls with empty arguments."""
+        model_output = "<|tool_call>call:get_status{}<tool_call|>"
+        result = parser.extract_tool_calls(model_output, mock_request)
+
+        assert result.tools_called is True
+        assert result.tool_calls[0].function.name == "get_status"
+        args = json.loads(result.tool_calls[0].function.arguments)
+        assert args == {}
+
+
+# ---------------------------------------------------------------------------
+# Streaming extraction tests
+# ---------------------------------------------------------------------------
+
+
+class TestStreamingExtraction:
+    """Tests for the streaming tool call extraction.
+
+    These simulate the token-by-token streaming that vLLM performs,
+    feeding incremental text to extract_tool_calls_streaming() and
+    verifying that the accumulated argument deltas form valid JSON.
+    """
+
+    def _simulate_streaming(
+        self, parser: Gemma4ToolParser, mock_request: Any, chunks: list[str]
+    ) -> list[tuple[Any, str]]:
+        """Feed chunks through the streaming parser and collect results.
+
+        Returns a list of (delta_message, accumulated_text) tuples.
+        """
+        results: list[tuple[Any, str]] = []
+        previous_text: str = ""
+        previous_token_ids: list[int] = []
+
+        for chunk in chunks:
+            current_text = previous_text + chunk
+            # Use token ID 48 for tool_call start, 49 for end, 0 otherwise
+            delta_token_ids: list[int] = []
+            if TOOL_CALL_START in chunk:
+                delta_token_ids.append(48)
+            elif TOOL_CALL_END in chunk:
+                delta_token_ids.append(49)
+            else:
+                delta_token_ids.append(0)
+
+            current_token_ids = previous_token_ids + delta_token_ids
+
+            delta = parser.extract_tool_calls_streaming(
+                previous_text=previous_text,
+                current_text=current_text,
+                delta_text=chunk,
+                previous_token_ids=tuple(previous_token_ids),
+                current_token_ids=tuple(current_token_ids),
+                delta_token_ids=tuple(delta_token_ids),
+                request=mock_request,
+            )
+            results.append((delta, current_text))
+            previous_text = current_text
+            previous_token_ids = list(current_token_ids)
+
+        return results
+
+    def _collect_arguments(self, results):
+        """Collect all argument deltas from streaming results into one string."""
+        args_text = ""
+        for delta, _ in results:
+            if delta and delta.tool_calls:
+                for tc in delta.tool_calls:
+                    func = tc.function if isinstance(tc.function, dict) else tc.function
+                    if isinstance(func, dict):
+                        arg = func.get("arguments", "")
+                    else:
+                        arg = getattr(func, "arguments", "") or ""
+                    if arg:
+                        args_text += arg
+        return args_text
+
+    def _collect_function_name(self, results):
+        """Extract the function name from streaming results."""
+        for delta, _ in results:
+            if delta and delta.tool_calls:
+                for tc in delta.tool_calls:
+                    func = tc.function if isinstance(tc.function, dict) else tc.function
+                    if isinstance(func, dict):
+                        name = func.get("name")
+                    else:
+                        name = getattr(func, "name", None)
+                    if name:
+                        return name
+        return None
+
+    def test_basic_streaming_single_tool(self, parser, mock_request):
+        """Simulate the exact streaming scenario from the bug report.
+
+        Model generates:
+        <|tool_call>call:get_weather{location:<|"|>Paris, France<|"|>}<tool_call|>
+
+        Expected: arguments should be valid JSON {"location": "Paris, France"}
+        """
+        chunks = [
+            "<|tool_call>",
+            "call:get_weather{",
+            'location:<|"|>Paris',
+            ", France",
+            '<|"|>}',
+            "<tool_call|>",
+        ]
+
+        results = self._simulate_streaming(parser, mock_request, chunks)
+
+        # Verify function name
+        name = self._collect_function_name(results)
+        assert name == "get_weather", f"Expected 'get_weather', got '{name}'"
+
+        # Verify arguments form valid JSON
+        args_text = self._collect_arguments(results)
+        assert args_text, "No arguments were streamed"
+        parsed_args = json.loads(args_text)
+        assert parsed_args == {"location": "Paris, France"}
+
+    def test_streaming_multi_arg(self, parser, mock_request):
+        """Streaming with multiple arguments."""
+        chunks = [
+            "<|tool_call>",
+            "call:get_weather{",
+            'location:<|"|>Tokyo<|"|>,',
+            'unit:<|"|>celsius<|"|>}',
+            "<tool_call|>",
+        ]
+
+        results = self._simulate_streaming(parser, mock_request, chunks)
+
+        name = self._collect_function_name(results)
+        assert name == "get_weather"
+
+        args_text = self._collect_arguments(results)
+        assert args_text
+        parsed_args = json.loads(args_text)
+        assert parsed_args == {"location": "Tokyo", "unit": "celsius"}
+
+    def test_streaming_no_extra_brace(self, parser, mock_request):
+        """Verify the closing } is NOT leaked into arguments (Bug #2)."""
+        chunks = [
+            "<|tool_call>",
+            "call:get_weather{",
+            'location:<|"|>London<|"|>}',
+            "<tool_call|>",
+        ]
+
+        results = self._simulate_streaming(parser, mock_request, chunks)
+        args_text = self._collect_arguments(results)
+        assert args_text
+
+        # The args text must be valid JSON (no extra })
+        parsed = json.loads(args_text)
+        assert parsed == {"location": "London"}
+
+        # Specifically assert no double-brace
+        assert args_text.count("}") <= 1, (
+            f"Arguments contain extra closing brace: {args_text!r}"
+        )
+
+    def test_streaming_no_unquoted_keys(self, parser, mock_request):
+        """Verify keys are properly quoted in JSON (Bug #1)."""
+        chunks = [
+            "<|tool_call>",
+            "call:get_weather{",
+            'location:<|"|>Paris<|"|>}',
+            "<tool_call|>",
+        ]
+
+        results = self._simulate_streaming(parser, mock_request, chunks)
+        args_text = self._collect_arguments(results)
+
+        # Must start with { and contain quoted key
+        assert args_text.lstrip().startswith("{"), (
+            f"Arguments don't start with '{{': {args_text!r}"
+        )
+        assert '"location"' in args_text, (
+            f"Key 'location' not properly quoted: {args_text!r}"
+        )
+
+    def test_streaming_name_no_call_prefix(self, parser, mock_request):
+        """Verify function name has no 'call:' prefix."""
+        chunks = [
+            "<|tool_call>",
+            "call:get_weather{",
+            'location:<|"|>Paris<|"|>}',
+            "<tool_call|>",
+        ]
+
+        results = self._simulate_streaming(parser, mock_request, chunks)
+        name = self._collect_function_name(results)
+        assert name == "get_weather"
+        assert not name.startswith("call:"), f"Name has 'call:' prefix: {name!r}"
+
+    def test_streaming_text_before_tool_call(self, parser, mock_request):
+        """Text before tool call should be emitted as content."""
+        chunks = [
+            "Let me check ",
+            "the weather. ",
+            "<|tool_call>",
+            "call:get_weather{",
+            'location:<|"|>London<|"|>}',
+            "<tool_call|>",
+        ]
+
+        results = self._simulate_streaming(parser, mock_request, chunks)
+
+        # First chunks should be content
+        content_parts = []
+        for delta, _ in results:
+            if delta and delta.content:
+                content_parts.append(delta.content)
+
+        assert "".join(content_parts).strip().startswith("Let me check")
+
+    def test_streaming_numeric_args(self, parser, mock_request):
+        """Streaming with numeric and boolean argument values."""
+        chunks = [
+            "<|tool_call>",
+            "call:set_config{",
+            "count:42,",
+            "active:true}",
+            "<tool_call|>",
+        ]
+
+        results = self._simulate_streaming(parser, mock_request, chunks)
+        args_text = self._collect_arguments(results)
+        if args_text:
+            parsed_args = json.loads(args_text)
+            assert parsed_args["count"] == 42
+            assert parsed_args["active"] is True
+
+    def test_streaming_boolean_split_across_chunks(self, parser, mock_request):
+        """Boolean value split across token boundaries must not corrupt JSON."""
+        chunks = [
+            "<|tool_call>",
+            "call:search{input:{all:" + "true"[:3],
+            "e}}",
+            "<tool_call|>",
+        ]
+
+        results = self._simulate_streaming(parser, mock_request, chunks)
+        args_text = self._collect_arguments(results)
+        assert args_text, "No arguments were streamed"
+        parsed_args = json.loads(args_text)
+        assert parsed_args["input"]["all"] is True
+
+    def test_streaming_false_split_across_chunks(self, parser, mock_request):
+        """Boolean false split across chunks."""
+        chunks = [
+            "<|tool_call>",
+            "call:set{flag:" + "false"[:4],
+            "e}",
+            "<tool_call|>",
+        ]
+
+        results = self._simulate_streaming(parser, mock_request, chunks)
+        args_text = self._collect_arguments(results)
+        assert args_text, "No arguments were streamed"
+        parsed_args = json.loads(args_text)
+        assert parsed_args["flag"] is False
+
+    def test_streaming_number_split_across_chunks(self, parser, mock_request):
+        """Number split across chunks must not change type."""
+        chunks = [
+            "<|tool_call>",
+            "call:set{count:4",
+            "2}",
+            "<tool_call|>",
+        ]
+
+        results = self._simulate_streaming(parser, mock_request, chunks)
+        args_text = self._collect_arguments(results)
+        assert args_text, "No arguments were streamed"
+        parsed_args = json.loads(args_text)
+        assert parsed_args["count"] == 42
+
+    def test_streaming_empty_args(self, parser, mock_request):
+        """Tool call with no arguments."""
+        chunks = [
+            "<|tool_call>",
+            "call:get_status{}",
+            "<tool_call|>",
+        ]
+
+        results = self._simulate_streaming(parser, mock_request, chunks)
+        name = self._collect_function_name(results)
+        assert name == "get_status"
+
+    def test_streaming_split_delimiter_no_invalid_json(self, parser, mock_request):
+        """Partial <|"|> delimiter chars must not leak into streamed JSON.
+
+        Reproduces the bug from https://github.com/vllm-project/vllm/issues/38946
+        where a token boundary splits the string delimiter, leaving fragments
+        like '<|' at the end of a parsed value which then corrupt the JSON.
+        """
+        chunks = [
+            "<|tool_call>",
+            "call:todowrite{",
+            'content:<|"|>Buy milk<|',
+            '"|>}',
+            "<tool_call|>",
+        ]
+
+        results = self._simulate_streaming(parser, mock_request, chunks)
+
+        args_text = self._collect_arguments(results)
+        assert args_text, "No arguments were streamed"
+
+        # Must be valid JSON — the original bug caused a JSON parse error
+        parsed_args = json.loads(args_text)
+        assert parsed_args["content"] == "Buy milk"
+
+        # Ensure no raw delimiter fragments leaked into the JSON
+        assert "<|" not in args_text, (
+            f"Partial delimiter leaked into JSON: {args_text!r}"
+        )
+
+    def test_streaming_does_not_duplicate_plain_text_after_tool_call(
+        self, parser, mock_request, monkeypatch
+    ):
+        """Buffered plain text after a tool call must not corrupt current_text."""
+        captured_current_texts: list[str] = []
+        original_extract_streaming = parser._extract_streaming
+
+        def wrapped_extract_streaming(previous_text, current_text, delta_text):
+            captured_current_texts.append(current_text)
+            return original_extract_streaming(previous_text, current_text, delta_text)
+
+        monkeypatch.setattr(parser, "_extract_streaming", wrapped_extract_streaming)
+
+        chunks = [
+            "<|tool_call>",
+            "call:get_weather{",
+            'location:<|"|>Paris<|"|>}',
+            "<tool_call|><",
+            "div>",
+        ]
+
+        results = self._simulate_streaming(parser, mock_request, chunks)
+        content_parts = [
+            delta.content for delta, _ in results if delta is not None and delta.content
+        ]
+        assert "".join(content_parts) == "<div>"
+        assert captured_current_texts[-1].endswith("<tool_call|><div>")
+        assert not captured_current_texts[-1].endswith("<tool_call|><<div>")
+
+    def test_streaming_html_argument_does_not_duplicate_tag_prefixes(
+        self, parser, mock_request
+    ):
+        """HTML content inside tool arguments must not be duplicated."""
+        chunks = [
+            "<|tool_call>",
+            "call:write_file{",
+            'path:<|"|>index.html<|"|>,',
+            'content:<|"|><!DOCTYPE html>\n<',
+            'html lang="zh-CN">\n<',
+            "head>\n    <",
+            'meta charset="UTF-8">\n    <',
+            'meta name="viewport" content="width=device-width">\n',
+            '<|"|>}',
+            "<tool_call|>",
+        ]
+
+        results = self._simulate_streaming(parser, mock_request, chunks)
+        args_text = self._collect_arguments(results)
+        assert args_text
+
+        parsed_args = json.loads(args_text)
+        assert parsed_args["path"] == "index.html"
+        assert (
+            parsed_args["content"] == "<!DOCTYPE html>\n"
+            '<html lang="zh-CN">\n'
+            "<head>\n"
+            '    <meta charset="UTF-8">\n'
+            '    <meta name="viewport" content="width=device-width">\n'
+        )
+
+    def test_streaming_trailing_bare_bool_not_duplicated(self, parser, mock_request):
+        """Trailing bare boolean must not be streamed twice."""
+        chunks = [
+            "<|tool_call>",
+            "call:Edit{",
+            'file_path:<|"|>src/env.py<|"|>,',
+            'old_string:<|"|>old_val<|"|>,',
+            'new_string:<|"|>new_val<|"|>,',
+            "replace_all:",
+            "false}",
+            "<tool_call|>",
+        ]
+
+        results = self._simulate_streaming(parser, mock_request, chunks)
+        args_text = self._collect_arguments(results)
+        assert args_text, "No arguments were streamed"
+
+        parsed_args = json.loads(args_text)
+        assert parsed_args == {
+            "file_path": "src/env.py",
+            "old_string": "old_val",
+            "new_string": "new_val",
+            "replace_all": False,
+        }
+
+        assert args_text.count("replace_all") == 1

tests/v1/attention/test_sparse_mla_backends.py

@@ -42,6 +42,7 @@ from vllm.v1.attention.backends.mla.flashmla_sparse import (
     FlashMLASparseBackend,
     triton_convert_req_index_to_global_index,
 )
+from vllm.v1.attention.backends.mla.indexer import split_indexer_prefill_chunks
 from vllm.v1.attention.backends.utils import split_prefill_chunks
 from vllm.v1.attention.ops import flashmla
 
@@ -716,6 +717,81 @@ def test_split_prefill_chunks(seq_lens, max_buf, expected):
     assert out == expected
 
 
+@pytest.mark.parametrize(
+    "seq_lens,query_lens,workspace_size,max_logits_bytes,expected",
+    [
+        # Logits constraint triggers split (M*N exceeds budget)
+        # req0: M=10, N=100 -> 1000 elems (4000 bytes) - fits in 5000
+        # req1: adding M=10, N=100 -> new_M=20, new_N=200 -> 4000 elems > 1250
+        (
+            torch.tensor([100, 100, 100]),
+            torch.tensor([10, 10, 10]),
+            1000,  # workspace allows all
+            5000,  # 1250 float32 elems -> forces split
+            [
+                (slice(0, 1), slice(0, 10)),
+                (slice(1, 2), slice(0, 10)),
+                (slice(2, 3), slice(0, 10)),
+            ],
+        ),
+        # Both constraints satisfied - all fit in one chunk
+        (
+            torch.tensor([10, 10, 10]),
+            torch.tensor([5, 5, 5]),
+            100,
+            10000,  # 2500 elems, M*N = 15*30 = 450 < 2500
+            [(slice(0, 3), slice(0, 15))],
+        ),
+        # Workspace constraint triggers first
+        (
+            torch.tensor([50, 50, 50]),
+            torch.tensor([1, 1, 1]),
+            50,  # workspace only fits one at a time
+            1000000,  # logits budget is huge
+            [
+                (slice(0, 1), slice(0, 1)),
+                (slice(1, 2), slice(0, 1)),
+                (slice(2, 3), slice(0, 1)),
+            ],
+        ),
+        # Greedy filling: first two fit, third doesn't
+        # req0: M=5, N=10 -> 50 elems
+        # req0+1: M=10, N=20 -> 200 elems <= 250
+        # req0+1+2: M=15, N=30 -> 450 elems > 250
+        (
+            torch.tensor([10, 10, 10]),
+            torch.tensor([5, 5, 5]),
+            100,
+            1000,  # 250 elems
+            [(slice(0, 2), slice(0, 10)), (slice(2, 3), slice(0, 5))],
+        ),
+    ],
+)
+def test_split_indexer_prefill_chunks(
+    seq_lens, query_lens, workspace_size, max_logits_bytes, expected
+):
+    out = split_indexer_prefill_chunks(
+        seq_lens,
+        query_lens,
+        workspace_size,
+        max_logits_bytes,
+    )
+    assert out == expected
+
+
+def test_split_indexer_prefill_chunks_single_request_overflow():
+    """Test that single request exceeding budget is sub-chunked on query dim."""
+    seq_lens = torch.tensor([1000, 50])
+    query_lens = torch.tensor([100, 5])
+
+    out = split_indexer_prefill_chunks(seq_lens, query_lens, 2000, 1000)
+    # max_logits_elems = 250, N=1000 -> max_q = 1 -> 100 query sub-chunks
+    expected = [(slice(0, 1), slice(i, i + 1)) for i in range(100)]
+    # req1: M=5, N=50 -> 250 elems fits budget
+    expected.append((slice(1, 2), slice(0, 5)))
+    assert out == expected
+
+
 def test_triton_convert_returns_valid_counts():
     """Test that return_valid_counts correctly counts non-negative indices."""
     device = torch.device("cuda")

tests/v1/e2e/spec_decode/test_spec_decode.py

@@ -542,12 +542,16 @@ def test_eagle_correctness_light(
             "auto",
             0.8,
         ),
-        (
+        pytest.param(
             ("eagle3", "Qwen/Qwen3-8B", "AngelSlim/Qwen3-8B_eagle3", 1),
             False,
             False,
             "transformers",
             0.8,
+            # TODO(hmellor): figure out why memory usage is so high
+            marks=pytest.mark.skip(
+                reason="Feature is experimental and uses too much memory in CI",
+            ),
         ),
         pytest.param(
             (

tests/v1/entrypoints/openai/test_thinking_token_budget.py

@@ -20,7 +20,7 @@ def server():
         "--reasoning-parser",
         "qwen3",
         "--reasoning-config",
-        '{"think_start_str": "<think>", "think_end_str": "</think>"}',
+        '{"reasoning_start_str": "<think>", "reasoning_end_str": "</think>"}',
         "--max-model-len",
         "2048",
         "--enforce-eager",

tests/v1/logits_processors/test_correctness.py

@@ -103,8 +103,8 @@ class LogitsProcsRequestParams:
 class MockReasoningConfig:
     """Mock reasoning config for testing ThinkingTokenBudgetLogitsProcessor."""
 
-    think_start_token_ids = [THINK_START_TOKEN_ID]
-    think_end_token_ids = [THINK_END_TOKEN_ID]
+    reasoning_start_token_ids = [THINK_START_TOKEN_ID]
+    reasoning_end_token_ids = [THINK_END_TOKEN_ID]
 
 
 def _generate_fake_sampling_metadata(
@@ -491,7 +491,7 @@ def _thinking_budget_validate(
 
         # Find if thinking has started in output tokens
         thinking_started = False
-        start_tokens = tb_processor.think_start_token_ids
+        start_tokens = tb_processor.reasoning_start_token_ids
 
         if len(start_tokens) > 0:
             for i in range(len(output_tokens) - len(start_tokens) + 1):
@@ -518,7 +518,7 @@ def _thinking_budget_validate(
                     )
 
                 # Validate that only end tokens are allowed
-                end_tokens = tb_processor.think_end_token_ids
+                end_tokens = tb_processor.reasoning_end_token_ids
                 if len(end_tokens) > 0:
                     expected_end_token_id = end_tokens[
                         min(state["end_count"], len(end_tokens) - 1)

tests/v1/simple_kv_offload/test_integration.py

@@ -0,0 +1,193 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Integration tests for SimpleCPUOffloadConnector with real models."""
+
+import time
+
+import pytest
+
+from vllm import LLM, SamplingParams, TokensPrompt
+from vllm.config import KVTransferConfig
+from vllm.platforms import current_platform
+
+if not current_platform.is_cuda():
+    pytest.skip("Requires CUDA", allow_module_level=True)
+
+# Small models for default CI / local runs (accuracy only).
+SMALL_MODELS = [
+    "meta-llama/Llama-3.2-1B-Instruct",
+    "google/gemma-3-1b-it",
+]
+
+# Large models for optional perf runs only (slow to load and execute).
+PERF_MODELS = [
+    "meta-llama/Llama-3.1-8B",
+    "openai/gpt-oss-20b",
+]
+
+
+def _make_llm(model: str, lazy: bool, cpu_bytes_to_use: int) -> LLM:
+    kv_transfer_config = KVTransferConfig(
+        kv_connector="SimpleCPUOffloadConnector",
+        kv_role="kv_both",
+        kv_connector_extra_config={
+            "cpu_bytes_to_use": cpu_bytes_to_use,
+            "lazy_offload": lazy,
+        },
+    )
+    return LLM(
+        model=model,
+        kv_cache_memory_bytes=40 << 30,  # 40 GiB
+        disable_hybrid_kv_cache_manager=False,
+        enable_prefix_caching=True,
+        kv_transfer_config=kv_transfer_config,
+    )
+
+
+def _flush_gpu_cache(llm: LLM, sampling_params: SamplingParams, seed: int = 0):
+    """Generate enough filler requests to allocate the entire GPU KV cache.
+
+    This pushes all prior blocks through the free queue so that the lazy
+    cursor offloads them to CPU before they are evicted.
+    """
+    cache_config = llm.llm_engine.vllm_config.cache_config
+    num_gpu_blocks = cache_config.num_gpu_blocks
+    block_size = cache_config.block_size
+    # Use 1.2x GPU capacity to give the lazy cursor enough scheduling steps
+    # to walk past all target blocks near the tail of the free queue.
+    total_tokens_needed = int(num_gpu_blocks * block_size * 1.5)
+
+    # Use token-id prompts so each filler is unique (no prefix sharing).
+    # Split into multiple requests to stay under max_model_len.
+    max_tokens_per_req = 4096
+    num_fillers = (total_tokens_needed + max_tokens_per_req - 1) // max_tokens_per_req
+    batch_size = 10
+    for i in range(0, num_fillers, batch_size):
+        batch_end = min(i + batch_size, num_fillers)
+        filler_prompts = []
+        for j in range(i, batch_end):
+            ids = [seed * num_fillers + j + 1] * max_tokens_per_req
+            filler_prompts.append(TokensPrompt(prompt_token_ids=ids))
+        llm.generate(filler_prompts, sampling_params, use_tqdm=False)
+
+
+def _accuracy_test(llm: LLM, lazy: bool = False):
+    """Verify that CPU-loaded KV produces correct output."""
+    sampling_params = SamplingParams(max_tokens=1, temperature=0)
+    prompt = "hi " * 2000 + "Let's count to ten. One, two, three, "
+
+    # Cold run — populate GPU cache and trigger CPU offload
+    cold_output = llm.generate(prompt, sampling_params, use_tqdm=False)[0]
+
+    # CPU hit runs
+    test_count = 10
+    success_count = 0
+    expected = cold_output.outputs[0].text
+    for i in range(test_count):
+        if lazy:
+            _flush_gpu_cache(llm, sampling_params, seed=i)
+        time.sleep(2)  # let engine core drain pending transfers
+
+        # Reset GPU prefix cache so next run must load from CPU
+        if not llm.reset_prefix_cache():
+            print(f"GPU prefix cache reset failed for iteration {i}")
+
+        output = llm.generate(prompt, sampling_params, use_tqdm=False)[0]
+        if output.outputs[0].text == expected:
+            success_count += 1
+
+    assert success_count >= 0.5 * test_count, (
+        f"Accuracy too low: {success_count}/{test_count} matched '{expected}'"
+    )
+
+
+def _latency_test(llm: LLM, lazy: bool = False):
+    """Verify CPU cache hit is faster than cold compute."""
+    sampling_params = SamplingParams(max_tokens=1, seed=42)
+    prompt_token_ids = [0] * 10001
+
+    num_times_cpu_better = 0
+    num_tests = 10
+    for i in range(num_tests):
+        prompt_token_ids[0] = i
+        prompts = [TokensPrompt(prompt_token_ids=prompt_token_ids)]
+
+        # Cold
+        time.sleep(2)  # let engine core drain pending transfers
+        if not llm.reset_prefix_cache():
+            print(f"GPU prefix cache reset failed for iteration {i}")
+        start = time.time()
+        llm.generate(prompts, sampling_params, use_tqdm=False)
+        cold_time = time.time() - start
+
+        if lazy:
+            _flush_gpu_cache(llm, sampling_params, seed=i)
+        else:
+            # Eager mode: GPU hit ensures store completion is processed.
+            llm.generate(prompts, sampling_params, use_tqdm=False)
+
+        time.sleep(2)  # let engine core drain pending transfers
+        if not llm.reset_prefix_cache():
+            print(f"GPU prefix cache reset failed for iteration {i}")
+
+        # CPU hit
+        start = time.time()
+        llm.generate(prompts, sampling_params, use_tqdm=False)
+        cpu_time = time.time() - start
+
+        if cpu_time < cold_time:
+            num_times_cpu_better += 1
+
+    assert num_times_cpu_better >= 0.8 * num_tests, (
+        f"CPU hit only faster {num_times_cpu_better}/{num_tests} times"
+    )
+
+
+@pytest.mark.optional
+@pytest.mark.slow_test
+@pytest.mark.parametrize("model", SMALL_MODELS)
+def test_simple_cpu_offload_accuracy(model: str):
+    """Store to CPU, reset GPU, load from CPU; verify output matches baseline."""
+    llm = _make_llm(model, False, 1 << 30)  # 1GB
+    try:
+        _accuracy_test(llm, lazy=False)
+    finally:
+        del llm
+
+
+@pytest.mark.optional
+@pytest.mark.slow_test
+@pytest.mark.parametrize("model", PERF_MODELS)
+def test_simple_cpu_offload_perf_latency(model: str):
+    """CPU KV hit should beat cold prefill on long context (large models only)."""
+    llm = _make_llm(model, False, 10 << 30)  # 10GB
+    try:
+        _latency_test(llm, lazy=False)
+    finally:
+        del llm
+
+
+@pytest.mark.optional
+@pytest.mark.slow_test
+@pytest.mark.parametrize("model", SMALL_MODELS)
+def test_simple_cpu_offload_accuracy_lazy(model: str):
+    """Lazy mode: flush GPU cache to trigger CPU offload, then verify hit."""
+    # CPU must be larger than GPU KV cache to avoid evicting offloaded blocks.
+    llm = _make_llm(model, True, 80 << 30)  # 80GB
+    try:
+        _accuracy_test(llm, lazy=True)
+    finally:
+        del llm
+
+
+@pytest.mark.optional
+@pytest.mark.slow_test
+@pytest.mark.parametrize("model", PERF_MODELS)
+def test_simple_cpu_offload_perf_latency_lazy(model: str):
+    """Lazy mode: CPU KV hit should beat cold prefill (large models only)."""
+    # CPU must be larger than GPU KV cache to avoid evicting offloaded blocks.
+    llm = _make_llm(model, True, 80 << 30)  # 80GB
+    try:
+        _latency_test(llm, lazy=True)
+    finally:
+        del llm

tests/v1/spec_decode/test_backup_token_async_spec.py

@@ -0,0 +1,147 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Regression tests for the backup token fix in prepare_next_token_ids_padded.
+
+Fixes #38098: with async scheduling, seq_lens_cpu is inflated by unaccepted
+draft token placeholders, causing get_token_id() to return -1.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+import torch
+
+
+class _FakeRequest:
+    def __init__(self, prompt_tokens: list[int], output_tokens: list[int]):
+        self.num_prompt_tokens = len(prompt_tokens)
+        self._prompt = prompt_tokens
+        self._output = output_tokens
+
+    @property
+    def num_tokens(self) -> int:
+        return self.num_prompt_tokens + len(self._output)
+
+    def get_token_id(self, idx: int) -> int:
+        if idx < self.num_prompt_tokens:
+            return self._prompt[idx]
+        out_idx = idx - self.num_prompt_tokens
+        if out_idx < len(self._output):
+            return self._output[out_idx]
+        return -1  # out of range
+
+
+class _FakeInputBatch:
+    def __init__(
+        self,
+        req_ids: list[str],
+        num_tokens_no_spec: list[int],
+        vocab_size: int = 32000,
+    ):
+        self.req_ids = req_ids
+        self.num_reqs = len(req_ids)
+        self.vocab_size = vocab_size
+        self.num_tokens_no_spec = np.array(num_tokens_no_spec, dtype=np.int64)
+
+
+def _make_requests(
+    req_ids: list[str],
+    prompt_lens: list[int],
+    output_lens: list[int],
+) -> dict[str, _FakeRequest]:
+    requests = {}
+    for rid, plen, olen in zip(req_ids, prompt_lens, output_lens):
+        requests[rid] = _FakeRequest(list(range(plen)), list(range(1000, 1000 + olen)))
+    return requests
+
+
+def _backup_buggy(
+    seq_lens_cpu: torch.Tensor,
+    requests: dict[str, _FakeRequest],
+    batch: _FakeInputBatch,
+) -> list[int]:
+    """Old logic: uses seq_lens_cpu directly (may be inflated)."""
+    n = batch.num_reqs
+    return [
+        requests[batch.req_ids[i]].get_token_id(int(seq_lens_cpu[i])) for i in range(n)
+    ]
+
+
+def _backup_fixed(
+    requests: dict[str, _FakeRequest],
+    batch: _FakeInputBatch,
+) -> list[int]:
+    """New logic: uses num_tokens_no_spec - 1 (last committed token)."""
+    n = batch.num_reqs
+    idx = (batch.num_tokens_no_spec[:n] - 1).tolist()
+    return [requests[batch.req_ids[i]].get_token_id(int(idx[i])) for i in range(n)]
+
+
+class TestBackupTokenAsyncSpec:
+    def test_no_inflation_fixed_returns_last_token(self):
+        req_ids = ["r0", "r1"]
+        requests = _make_requests(req_ids, [3, 3], [2, 2])
+        batch = _FakeInputBatch(req_ids, [5, 5])
+        # idx = 5-1 = 4 → output[1] = 1001
+        assert _backup_fixed(requests, batch) == [1001, 1001]
+
+    def test_inflation_buggy_returns_placeholder(self):
+        req_ids = ["r0", "r1"]
+        requests = _make_requests(req_ids, [3, 3], [2, 2])
+        batch = _FakeInputBatch(req_ids, [5, 5])
+        # inflated by 3 spec tokens → idx 8 is out of range
+        seq_lens = torch.tensor([8, 8], dtype=torch.int64)
+        assert _backup_buggy(seq_lens, requests, batch) == [-1, -1]
+
+    def test_inflation_fixed_returns_correct_token(self):
+        req_ids = ["r0", "r1"]
+        requests = _make_requests(req_ids, [3, 3], [2, 2])
+        batch = _FakeInputBatch(req_ids, [5, 5])
+        assert _backup_fixed(requests, batch) == [1001, 1001]
+
+    def test_mixed_inflation_per_request(self):
+        req_ids = ["r0", "r1", "r2"]
+        requests = {
+            "r0": _FakeRequest([0, 1], [1000, 1001, 1002]),
+            "r1": _FakeRequest([0, 1, 2, 3], [2000]),
+            "r2": _FakeRequest([0], [3000, 3001, 3002, 3003]),
+        }
+        batch = _FakeInputBatch(req_ids, [5, 5, 5])
+        seq_lens = torch.tensor([7, 9, 5], dtype=torch.int64)
+
+        assert _backup_buggy(seq_lens, requests, batch) == [-1, -1, -1]
+        assert _backup_fixed(requests, batch) == [1002, 2000, 3003]
+
+    def test_prefill_only_request(self):
+        """No output tokens yet — backup should be the last prompt token."""
+        req_ids = ["r0"]
+        requests = {"r0": _FakeRequest([10, 20, 30], [])}
+        batch = _FakeInputBatch(req_ids, [3])
+        # idx = 3-1 = 2 → prompt[2] = 30
+        assert _backup_fixed(requests, batch) == [30]
+
+    @pytest.mark.parametrize("num_spec_tokens", [1, 2, 3, 4, 5])
+    def test_various_spec_token_counts(self, num_spec_tokens: int):
+        req_ids = ["r0"]
+        requests = {"r0": _FakeRequest([0, 1, 2], list(range(1000, 1005)))}
+        batch = _FakeInputBatch(req_ids, [8])
+        # idx = 8-1 = 7 → output[4] = 1004
+        assert _backup_fixed(requests, batch) == [1004]
+
+    def test_buggy_code_was_always_off_by_one(self):
+        """The original code used seq_len as index, which is always one past
+        the end of output_token_ids even without async inflation."""
+        req_ids = ["r0"]
+        requests = {"r0": _FakeRequest([0, 1, 2], [1000, 1001])}
+        batch = _FakeInputBatch(req_ids, [5])
+
+        # no inflation: seq_len == num_tokens == 5 → idx 5 is out of range
+        seq_lens = torch.tensor([5], dtype=torch.int64)
+        assert _backup_buggy(seq_lens, requests, batch) == [-1]
+        assert _backup_fixed(requests, batch) == [1001]
+
+        # with inflation: still -1, fixed still correct
+        seq_lens_inf = torch.tensor([8], dtype=torch.int64)
+        assert _backup_buggy(seq_lens_inf, requests, batch) == [-1]
+        assert _backup_fixed(requests, batch) == [1001]

tests/v1/spec_decode/test_eagle.py

@@ -3,6 +3,7 @@
 
 from unittest import mock
 
+import numpy as np
 import pytest
 import torch
 
@@ -111,16 +112,14 @@ def test_prepare_next_token_ids():
 
     num_requests = 4
     num_speculative_tokens = 4
-    batch_spec = BatchSpec(
-        seq_lens=[num_speculative_tokens + 1] * num_requests,
-        query_lens=[num_speculative_tokens + 1] * num_requests,
-    )
-
     req_ids = [f"req_{i + 1}" for i in range(num_requests)]
     mock_input_batch = mock.MagicMock(spec=InputBatch)
     mock_input_batch.req_ids = req_ids
     mock_input_batch.num_reqs = num_requests
     mock_input_batch.vocab_size = 100
+    mock_input_batch.num_tokens_no_spec = np.array(
+        [num_speculative_tokens + 1] * num_requests
+    )
 
     mock_num_scheduled_tokens = {req_id: 0 for req_id in req_ids}
     mock_requests = {}
@@ -165,19 +164,12 @@ def test_prepare_next_token_ids():
 
     assert torch.equal(next_token_ids_from_cpu, expected_next_token_ids_tensor)
 
-    common_attn_metadata = create_common_attn_metadata(
-        batch_spec,
-        block_size=BLOCK_SIZE,
-        device=device,
-    )
-
     expected_valid_sampled_tokens_count = torch.tensor(
         [2, 5, 0, 0], dtype=torch.int32, device=device
     )
 
     next_token_ids_from_padded, valid_sampled_tokens_count = (
         proposer.prepare_next_token_ids_padded(
-            common_attn_metadata.seq_lens_cpu,
             sampled_token_ids_tensor,
             mock_requests,
             mock_input_batch,

tests/v1/spec_decode/test_extract_hidden_states.py

@@ -3,6 +3,7 @@
 
 from unittest import mock
 
+import numpy as np
 import pytest
 import torch
 
@@ -132,16 +133,12 @@ def test_prepare_next_token_ids_padded():
     device = torch.device(current_platform.device_type)
 
     num_requests = 4
-    batch_spec = BatchSpec(
-        seq_lens=[5] * num_requests,
-        query_lens=[5] * num_requests,
-    )
-
     req_ids = [f"req_{i + 1}" for i in range(num_requests)]
     mock_input_batch = mock.MagicMock(spec=InputBatch)
     mock_input_batch.req_ids = req_ids
     mock_input_batch.num_reqs = num_requests
     mock_input_batch.vocab_size = 100
+    mock_input_batch.num_tokens_no_spec = np.array([5] * num_requests)
 
     mock_requests = {}
     for req_id in req_ids:
@@ -174,12 +171,6 @@ def test_prepare_next_token_ids_padded():
 
     proposer = _create_proposer(num_speculative_tokens=1)
 
-    common_attn_metadata = create_common_attn_metadata(
-        batch_spec,
-        block_size=16,
-        device=device,
-    )
-
     # valid_sampled_tokens_count tracks if token is valid (not -1 and in vocab range)
     # It doesn't depend on whether the request is discarded
     expected_valid_sampled_tokens_count = torch.tensor(
@@ -187,7 +178,6 @@ def test_prepare_next_token_ids_padded():
     )
 
     next_token_ids, valid_sampled_tokens_count = proposer.prepare_next_token_ids_padded(
-        common_attn_metadata.seq_lens_cpu,
         sampled_token_ids,
         mock_requests,
         mock_input_batch,

vllm/_custom_ops.py

@@ -3397,3 +3397,38 @@ if hasattr(torch.ops._C, "hadacore_transform"):
     @register_fake("_C::hadacore_transform")
     def _hadacore_transform_fake(x: torch.Tensor, inplace: bool) -> torch.Tensor:
         return torch.empty_like(x) if not inplace else x
+
+
+if hasattr(torch.ops._C, "minimax_allreduce_rms"):
+
+    @register_fake("_C::minimax_allreduce_rms")
+    def _minimax_allreduce_rms_fake(
+        input: torch.Tensor,
+        norm_weight: torch.Tensor,
+        workspace: torch.Tensor,
+        rank: int,
+        nranks: int,
+        eps: float,
+    ) -> torch.Tensor:
+        return torch.empty_like(input)
+
+
+if hasattr(torch.ops._C, "minimax_allreduce_rms_qk"):
+
+    @register_fake("_C::minimax_allreduce_rms_qk")
+    def _minimax_allreduce_rms_qk_fake(
+        qkv: torch.Tensor,
+        norm_weight_q: torch.Tensor,
+        norm_weight_k: torch.Tensor,
+        workspace: torch.Tensor,
+        q_size: int,
+        kv_size: int,
+        rank: int,
+        nranks: int,
+        eps: float,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        token_num = qkv.shape[0]
+        return (
+            torch.empty([token_num, q_size], dtype=qkv.dtype, device=qkv.device),
+            torch.empty([token_num, kv_size], dtype=qkv.dtype, device=qkv.device),
+        )

vllm/compilation/decorators.py

@@ -205,6 +205,8 @@ def support_torch_compile(
                 if v.annotation in [
                     torch.Tensor,
                     torch.Tensor | None,
+                    torch.FloatTensor,
+                    torch.FloatTensor | None,
                     IntermediateTensors,
                     IntermediateTensors | None,
                 ]:
@@ -346,7 +348,7 @@ def _support_torch_compile(
 
     def __init__(
         self: _T,
-        *,
+        *args,
         vllm_config: VllmConfig | None = None,
         prefix: str = "",
         **kwargs: Any,
@@ -357,11 +359,24 @@ def _support_torch_compile(
         # NOTE: to support multimodal models (such as encoder),
         # we may not have vllm_config so we may need to patch it
         sig = inspect.signature(old_init)
+        # Check that any positional arguments match the old_init method signature
+        annotations = [p.annotation for p in sig.parameters.values()]
+        for arg, annotation in zip(args, annotations):
+            if annotation is inspect._empty:
+                continue
+            if not isinstance(arg, annotation):
+                init = f"'{type(self).__name__}.__init__'"
+                arg_type = f"'{type(arg).__name__}'"
+                raise TypeError(
+                    f"{init} received a positional argument of type {arg_type}, "
+                    "but no parameter of that type was found in the method signature. "
+                    f"Please either annotate {init} or pass it as a keyword argument."
+                )
         if "vllm_config" in sig.parameters:
             kwargs["vllm_config"] = vllm_config
         if "prefix" in sig.parameters:
             kwargs["prefix"] = prefix
-        old_init(self, **kwargs)
+        old_init(self, *args, **kwargs)
 
         self.vllm_config = vllm_config
         self.compilation_config = self.vllm_config.compilation_config

vllm/compilation/passes/fusion/minimax_qk_norm_fusion.py

@@ -0,0 +1,340 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Fusion pass: replace MiniMax QK allreduce + RMS norm with the Lamport
+fused kernel (minimax_allreduce_rms_qk) for decode-size batches.
+
+Pattern (inlined forward_qk in compiled graph):
+    q, k, v = qkv.split([q_size, kv_size, kv_size], -1)
+    q_fp32 = q.to(float32); k_fp32 = k.to(float32)
+    q_var = q_fp32.pow(2).mean(-1, keepdim=True)
+    k_var = k_fp32.pow(2).mean(-1, keepdim=True)
+    qk_var = cat([q_var, k_var], -1)
+    qk_var = allreduce(qk_var) / tp_world
+    q_var, k_var = qk_var.chunk(2, -1)
+    q_out = (q_fp32 * rsqrt(q_var + eps) * q_weight).to(orig_dtype)
+    k_out = (k_fp32 * rsqrt(k_var + eps) * k_weight).to(orig_dtype)
+    return q_out, k_out, v
+
+Replacement (pure, no in-place on qkv/q/k):
+    q_out, k_out = minimax_qk_norm_fused(qkv, q_weight, k_weight, workspace, ...)
+    v = qkv.split([q_size, kv_size, kv_size], -1)[2]
+    return q_out, k_out, v
+
+is_applicable_for_range: only fires for compile_range.end <= max_decode_tokens
+so that large prefill batches fall through to the original forward_qk (= main).
+"""
+
+import torch
+import torch._inductor.pattern_matcher as pm
+import torch.fx as fx
+from torch._inductor.pattern_matcher import PatternMatcherPass
+
+from vllm.config import VllmConfig
+from vllm.config.utils import Range
+from vllm.distributed import tensor_model_parallel_all_reduce
+from vllm.distributed.parallel_state import (
+    get_tensor_model_parallel_rank,
+    get_tensor_model_parallel_world_size,
+)
+from vllm.logger import init_logger
+from vllm.utils.torch_utils import direct_register_custom_op
+
+from ..inductor_pass import enable_fake_mode
+from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
+
+logger = init_logger(__name__)
+
+MAX_TOKEN_NUM = 2048
+
+_MINIMAX_QK_NORM_FUSED_OP = None
+if hasattr(torch.ops._C, "minimax_allreduce_rms_qk"):
+
+    def _minimax_qk_norm_fused(
+        qkv: torch.Tensor,
+        norm_weight_q: torch.Tensor,
+        norm_weight_k: torch.Tensor,
+        q_size: int,
+        kv_size: int,
+        rank: int,
+        nranks: int,
+        eps: float,
+        max_tokens: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        from vllm.distributed.parallel_state import get_tp_group
+        from vllm.model_executor.layers.mamba.lamport_workspace import (
+            get_allreduce_workspace,
+        )
+
+        workspace = get_allreduce_workspace(
+            rank=rank,
+            world_size=nranks,
+            max_tokens=max_tokens,
+            process_group=get_tp_group().cpu_group,
+        )
+        return torch.ops._C.minimax_allreduce_rms_qk(
+            qkv,
+            norm_weight_q,
+            norm_weight_k,
+            workspace,
+            q_size,
+            kv_size,
+            rank,
+            nranks,
+            eps,
+        )
+
+    def _minimax_qk_norm_fused_fake(
+        qkv: torch.Tensor,
+        norm_weight_q: torch.Tensor,
+        norm_weight_k: torch.Tensor,
+        q_size: int,
+        kv_size: int,
+        rank: int,
+        nranks: int,
+        eps: float,
+        max_tokens: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        T = qkv.shape[0]
+        return (
+            torch.empty([T, q_size], dtype=qkv.dtype, device=qkv.device),
+            torch.empty([T, kv_size], dtype=qkv.dtype, device=qkv.device),
+        )
+
+    direct_register_custom_op(
+        op_name="minimax_qk_norm_fused",
+        op_func=_minimax_qk_norm_fused,
+        fake_impl=_minimax_qk_norm_fused_fake,
+        mutates_args=[],
+    )
+    _MINIMAX_QK_NORM_FUSED_OP = torch.ops.vllm.minimax_qk_norm_fused.default
+
+
+class MiniMaxQKNormPattern:
+    """
+    Match the forward_qk allreduce+rms pattern and replace with Lamport kernel.
+    """
+
+    def __init__(
+        self,
+        q_size: int,
+        kv_size: int,
+        eps: float,
+        tp_world: int,
+        tp_rank: int,
+        max_tokens: int,
+        dtype: torch.dtype,
+        device: str | None,
+    ) -> None:
+        self.q_size = q_size
+        self.kv_size = kv_size
+        self.eps = eps
+        self.tp_world = tp_world
+        self.tp_rank = tp_rank
+        self.max_tokens = max_tokens
+        self.dtype = dtype
+        self.device = device
+
+    def get_inputs(self) -> list[torch.Tensor]:
+        T = 4
+        qkv = torch.empty(
+            [T, self.q_size + 2 * self.kv_size],
+            device=self.device,
+            dtype=self.dtype,
+        )
+        q_weight = torch.empty([self.q_size], device=self.device, dtype=self.dtype)
+        k_weight = torch.empty([self.kv_size], device=self.device, dtype=self.dtype)
+        return [qkv, q_weight, k_weight]
+
+    def register(self, pm_pass: PatternMatcherPass) -> None:
+        q_size = self.q_size
+        kv_size = self.kv_size
+        eps = self.eps
+        tp_world = self.tp_world
+        max_tokens = self.max_tokens
+        tp_rank = self.tp_rank
+        dtype = self.dtype
+
+        def pattern(
+            qkv: torch.Tensor,
+            q_weight: torch.Tensor,
+            k_weight: torch.Tensor,
+        ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+            q, k, v = qkv.split([q_size, kv_size, kv_size], dim=-1)
+            q_fp32 = q.to(torch.float32)
+            k_fp32 = k.to(torch.float32)
+            q_var = q_fp32.pow(2).mean(dim=-1, keepdim=True)
+            k_var = k_fp32.pow(2).mean(dim=-1, keepdim=True)
+            qk_var = torch.cat([q_var, k_var], dim=-1)
+            qk_var = tensor_model_parallel_all_reduce(qk_var) / tp_world
+            q_var, k_var = qk_var.chunk(2, dim=-1)
+            q_out = (q_fp32 * torch.rsqrt(q_var + eps) * q_weight).to(dtype)
+            k_out = (k_fp32 * torch.rsqrt(k_var + eps) * k_weight).to(dtype)
+            return q_out, k_out, v
+
+        def replacement(
+            qkv: torch.Tensor,
+            q_weight: torch.Tensor,
+            k_weight: torch.Tensor,
+        ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+            assert _MINIMAX_QK_NORM_FUSED_OP is not None
+            q_out, k_out = torch.ops.vllm.minimax_qk_norm_fused(
+                qkv,
+                q_weight,
+                k_weight,
+                q_size,
+                kv_size,
+                tp_rank,
+                tp_world,
+                eps,
+                max_tokens,
+            )
+            _, _, v = qkv.split([q_size, kv_size, kv_size], dim=-1)
+            return q_out, k_out, v
+
+        pm.register_replacement(
+            pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
+        )
+
+        # Second pattern: three separate split_with_sizes nodes (one per output),
+        # each with _users=1. This occurs when the QKV projection uses a
+        # functional GEMM kernel (e.g. cutlass_scaled_mm via auto_functionalized),
+        # which causes inductor to generate one split per consumer.
+        def pattern_split3(
+            qkv: torch.Tensor,
+            q_weight: torch.Tensor,
+            k_weight: torch.Tensor,
+        ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+            q = qkv.split([q_size, kv_size, kv_size], dim=-1)[0]
+            k = qkv.split([q_size, kv_size, kv_size], dim=-1)[1]
+            v = qkv.split([q_size, kv_size, kv_size], dim=-1)[2]
+            q_fp32 = q.to(torch.float32)
+            k_fp32 = k.to(torch.float32)
+            q_var = q_fp32.pow(2).mean(dim=-1, keepdim=True)
+            k_var = k_fp32.pow(2).mean(dim=-1, keepdim=True)
+            qk_var = torch.cat([q_var, k_var], dim=-1)
+            qk_var = tensor_model_parallel_all_reduce(qk_var) / tp_world
+            q_var, k_var = qk_var.chunk(2, dim=-1)
+            q_out = (q_fp32 * torch.rsqrt(q_var + eps) * q_weight).to(dtype)
+            k_out = (k_fp32 * torch.rsqrt(k_var + eps) * k_weight).to(dtype)
+            return q_out, k_out, v
+
+        pm.register_replacement(
+            pattern_split3, replacement, self.get_inputs(), pm.fwd_only, pm_pass
+        )
+
+
+class MiniMaxQKNormPass(VllmPatternMatcherPass):
+    """
+    Replace forward_qk allreduce+norm with the Lamport fused kernel.
+    Only applied for decode-size compile ranges (small token counts).
+    """
+
+    def __init__(self, config: VllmConfig) -> None:
+        super().__init__(config)
+        self.disabled = True
+
+        if _MINIMAX_QK_NORM_FUSED_OP is None:
+            logger.warning_once(
+                "minimax_allreduce_rms_qk op not found, MiniMaxQKNormPass disabled."
+            )
+            return
+
+        tp_world = get_tensor_model_parallel_world_size()
+        if tp_world <= 1:
+            logger.warning_once("MiniMaxQKNormPass disabled: tp_size <= 1.")
+            return
+
+        if config.model_config is None:
+            logger.warning_once("MiniMaxQKNormPass disabled: no model_config.")
+            return
+
+        hf_cfg = config.model_config.hf_config
+
+        model_name = getattr(hf_cfg, "architectures", "")[0]
+        if model_name != "MiniMaxM2ForCausalLM":
+            return
+
+        num_attention_heads = getattr(hf_cfg, "num_attention_heads", 0)
+        num_key_value_heads = getattr(hf_cfg, "num_key_value_heads", 0)
+        hidden_size = getattr(hf_cfg, "hidden_size", 0)
+        head_dim = getattr(hf_cfg, "head_dim", 0)
+        eps: float = getattr(hf_cfg, "rms_norm_eps", 1e-6)
+
+        if (
+            num_attention_heads != 48
+            or num_key_value_heads != 8
+            or hidden_size != 3072
+            or head_dim != 128
+        ):
+            logger.warning_once(
+                "MiniMaxQKNormPass disabled: cannot infer model info from hf_config."
+            )
+            return
+
+        num_heads_per_rank = num_attention_heads // tp_world
+        num_kv_heads_per_rank = max(1, num_key_value_heads // tp_world)
+        q_size = num_heads_per_rank * head_dim
+        kv_size = num_kv_heads_per_rank * head_dim
+
+        self.max_token_num = min(
+            MAX_TOKEN_NUM, config.scheduler_config.max_num_batched_tokens
+        )
+
+        tp_rank = get_tensor_model_parallel_rank()
+        # Allocate Lamport workspace first.
+        from vllm.distributed.parallel_state import get_tp_group
+        from vllm.model_executor.layers.mamba.lamport_workspace import (
+            get_allreduce_workspace,
+        )
+
+        get_allreduce_workspace(
+            rank=tp_rank,
+            world_size=tp_world,
+            max_tokens=self.max_token_num,
+            process_group=get_tp_group().cpu_group,
+        )
+
+        self.patterns: PatternMatcherPass = PatternMatcherPass(
+            pass_name="minimax_qk_norm_pass"
+        )
+        self._register_patterns(q_size, kv_size, eps, tp_world, tp_rank)
+        self.dump_patterns(config, self.patterns)
+        self.disabled = False
+
+    @enable_fake_mode
+    def _register_patterns(
+        self,
+        q_size: int,
+        kv_size: int,
+        eps: float,
+        tp_world: int,
+        tp_rank: int,
+    ) -> None:
+        MiniMaxQKNormPattern(
+            q_size=q_size,
+            kv_size=kv_size,
+            eps=eps,
+            tp_world=tp_world,
+            tp_rank=tp_rank,
+            max_tokens=self.max_token_num,
+            dtype=self.model_dtype,
+            device=self.device,
+        ).register(self.patterns)
+
+    def is_applicable_for_range(self, compile_range: Range) -> bool:
+        if self.disabled:
+            return False
+
+        return bool(compile_range.end <= self.max_token_num)
+
+    @VllmInductorPass.time_and_log
+    def __call__(self, graph: fx.Graph) -> None:
+        if self.disabled:
+            return
+        self.matched_count = self.patterns.apply(graph)
+        logger.debug("MiniMaxQKNormPass replaced %s patterns", self.matched_count)
+
+    def uuid(self) -> str:
+        return VllmInductorPass.hash_source(self, MiniMaxQKNormPattern)

vllm/compilation/passes/pass_manager.py

@@ -36,6 +36,7 @@ if current_platform.is_cuda_alike():
 if current_platform.is_cuda():
     from .fusion.allreduce_rms_fusion import AllReduceFusionPass
     from .fusion.collective_fusion import AsyncTPPass
+    from .fusion.minimax_qk_norm_fusion import MiniMaxQKNormPass
 
 from .inductor_pass import (
     CustomGraphPass,
@@ -124,6 +125,9 @@ class PostGradPassManager(CustomGraphPass):  # type: ignore[misc]
             if self.pass_config.fuse_allreduce_rms:
                 self.passes += [AllReduceFusionPass(config)]
 
+            if self.pass_config.fuse_minimax_qk_norm:
+                self.passes += [MiniMaxQKNormPass(config)]
+
             if self.pass_config.fuse_norm_quant:
                 self.passes += [RMSNormQuantFusionPass(config)]
                 if rocm_aiter_ops.is_enabled():

vllm/config/compilation.py

@@ -132,6 +132,8 @@ class PassConfig:
     """Enable async TP."""
     fuse_allreduce_rms: bool = None  # type: ignore[assignment]
     """Enable flashinfer allreduce fusion."""
+    fuse_minimax_qk_norm: bool = None  # type: ignore[assignment]
+    """Enable fused allreduce+RMSNorm for MiniMax QK norm."""
     enable_qk_norm_rope_fusion: bool = False
     """Enable fused Q/K RMSNorm + RoPE pass."""
 
@@ -282,7 +284,7 @@ class PassConfig:
         """
         enabled_fusions = [
             f.name[len("fuse_") :]
-            for f in fields(self)
+            for f in fields(self)  # type: ignore[arg-type]
             if getattr(self, f.name) and f.name.startswith("fuse_")
         ]
 
@@ -486,9 +488,10 @@ class CompilationConfig:
     If empty list [], no ops are excluded (suitable for full cudagraphs)."""
     compile_mm_encoder: bool = False
     """Whether or not to compile the multimodal encoder.
-    Currently, this only works for `Qwen2_5_vl` and `mLLaMa4` models
-    on selected platforms. Disabled by default until more models
-    are supported/tested to work."""
+    Currently, this only works for `Qwen2_5_vl` and `mLLaMa4` models on selected
+    platforms. It may also work for models loaded with the Transformers modeling backend
+    if the encoder is compilable. Disabled by default until more models are
+    supported/tested to work."""
 
     # Vision encoder CUDA graph
     cudagraph_mm_encoder: bool = False

vllm/config/reasoning.py

@@ -12,7 +12,7 @@ from vllm.tokenizers import cached_tokenizer_from_config
 class ReasoningConfig:
     """Configuration for reasoning models.
 
-    Set `think_start_str` and `think_end_str` to the strings that delimit
+    Set `reasoning_start_str` and `reasoning_end_str` to the strings that delimit
     the reasoning block (e.g. `"<think>"` and `"</think>"`).  The
     corresponding token IDs are derived automatically via
     `initialize_token_ids` and are not intended to be set directly.
@@ -20,53 +20,55 @@ class ReasoningConfig:
 
     # NOTE: These parameters are temporary, the intent is to derive them
     # automatically from the reasoning parser in a future version.
-    think_start_str: str = "<think>"
+    reasoning_start_str: str = "<think>"
     """String that indicates the start of reasoning."""
-    think_end_str: str = "</think>"
+    reasoning_end_str: str = "</think>"
     """String that indicates the end of reasoning content."""
 
-    _think_start_token_ids: list[int] | None = field(
+    _reasoning_start_token_ids: list[int] | None = field(
         default=None, init=False, repr=False
     )
-    """Private backing field for `think_start_token_ids`. Set by
+    """Private backing field for `reasoning_start_token_ids`. Set by
     `initialize_token_ids`. Not intended to be configured directly."""
-    _think_end_token_ids: list[int] | None = field(default=None, init=False, repr=False)
-    """Private backing field for `think_end_token_ids`. Set by
+    _reasoning_end_token_ids: list[int] | None = field(
+        default=None, init=False, repr=False
+    )
+    """Private backing field for `reasoning_end_token_ids`. Set by
     `initialize_token_ids`. Not intended to be configured directly."""
 
     @property
-    def think_start_token_ids(self) -> list[int] | None:
-        """Token IDs derived from `think_start_str`. Set automatically by
+    def reasoning_start_token_ids(self) -> list[int] | None:
+        """Token IDs derived from `reasoning_start_str`. Set automatically by
         `initialize_token_ids`. Not intended to be configured directly."""
-        return self._think_start_token_ids
+        return self._reasoning_start_token_ids
 
     @property
-    def think_end_token_ids(self) -> list[int] | None:
-        """Token IDs derived from `think_end_str`. Set automatically by
+    def reasoning_end_token_ids(self) -> list[int] | None:
+        """Token IDs derived from `reasoning_end_str`. Set automatically by
         `initialize_token_ids`. Not intended to be configured directly."""
-        return self._think_end_token_ids
+        return self._reasoning_end_token_ids
 
     def initialize_token_ids(self, model_config: ModelConfig) -> None:
         """Initialize reasoning token IDs from strings using the tokenizer."""
         if (
-            self._think_start_token_ids is not None
-            and self._think_end_token_ids is not None
+            self._reasoning_start_token_ids is not None
+            and self._reasoning_end_token_ids is not None
         ):
             return
 
         tokenizer = cached_tokenizer_from_config(model_config=model_config)
 
-        self._think_start_token_ids = tokenizer.encode(
-            self.think_start_str, add_special_tokens=False
+        self._reasoning_start_token_ids = tokenizer.encode(
+            self.reasoning_start_str, add_special_tokens=False
         )
-        self._think_end_token_ids = tokenizer.encode(
-            self.think_end_str, add_special_tokens=False
+        self._reasoning_end_token_ids = tokenizer.encode(
+            self.reasoning_end_str, add_special_tokens=False
         )
 
-        if not self._think_start_token_ids or not self._think_end_token_ids:
+        if not self._reasoning_start_token_ids or not self._reasoning_end_token_ids:
             raise ValueError(
                 f"ReasoningConfig: failed to tokenize reasoning strings: "
-                f"think_start_str='{self.think_start_str}', "
-                f"think_end_str='{self.think_end_str}'. "
+                f"reasoning_start_str='{self.reasoning_start_str}', "
+                f"reasoning_end_str='{self.reasoning_end_str}'. "
                 "Ensure the strings are valid tokens in the model's vocabulary."
             )

vllm/config/speculative.py

@@ -805,6 +805,8 @@ class SpeculativeConfig:
             "deepseek_v3",
             "kimi_k2",
             "kimi_k25",
+            "minimax_m2",
+            "gemma4",
         ]
         if (
             self.method in ("eagle3", "extract_hidden_states")

vllm/config/vllm.py

@@ -657,7 +657,11 @@ class VllmConfig:
         )
 
         if kv_offloading_backend == "native":
-            self.kv_transfer_config.kv_connector = "OffloadingConnector"
+            if envs.VLLM_USE_SIMPLE_KV_OFFLOAD:
+                config_connector = "SimpleCPUOffloadConnector"
+            else:
+                config_connector = "OffloadingConnector"
+            self.kv_transfer_config.kv_connector = config_connector
             self.kv_transfer_config.kv_connector_extra_config.update(
                 {"cpu_bytes_to_use": kv_offloading_size * (1 << 30)}
             )
@@ -1573,6 +1577,22 @@ class VllmConfig:
                         compile_range_end,
                     )
 
+        if compilation_config.pass_config.fuse_minimax_qk_norm:
+            from vllm.compilation.passes.fusion.minimax_qk_norm_fusion import (
+                MAX_TOKEN_NUM,
+            )
+
+            max_token_num = min(
+                MAX_TOKEN_NUM, self.scheduler_config.max_num_batched_tokens
+            )
+            if compile_range_end is not None and max_token_num < compile_range_end:
+                computed_compile_ranges_endpoints.append(max_token_num)
+            else:
+                logger.debug(
+                    "Max num batched tokens below MiniMax QK norm fusion threshold, "
+                    "MiniMax QK norm fusion enabled for all num_tokens."
+                )
+
         if compilation_config.compile_ranges_endpoints is not None:
             for x in compilation_config.compile_ranges_endpoints:
                 assert isinstance(x, int)

vllm/distributed/kv_transfer/kv_connector/factory.py

@@ -202,6 +202,7 @@ KVConnectorFactory.register_connector(
     "vllm.distributed.kv_transfer.kv_connector.v1.decode_bench_connector",
     "DecodeBenchConnector",
 )
+
 KVConnectorFactory.register_connector(
     "MooncakeConnector",
     "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector",
@@ -213,3 +214,9 @@ KVConnectorFactory.register_connector(
     "vllm.distributed.kv_transfer.kv_connector.v1.flexkv_connector",
     "FlexKVConnectorV1",
 )
+
+KVConnectorFactory.register_connector(
+    "SimpleCPUOffloadConnector",
+    "vllm.distributed.kv_transfer.kv_connector.v1.simple_cpu_offload_connector",
+    "SimpleCPUOffloadConnector",
+)

vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py

@@ -0,0 +1,247 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""SimpleCPUOffloadConnector: minimal CPU KV cache offloading."""
+
+from collections.abc import Iterable
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+from vllm.config import VllmConfig
+from vllm.distributed.kv_events import KVCacheEvent
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
+    KVConnectorBase_V1,
+    KVConnectorMetadata,
+    KVConnectorRole,
+    SupportsHMA,
+)
+from vllm.logger import init_logger
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.outputs import KVConnectorOutput
+from vllm.v1.simple_kv_offload.manager import (
+    SimpleCPUOffloadScheduler,
+)
+from vllm.v1.simple_kv_offload.metadata import (
+    SimpleCPUOffloadMetadata,
+)
+from vllm.v1.simple_kv_offload.worker import (
+    SimpleCPUOffloadWorker,
+)
+
+if TYPE_CHECKING:
+    from vllm.forward_context import ForwardContext
+    from vllm.v1.attention.backend import AttentionMetadata
+    from vllm.v1.core.block_pool import BlockPool
+    from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+    from vllm.v1.kv_cache_interface import KVCacheConfig
+    from vllm.v1.request import Request
+
+logger = init_logger(__name__)
+
+# Default CPU capacity: 8 GB
+DEFAULT_CPU_CAPACITY_BYTES = 8 * (1024**3)
+
+
+class SimpleCPUOffloadConnector(KVConnectorBase_V1, SupportsHMA):
+    """CPU KV cache offloading with custom kernel transfers and BlockPool LRU."""
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        role: KVConnectorRole,
+        kv_cache_config: "KVCacheConfig | None" = None,
+    ):
+        super().__init__(vllm_config, role, kv_cache_config)
+
+        enable_prefix_caching = vllm_config.cache_config.enable_prefix_caching
+        extra_config = self._kv_transfer_config.kv_connector_extra_config or {}
+
+        cpu_capacity_bytes = int(
+            extra_config.get("cpu_bytes_to_use", DEFAULT_CPU_CAPACITY_BYTES)
+        )
+        # cpu_bytes_to_use is server-wide for compatibility;
+        # cpu_bytes_to_use_per_rank overrides for per-rank capacity.
+        world_size = vllm_config.parallel_config.world_size
+        cpu_capacity_per_rank = cpu_capacity_bytes // world_size
+        if "cpu_bytes_to_use_per_rank" in extra_config:
+            explicit = int(extra_config["cpu_bytes_to_use_per_rank"])
+            if explicit != cpu_capacity_per_rank:
+                logger.warning(
+                    "cpu_bytes_to_use_per_rank (%.2f GB) != "
+                    "cpu_bytes_to_use/world_size (%.2f GB). Using per-rank value.",
+                    explicit / (1024**3),
+                    cpu_capacity_per_rank / (1024**3),
+                )
+            cpu_capacity_per_rank = explicit
+
+        lazy_offload = bool(extra_config.get("lazy_offload", False))
+
+        self.scheduler_manager: SimpleCPUOffloadScheduler | None = None
+        self.worker_handler: SimpleCPUOffloadWorker | None = None
+
+        if not enable_prefix_caching:
+            logger.warning(
+                "Detected prefix caching disabled, disabling CPU offload "
+                "since it requires prefix caching."
+            )
+            return
+
+        logger.info(
+            "SimpleCPUOffloadConnector: role=%s, "
+            "per_rank=%.2f GB, world_size=%d, mode=%s",
+            role.name,
+            cpu_capacity_per_rank / (1024**3),
+            world_size,
+            "lazy" if lazy_offload else "eager",
+        )
+
+        if role == KVConnectorRole.SCHEDULER:
+            self.scheduler_manager = SimpleCPUOffloadScheduler(
+                vllm_config,
+                kv_cache_config,
+                cpu_capacity_per_rank,
+                lazy_offload=lazy_offload,
+            )
+        elif role == KVConnectorRole.WORKER:
+            self.worker_handler = SimpleCPUOffloadWorker(
+                vllm_config, kv_cache_config, cpu_capacity_per_rank
+            )
+
+    # --- Worker-side methods ---
+
+    def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None:
+        if self.worker_handler is not None:
+            self.worker_handler.register_kv_caches(kv_caches)
+
+    def bind_connector_metadata(
+        self,
+        connector_metadata: KVConnectorMetadata,
+    ) -> None:
+        super().bind_connector_metadata(connector_metadata)
+        if self.worker_handler is not None:
+            assert isinstance(connector_metadata, SimpleCPUOffloadMetadata)
+            self.worker_handler.bind_connector_metadata(connector_metadata)
+
+    def clear_connector_metadata(self) -> None:
+        super().clear_connector_metadata()
+        if self.worker_handler is not None:
+            self.worker_handler.clear_connector_metadata()
+
+    def handle_preemptions(self, kv_connector_metadata: KVConnectorMetadata) -> None:
+        if self.worker_handler is not None:
+            assert isinstance(kv_connector_metadata, SimpleCPUOffloadMetadata)
+            self.worker_handler.handle_preemptions(kv_connector_metadata)
+
+    def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> None:
+        pass  # Launch loads ops in get_finished() after launching model execution
+
+    def wait_for_layer_load(self, layer_name: str) -> None:
+        pass  # Always load asynchronously and deferred to get_finished()
+
+    def save_kv_layer(
+        self,
+        layer_name: str,
+        kv_layer: torch.Tensor,
+        attn_metadata: "AttentionMetadata",
+        **kwargs: Any,
+    ) -> None:
+        pass  # Always save asynchronously and deferred to get_finished()
+
+    def wait_for_save(self) -> None:
+        pass  # All stores are driven by get_finished() and no wait needed
+
+    def get_finished(
+        self,
+        finished_req_ids: set[str],
+    ) -> tuple[set[str] | None, set[str] | None]:
+        if self.worker_handler is not None:
+            return self.worker_handler.get_finished(finished_req_ids)
+        return None, None
+
+    def build_connector_worker_meta(self):
+        if self.worker_handler is not None:
+            return self.worker_handler.build_connector_worker_meta()
+        return None
+
+    # --- Scheduler-side methods ---
+
+    # NOTE: New API only for SimpleCPUOffloadConnector.
+    def bind_gpu_block_pool(self, gpu_block_pool: "BlockPool") -> None:
+        if self.scheduler_manager is not None:
+            self.scheduler_manager.bind_gpu_block_pool(gpu_block_pool)
+
+    def get_num_new_matched_tokens(
+        self,
+        request: "Request",
+        num_computed_tokens: int,
+    ) -> tuple[int | None, bool]:
+        if self.scheduler_manager is not None:
+            return self.scheduler_manager.get_num_new_matched_tokens(
+                request, num_computed_tokens
+            )
+        return 0, False
+
+    def update_state_after_alloc(
+        self,
+        request: "Request",
+        blocks: "KVCacheBlocks",
+        num_external_tokens: int,
+    ) -> None:
+        if self.scheduler_manager is not None:
+            self.scheduler_manager.update_state_after_alloc(
+                request, blocks, num_external_tokens
+            )
+
+    def build_connector_meta(
+        self,
+        scheduler_output: SchedulerOutput,
+    ) -> KVConnectorMetadata:
+        if self.scheduler_manager is not None:
+            return self.scheduler_manager.build_connector_meta(scheduler_output)
+        return SimpleCPUOffloadMetadata()
+
+    def update_connector_output(
+        self,
+        connector_output: KVConnectorOutput,
+    ) -> None:
+        if self.scheduler_manager is not None:
+            self.scheduler_manager.update_connector_output(connector_output)
+
+    def request_finished(
+        self,
+        request: "Request",
+        block_ids: list[int],
+    ) -> tuple[bool, dict[str, Any] | None]:
+        if self.scheduler_manager is not None:
+            return self.scheduler_manager.request_finished(request, block_ids)
+        return False, None
+
+    def request_finished_all_groups(
+        self,
+        request: "Request",
+        block_ids: tuple[list[int], ...],
+    ) -> tuple[bool, dict[str, Any] | None]:
+        if self.scheduler_manager is not None:
+            return self.scheduler_manager.request_finished_all_groups(
+                request, block_ids
+            )
+        return False, None
+
+    # NOTE: New API only for SimpleCPUOffloadConnector.
+    def has_pending_transfers(self) -> bool:
+        if self.scheduler_manager is not None:
+            return self.scheduler_manager.has_pending_stores()
+        return False
+
+    def take_events(self) -> Iterable[KVCacheEvent]:
+        if self.scheduler_manager is not None:
+            return self.scheduler_manager.take_events()
+        return []
+
+    def reset_cache(self) -> bool | None:
+        raise NotImplementedError(
+            "SimpleCPUOffloadConnector does not support reset_cache(). "
+            "reset_prefix_cache() requires synchronizing all pending "
+            "CPU offload transfers before clearing GPU prefix cache blocks, "
+            "which is not yet implemented."
+        )

vllm/entrypoints/anthropic/serving.py

@@ -170,7 +170,8 @@ class AnthropicServingMessages(OpenAIServingChat):
             else:
                 cls._convert_message_content(msg, openai_msg, openai_messages)
 
-            openai_messages.append(openai_msg)
+            if not (msg.role == "user" and "content" not in openai_msg):
+                openai_messages.append(openai_msg)
 
     @classmethod
     def _convert_message_content(

vllm/entrypoints/openai/api_server.py

@@ -372,6 +372,7 @@ async def init_app_state(
         enable_auto_tools=args.enable_auto_tool_choice,
         exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none,
         tool_parser=args.tool_call_parser,
+        reasoning_parser=args.structured_outputs_config.reasoning_parser,
         default_chat_template_kwargs=args.default_chat_template_kwargs,
         log_error_stack=args.log_error_stack,
     )
@@ -467,6 +468,7 @@ async def init_render_app_state(
         enable_auto_tools=args.enable_auto_tool_choice,
         exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none,
         tool_parser=args.tool_call_parser,
+        reasoning_parser=args.structured_outputs_config.reasoning_parser,
         default_chat_template_kwargs=args.default_chat_template_kwargs,
         log_error_stack=args.log_error_stack,
     )

vllm/entrypoints/openai/responses/serving.py

@@ -594,6 +594,7 @@ class OpenAIServingResponses(OpenAIServing):
             default_template_kwargs=None,
             tool_dicts=tool_dicts,
             tool_parser=self.parser.tool_parser_cls if self.parser else None,
+            reasoning_parser=self.parser.reasoning_parser_cls if self.parser else None,
         )
         return messages, engine_inputs
 
@@ -618,6 +619,7 @@ class OpenAIServingResponses(OpenAIServing):
             default_template_kwargs=None,
             tool_dicts=tool_dicts,
             tool_parser=tool_parser,
+            reasoning_parser=self.parser.reasoning_parser_cls if self.parser else None,
         )
         return engine_inputs
 

vllm/entrypoints/serve/render/serving.py

@@ -44,6 +44,7 @@ from vllm.inputs import (
 )
 from vllm.logger import init_logger
 from vllm.parser import ParserManager
+from vllm.reasoning.abs_reasoning_parsers import ReasoningParser
 from vllm.renderers import BaseRenderer, merge_kwargs
 from vllm.renderers.inputs.preprocess import (
     extract_prompt_components,
@@ -74,6 +75,7 @@ class OpenAIServingRender:
         enable_auto_tools: bool = False,
         exclude_tools_when_tool_choice_none: bool = False,
         tool_parser: str | None = None,
+        reasoning_parser: str | None = None,
         default_chat_template_kwargs: dict[str, Any] | None = None,
         log_error_stack: bool = False,
     ) -> None:
@@ -94,6 +96,11 @@ class OpenAIServingRender:
             enable_auto_tools=enable_auto_tools,
             model_name=model_config.model,
         )
+        self.reasoning_parser: type[ReasoningParser] | None = (
+            ParserManager.get_reasoning_parser(
+                reasoning_parser_name=reasoning_parser,
+            )
+        )
         self.default_chat_template_kwargs: dict[str, Any] = (
             default_chat_template_kwargs or {}
         )
@@ -245,6 +252,7 @@ class OpenAIServingRender:
                 default_template_kwargs=self.default_chat_template_kwargs,
                 tool_dicts=tool_dicts,
                 tool_parser=tool_parser,
+                reasoning_parser=self.reasoning_parser,
             )
         else:
             # For GPT-OSS.
@@ -498,6 +506,9 @@ class OpenAIServingRender:
         default_template_kwargs: dict[str, Any] | None,
         tool_dicts: list[dict[str, Any]] | None = None,
         tool_parser: type[ToolParser] | None = None,
+        reasoning_parser: type[ReasoningParser] | None = None,
+        *,
+        skip_mm_cache: bool = False,
     ) -> tuple[list[ConversationMessage], list[EngineInput]]:
         """Copied from OpenAIServing._preprocess_chat."""
         renderer = self.renderer
@@ -531,6 +542,10 @@ class OpenAIServingRender:
             },
         )
 
+        if reasoning_parser is not None:
+            tokenizer = renderer.get_tokenizer()
+            request = reasoning_parser(tokenizer).adjust_request(request=request)
+
         # tool parsing is done only if a tool_parser has been set and if
         # tool_choice is not "none" (if tool_choice is "none" but a tool_parser
         # is set, we want to prevent parsing a tool_call hallucinated by the LLM

vllm/envs.py

@@ -54,6 +54,7 @@ if TYPE_CHECKING:
     VLLM_ZENTORCH_WEIGHT_PREPACK: bool = True
     VLLM_XLA_CACHE_PATH: str = os.path.join(VLLM_CACHE_ROOT, "xla_cache")
     VLLM_XLA_CHECK_RECOMPILATION: bool = False
+    VLLM_SPARSE_INDEXER_MAX_LOGITS_MB: int = 512
     VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE: Literal["auto", "nccl", "shm"] = "auto"
     VLLM_USE_RAY_COMPILED_DAG_OVERLAP_COMM: bool = False
     VLLM_USE_RAY_WRAPPED_PP_COMM: bool = True
@@ -842,6 +843,12 @@ environment_variables: dict[str, Callable[[], Any]] = {
     ),
     # Enable SPMD mode for TPU backend.
     "VLLM_XLA_USE_SPMD": lambda: bool(int(os.getenv("VLLM_XLA_USE_SPMD", "0"))),
+    # Maximum size (in MB) for logits tensor in sparse MLA indexer prefill chunks.
+    # Bounds the [M, N] float32 logits tensor to prevent CUDA OOM.
+    # Default: 512 MB
+    "VLLM_SPARSE_INDEXER_MAX_LOGITS_MB": lambda: int(
+        os.getenv("VLLM_SPARSE_INDEXER_MAX_LOGITS_MB", "512")
+    ),
     # If set, the OpenAI API server will stay alive even after the underlying
     # AsyncLLMEngine errors and stops serving requests
     "VLLM_KEEP_ALIVE_ON_ENGINE_DEATH": lambda: bool(
@@ -1655,6 +1662,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
     "VLLM_XPU_ENABLE_XPU_GRAPH": lambda: bool(
         int(os.getenv("VLLM_XPU_ENABLE_XPU_GRAPH", "0"))
     ),
+    # Enable simple KV offload.
+    "VLLM_USE_SIMPLE_KV_OFFLOAD": lambda: bool(
+        int(os.getenv("VLLM_USE_SIMPLE_KV_OFFLOAD", "0"))
+    ),
 }
 
 

vllm/model_executor/layers/mamba/lamport_workspace.py

@@ -0,0 +1,302 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+import array
+import contextlib
+import struct
+import sys
+import threading
+
+import torch
+
+try:
+    from cuda.bindings import runtime as cudart
+except ImportError:
+    from cuda import cudart
+
+_ALIGN = 1 << 21  # 2 MiB — CUDA IPC allocation alignment
+
+
+# ---------------------------------------------------------------------------
+# CUDA helpers
+# ---------------------------------------------------------------------------
+
+
+def _check(error):
+    """Raise on CUDA runtime error."""
+    success = getattr(cudart.cudaError_t, "cudaSuccess", None) or cudart.cudaError_t(0)
+    if error != success:
+        raise RuntimeError(f"CUDA runtime error: {error}")
+
+
+def _cuda_malloc(size: int):
+    aligned = ((size + _ALIGN - 1) >> 21) << 21
+    err, ptr = cudart.cudaMalloc(aligned)
+    _check(err)
+    return ptr, aligned
+
+
+def _cuda_free(ptr: int):
+    if ptr:
+        _check(cudart.cudaFree(ptr)[0])
+
+
+def _cuda_memset_zero(ptr: int, size: int):
+    _check(cudart.cudaMemset(ptr, 0, size)[0])
+
+
+def _cuda_memcpy_d2d(dst: int, src: int, size: int):
+    _check(
+        cudart.cudaMemcpy(
+            dst, src, size, cudart.cudaMemcpyKind.cudaMemcpyDeviceToDevice
+        )[0]
+    )
+
+
+# ---------------------------------------------------------------------------
+# IPC buffer
+# ---------------------------------------------------------------------------
+
+
+class IpcBuffer:
+    """
+    Allocates CUDA device memory and exchanges IPC handles with all ranks
+    so that every rank holds a valid device pointer to every other rank's buffer.
+    """
+
+    def __init__(self, rank: int, world_size: int, size: int, process_group=None):
+        self.rank = rank
+        self.world_size = world_size
+        self.peer_ptrs: list[int] = [0] * world_size
+        self.local_ptr: int = 0
+        self._alive = False
+
+        if size <= 0:
+            return
+
+        self.local_ptr, _ = _cuda_malloc(size)
+        _cuda_memset_zero(self.local_ptr, size)
+        self._alive = True
+
+        # --- exchange IPC handles via torch.distributed ---
+        err, local_handle = cudart.cudaIpcGetMemHandle(self.local_ptr)
+        _check(err)
+
+        all_handles: list[bytes | None] = [None] * world_size
+        torch.distributed.all_gather_object(
+            all_handles, bytes(local_handle.reserved), group=process_group
+        )
+
+        for r in range(world_size):
+            if r == rank:
+                self.peer_ptrs[r] = self.local_ptr
+            else:
+                handle = cudart.cudaIpcMemHandle_t()
+                handle.reserved = all_handles[r]
+                err, ptr = cudart.cudaIpcOpenMemHandle(
+                    handle, cudart.cudaIpcMemLazyEnablePeerAccess
+                )
+                _check(err)
+                self.peer_ptrs[r] = ptr
+
+    def serialize(self) -> list[int]:
+        """Return peer pointers as a list of int64 values (one per rank)."""
+        raw = b""
+        for ptr in self.peer_ptrs:
+            raw += struct.pack("P", ptr)
+        return array.array("Q", raw).tolist()
+
+    def cleanup(self):
+        if not self._alive:
+            return
+        self._alive = False
+        for r in range(self.world_size):
+            if self.peer_ptrs[r] == 0:
+                continue
+            if r == self.rank:
+                _cuda_free(self.peer_ptrs[r])
+            else:
+                with contextlib.suppress(RuntimeError):
+                    _check(cudart.cudaIpcCloseMemHandle(self.peer_ptrs[r])[0])
+            self.peer_ptrs[r] = 0
+        self.local_ptr = 0
+
+    def __del__(self):
+        if not sys.is_finalizing():
+            self.cleanup()
+
+
+# ---------------------------------------------------------------------------
+# Lamport negative-zero initialization
+# ---------------------------------------------------------------------------
+
+
+def _lamport_fill_neg_zero(device_ptr: int, size_bytes: int):
+    """
+    Fill device memory with IEEE-754 negative zero (-0.0f = 0x80000000).
+    This is the "slot empty" sentinel for the Lamport protocol: the kernel
+    spin-waits until a value is *not* negative zero.
+    """
+    if size_bytes == 0 or device_ptr == 0:
+        return
+    n_floats = size_bytes // 4
+    # torch preserves -0.0 in IEEE-754
+    fill = torch.full((n_floats,), -0.0, dtype=torch.float32, device="cuda")
+    _cuda_memcpy_d2d(device_ptr, fill.data_ptr(), size_bytes)
+    del fill
+
+
+# ---------------------------------------------------------------------------
+# LamportWorkspace — the main class
+# ---------------------------------------------------------------------------
+
+
+class LamportWorkspace:
+    """
+    Self-contained workspace for Lamport-based cross-GPU AllReduce.
+
+    Parameters
+    ----------
+    rank : int
+        Local rank (0-based).
+    world_size : int
+        Total number of ranks in the TP group.
+    comm_size : int
+        Size in bytes of *one* Lamport buffer slot. The total IPC allocation
+        per rank is ``3 * comm_size`` (triple-buffering). Must be large enough
+        to hold the per-slot data written by the kernel.  Use
+        ``compute_comm_size_for_minimax()`` for a safe default.
+    process_group : optional
+        ``torch.distributed`` process group for IPC handle exchange.
+        ``None`` uses the default group.
+    """
+
+    def __init__(self, rank: int, world_size: int, comm_size: int, process_group=None):
+        assert world_size >= 2, "Lamport workspace requires at least 2 ranks"
+        assert comm_size > 0, "comm_size must be positive"
+
+        self.rank = rank
+        self.world_size = world_size
+        self.comm_size = comm_size
+
+        # 1) Lamport triple-buffer (the only IPC memory the kernel reads/writes)
+        lamport_total = 3 * comm_size
+        self._lamport = IpcBuffer(rank, world_size, lamport_total, process_group)
+        _lamport_fill_neg_zero(self._lamport.local_ptr, lamport_total)
+
+        # 2) flag_buffer on device: int32[3] = {counter, unused, lamport_flag}
+        #    counter  — used for block-level sync inside the kernel
+        #    unused   — reserved (index 1)
+        #    lamport_flag — triple-buffer rotation index (0 → 1 → 2 → 0 …)
+        self._flag_buf = torch.zeros(3, dtype=torch.int32, device="cuda")
+
+        # 3) layout_buffer on device: int64[2] = {clear_size, comm_size}
+        #    clear_size — bytes to clear from *previous* slot (set by kernel)
+        #    comm_size  — size of one triple-buffer slot
+        self._layout_buf = torch.tensor(
+            [0, comm_size], dtype=torch.int64, device="cuda"
+        )
+
+        # 4) Assemble device-side void* pointer array
+        N = world_size
+        ptrs: list[int] = []
+        ptrs += [0] * N  # [0   .. N-1]   ipc_buffers  (placeholder)
+        ptrs += [0] * N  # [N   .. 2N-1]  ipc_barriers (placeholder)
+        ptrs += self._lamport.serialize()  # [2N  .. 3N-1]  lamport peer ptrs
+        ptrs.append(self._flag_buf.data_ptr())  # [3N]           flag_buffer
+        ptrs.append(self._layout_buf.data_ptr())  # [3N+1]       layout_buffer
+
+        self._workspace = torch.tensor(ptrs, dtype=torch.int64, device="cuda")
+
+    @property
+    def workspace(self) -> torch.Tensor:
+        """Device tensor (int64) that can be passed to the kernel
+        as ``void** workspace``."""
+        return self._workspace
+
+    # ------------------------------------------------------------------
+    # Helpers
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def compute_comm_size_for_minimax(
+        max_tokens: int,
+        world_size: int,
+        fused_qk: bool = True,
+    ) -> int:
+        """
+        Return a safe ``comm_size`` (in bytes) for MiniMaxReduceRMSKernel.
+
+        The kernel stores per-token variance scalars in the Lamport buffer:
+          - single-matrix path: ``world_size × max_tokens × 4`` bytes per slot
+          - fused Q+K path: ``world_size × 2 × ceil(max_tokens/4) × 16`` bytes per slot
+
+        The returned value is rounded up to 2 MiB alignment.
+        """
+        if fused_qk:
+            groups = (max_tokens + 3) // 4
+            slot_bytes = world_size * 2 * groups * 16  # 16 = sizeof(float4)
+        else:
+            slot_bytes = world_size * max_tokens * 4  # 4  = sizeof(float)
+        return ((slot_bytes + _ALIGN - 1) >> 21) << 21
+
+    def cleanup(self):
+        if hasattr(self, "_lamport"):
+            self._lamport.cleanup()
+
+    def __del__(self):
+        if not sys.is_finalizing():
+            self.cleanup()
+
+    def __repr__(self):
+        return (
+            f"LamportWorkspace(rank={self.rank}, world_size={self.world_size}, "
+            f"comm_size={self.comm_size})"
+        )
+
+
+# ---------------------------------------------------------------------------
+# Cached convenience function (mirrors TRT-LLM's get_allreduce_workspace)
+# ---------------------------------------------------------------------------
+
+_cache_lock = threading.Lock()
+_workspace_cache: dict = {}
+
+
+def get_allreduce_workspace(
+    rank: int,
+    world_size: int,
+    comm_size: int | None = None,
+    max_tokens: int = 16384,
+    process_group=None,
+) -> torch.Tensor:
+    """
+    Return a cached workspace tensor for the given (rank, world_size) pair.
+
+    On first call the workspace is allocated and IPC handles are exchanged;
+    subsequent calls with the same arguments return the cached tensor.
+
+    Parameters
+    ----------
+    rank, world_size : int
+        TP rank and TP size.
+    comm_size : int, optional
+        Explicit slot size in bytes.  If ``None``, computed automatically
+        from ``max_tokens`` and ``world_size`` (fused Q+K path).
+    max_tokens : int
+        Maximum number of tokens per batch (used when ``comm_size is None``).
+    process_group : optional
+        ``torch.distributed`` process group.
+    """
+    if comm_size is None:
+        comm_size = LamportWorkspace.compute_comm_size_for_minimax(
+            max_tokens, world_size, fused_qk=True
+        )
+    pg_id = id(process_group) if process_group is not None else 0
+    key = (rank, world_size, comm_size, pg_id)
+    with _cache_lock:
+        if key not in _workspace_cache:
+            ws = LamportWorkspace(rank, world_size, comm_size, process_group)
+            _workspace_cache[key] = ws
+        return _workspace_cache[key].workspace

vllm/model_executor/layers/rotary_embedding/__init__.py

@@ -12,6 +12,7 @@ from .dual_chunk_rope import DualChunkRotaryEmbedding
 from .dynamic_ntk_alpha_rope import DynamicNTKAlphaRotaryEmbedding
 from .dynamic_ntk_scaling_rope import DynamicNTKScalingRotaryEmbedding
 from .fope import FourierRotaryEmbedding
+from .gemma4_rope import Gemma4RotaryEmbedding
 from .linear_scaling_rope import LinearScalingRotaryEmbedding
 from .llama3_rope import Llama3RotaryEmbedding
 from .llama4_vision_rope import Llama4VisionRotaryEmbedding
@@ -134,6 +135,17 @@ def get_rope(
                 is_neox_style,
                 dtype,
             )
+    elif scaling_type == "proportional":
+        # Proportional RoPE is used by Gemma4 for global (full) attention.
+        # Gemma4 uses a sparse/fractional RoPE with cross-mixing between halves.
+        rotary_emb = Gemma4RotaryEmbedding(
+            head_size,
+            rotary_dim,
+            max_position,
+            base,
+            is_neox_style,
+            dtype,
+        )
     elif scaling_type == "llama3":
         scaling_factor = rope_parameters["factor"]
         low_freq_factor = rope_parameters["low_freq_factor"]

vllm/model_executor/layers/rotary_embedding/gemma4_rope.py

@@ -0,0 +1,84 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Gemma4-specific Rotary Positional Embeddings (proportional scaling).
+
+Gemma4 uses "proportional" RoPE which computes inv_freq frequencies scaled
+by head_dim (not rotary_dim), and zero-pads for non-rotated dimensions when
+partial_rotary_factor < 1. The actual rotation uses standard neox-style
+rotate_half, matching HF transformers' apply_rotary_pos_emb.
+"""
+
+import torch
+
+from .base import RotaryEmbedding
+
+
+class Gemma4RotaryEmbedding(RotaryEmbedding):
+    """Gemma4 proportional RoPE.
+
+    Extends RotaryEmbedding (which provides standard neox-style rotation
+    via ops.rotary_embedding CUDA kernel) but overrides the inv_freq
+    computation to match HF's _compute_proportional_rope_parameters:
+    - Frequency exponents use head_dim (not rotary_dim) as denominator
+    - Non-rotated dims are zero-padded (cos=1, sin=0 = identity rotation)
+
+    When partial_rotary_factor=1.0 (the default for some variants), ALL dims are
+    rotated and this is equivalent to standard RotaryEmbedding with
+    head_dim-scaled frequencies.
+    """
+
+    def __init__(
+        self,
+        head_size: int,
+        rotary_dim: int,
+        max_position_embeddings: int,
+        base: float,
+        is_neox_style: bool,
+        dtype: torch.dtype,
+    ) -> None:
+        # Number of rotation angle pairs (from partial_rotary_factor)
+        self.rope_angles = rotary_dim // 2
+        # Non-rotated angle pairs per half
+        self.nope_angles = (head_size // 2) - self.rope_angles
+
+        # Important: set rotary_dim = head_size so the base class's
+        # forward_static applies rotation to ALL dims of the cos/sin cache.
+        # The non-rotated dims will have cos=1, sin=0 (identity) thanks
+        # to our _compute_inv_freq zero-padding.
+        super().__init__(
+            head_size,
+            head_size,  # rotary_dim = head_size (full application)
+            max_position_embeddings,
+            base,
+            is_neox_style,
+            dtype,
+        )
+
+    def _compute_inv_freq(self, base: float) -> torch.Tensor:
+        """Compute frequencies matching HF proportional RoPE.
+
+        Key difference from base: exponent denominator is head_size (not
+        rotary_dim), and non-rotated dims are zero-padded.
+        """
+        # HF formula: base ** (arange(0, 2*rope_angles, 2) / head_dim)
+        freq_exponents = (
+            torch.arange(0, 2 * self.rope_angles, 2, dtype=torch.float) / self.head_size
+        )
+        inv_freq = 1.0 / (base**freq_exponents)
+
+        # Zero-pad for non-rotated dims (identity rotation: cos=1, sin=0)
+        if self.nope_angles > 0:
+            inv_freq = torch.cat(
+                [
+                    inv_freq,
+                    torch.zeros(self.nope_angles, dtype=torch.float),
+                ]
+            )
+        return inv_freq
+
+    def extra_repr(self) -> str:
+        s = f"head_size={self.head_size}, rotary_dim={self.rotary_dim}"
+        s += f", rope_angles={self.rope_angles}, nope_angles={self.nope_angles}"
+        s += f", max_position_embeddings={self.max_position_embeddings}"
+        s += f", base={self.base}, is_neox_style={self.is_neox_style}"
+        return s

vllm/model_executor/layers/sparse_attn_indexer.py

@@ -4,6 +4,7 @@
 
 import torch
 
+import vllm.envs as envs
 from vllm._aiter_ops import rocm_aiter_ops
 from vllm.forward_context import get_forward_context
 from vllm.logger import init_logger
@@ -51,6 +52,14 @@ def sparse_attn_indexer(
             ((total_seq_lens, head_dim), torch.float8_e4m3fn),
             ((total_seq_lens, 4), torch.uint8),
         )
+
+        # Dummy allocation to simulate for peak logits tensor memory during inference.
+        # FP8 elements so elements == bytes
+        max_logits_elems = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024
+        _ = torch.empty(
+            max_logits_elems, dtype=torch.uint8, device=hidden_states.device
+        )
+
         return sparse_attn_indexer_fake(
             hidden_states,
             k_cache_prefix,
@@ -101,13 +110,16 @@ def sparse_attn_indexer(
         for chunk in prefill_metadata.chunks:
             k_fp8 = k_fp8_full[: chunk.total_seq_lens]
             k_scale = k_scale_full[: chunk.total_seq_lens]
-            ops.cp_gather_indexer_k_quant_cache(
-                kv_cache,
-                k_fp8,
-                k_scale,
-                chunk.block_table,
-                chunk.cu_seq_lens,
-            )
+
+            if not chunk.skip_kv_gather:
+                ops.cp_gather_indexer_k_quant_cache(
+                    kv_cache,
+                    k_fp8,
+                    k_scale,
+                    chunk.block_table,
+                    chunk.cu_seq_lens,
+                )
+
             logits = fp8_mqa_logits(
                 q_fp8[chunk.token_start : chunk.token_end],
                 (k_fp8, k_scale.view(torch.float32).flatten()),

vllm/model_executor/model_loader/gguf_loader.py

@@ -209,12 +209,24 @@ class GGUFModelLoader(BaseModelLoader):
                 GGUF tensor name with suffix (e.g., 'mm.soft_emb_norm.weight')
                 or None if no mapping found
             """
+            # In transformers v5, multimodal models (e.g. Gemma3) wrap
+            # all sub-models under an outer 'model.' attribute, producing
+            # state_dict keys like 'model.language_model.layers.0...' and
+            # 'model.vision_tower.vision_model...'.  Strip this outer
+            # prefix so the keys match what gguf-py expects.
+            if is_multimodal and hf_name.startswith("model."):
+                hf_name = hf_name[6:]  # Remove outer 'model.'
+
             # Strip 'language_model.' prefix for multimodal models - gguf-py
             # tensor mappings expect parameter names without this prefix.
             # Note: 'model.' prefix should be KEPT for text-only models as
             # gguf-py expects it.
             if hf_name.startswith("language_model."):
                 hf_name = hf_name[15:]  # Remove 'language_model.'
+                # Re-add 'model.' prefix because gguf-py text tensor maps
+                # expect 'model.layers...' format.
+                if is_multimodal:
+                    hf_name = "model." + hf_name
 
             # Parse parameter name and suffix
             if hf_name.endswith((".weight", ".bias")):

vllm/model_executor/models/config.py

@@ -14,6 +14,7 @@ from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec, MLAAttentio
 if TYPE_CHECKING:
     from vllm.config import ModelConfig, VllmConfig
 
+
 logger = init_logger(__name__)
 
 
@@ -57,6 +58,58 @@ class Gemma3TextModelConfig(VerifyAndUpdateConfig):
         hf_config.is_causal = not hf_config.use_bidirectional_attention
 
 
+class Gemma4Config(VerifyAndUpdateConfig):
+    @staticmethod
+    def verify_and_update_config(vllm_config: "VllmConfig") -> None:
+        """Force unified attention backend for models with heterogeneous
+        head dimensions.
+
+        Some Gemma4 variants use different head dimensions for
+        sliding window (head_dim) vs full attention (global_head_dim) layers.
+        When global_head_dim > 256, FlashAttention rejects those layers
+        (head_size <= 256 kernel limit), causing vLLM to select a different
+        backend for each layer type. This mixed-backend execution produces
+        numerical divergence and output corruption.
+
+        The fix detects heterogeneous head dimensions from the model config
+        and forces TRITON_ATTN (which has no head_size ceiling) for all
+        layers when the user hasn't explicitly chosen a backend.
+
+        TODO: Heterogeneous head_sizes (head_dim != global_head_dim)
+        require NixlConnector changes to support per-layer KV transfer
+        with different head dimensions for prefill-decode disaggregation.
+        """
+        hf_text_config = vllm_config.model_config.hf_text_config
+        head_dim = getattr(hf_text_config, "head_dim", None)
+        global_head_dim = getattr(hf_text_config, "global_head_dim", None)
+
+        # Only force Triton when head dimensions actually differ AND the
+        # larger one exceeds FlashAttention's kernel limit (head_size <= 256).
+        # This avoids unnecessary backend forcing on smaller models where
+        # the config carries global_head_dim but all layers can still use
+        # the same FA backend.
+        max_head_dim = max(head_dim or 0, global_head_dim or 0)
+        if (
+            head_dim is not None
+            and global_head_dim is not None
+            and head_dim != global_head_dim
+            and max_head_dim > 256
+            and vllm_config.attention_config.backend is None
+        ):
+            from vllm.v1.attention.backends.registry import (
+                AttentionBackendEnum,
+            )
+
+            vllm_config.attention_config.backend = AttentionBackendEnum.TRITON_ATTN
+            logger.info(
+                "Gemma4 model has heterogeneous head dimensions "
+                "(head_dim=%d, global_head_dim=%d). Forcing TRITON_ATTN "
+                "backend to prevent mixed-backend numerical divergence.",
+                head_dim,
+                global_head_dim,
+            )
+
+
 class GptOssForCausalLMConfig(VerifyAndUpdateConfig):
     @staticmethod
     def verify_and_update_config(vllm_config: "VllmConfig") -> None:
@@ -668,6 +721,8 @@ MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = {
     "Ernie4_5_VLMoeForConditionalGeneration": Ernie4_5_VLMoeForConditionalGenerationConfig,  # noqa: E501
     "FalconMambaForCausalLM": MambaModelConfig,
     "Gemma3TextModel": Gemma3TextModelConfig,
+    "Gemma4ForCausalLM": Gemma4Config,
+    "Gemma4ForConditionalGeneration": Gemma4Config,
     "GptOssForCausalLM": GptOssForCausalLMConfig,
     "GteModel": SnowflakeGteNewModelConfig,
     "GteNewForSequenceClassification": GteNewModelConfig,

vllm/model_executor/models/gemma4_utils.py

@@ -0,0 +1,292 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Copyright 2025 Google Inc. HuggingFace Inc. team. All rights reserved.
+
+"""Gemma4 output parsing utilities for offline inference.
+
+Standalone functions that parse decoded model text to extract structured
+thinking content and tool calls from Gemma4 models. These are pure-Python
+utilities with zero heavy dependencies — they work on raw decoded strings
+from any inference backend (vLLM, HuggingFace, TGI, etc.).
+
+Usage with vLLM offline inference::
+
+    from vllm import LLM, SamplingParams
+    from vllm.model_executor.models.gemma4_utils import (
+        parse_output,
+        parse_tool_calls,
+    )
+
+    llm = LLM(model="google/gemma-4-it")
+    outputs = llm.generate(prompt, SamplingParams(...))
+    text = tokenizer.decode(outputs[0].outputs[0].token_ids, skip_special_tokens=False)
+
+    # Extract thinking / answer (works with or without enable_thinking)
+    result = parse_output(text)
+    print(result["thinking"])  # chain-of-thought or None
+    print(result["answer"])  # final answer
+
+    # Extract tool calls
+    tool_calls = parse_tool_calls(text)
+    for tc in tool_calls:
+        print(f"{tc['name']}({tc['arguments']})")
+
+Ported from ``transformers.models.gemma4.utils_gemma4`` so that vLLM users
+do not need a transformers dependency for output parsing.
+"""
+
+import json
+
+import regex as re
+
+# ---- Thinking Mode Utility ----
+
+# Thinking delimiter tokens as they appear in decoded text.
+# Gemma4 uses <|channel> (start) and <channel|> (end) as thinking delimiters.
+_THINKING_START_TAG = "<|channel>"
+_THINKING_END_TAG = "<channel|>"
+
+# Sentinel tokens that may appear in decoded output.
+_TURN_END_TAG = "<turn|>"
+
+
+def parse_thinking_output(text: str) -> dict[str, str | None]:
+    """Parse decoded Gemma4 model output.
+
+    Use this on **all** Gemma4 output regardless of whether thinking mode
+    was enabled.  It handles three cases:
+
+    1. **Thinking enabled, tags present** — splits on ``<|channel>``/
+       ``<channel|>`` to separate chain-of-thought from the answer and
+       strips the ``thought\\n`` role label.
+    2. **Thinking disabled, spurious label** — strips the bare
+       ``thought\\n`` prefix that some Gemma4 models emit even
+       without thinking mode.
+    3. **Clean output** — returns the text unchanged.
+
+    The answer text is always cleaned of trailing sentinel tokens
+    (``<turn|>``, ``<eos>``, etc.).
+
+    Args:
+        text: Decoded model output text (from ``tokenizer.decode(...)``).
+
+    Returns:
+        A dict with keys:
+            - ``"thinking"``: The chain-of-thought text, or ``None`` if no
+              thinking delimiters were found.
+            - ``"answer"``: The final answer text.
+
+    Example::
+
+        >>> from vllm.model_executor.models.gemma4_utils import parse_thinking_output
+        >>> output_text = tokenizer.decode(outputs[0], skip_special_tokens=False)
+        >>> result = parse_thinking_output(output_text)
+        >>> print(result["thinking"])  # chain-of-thought reasoning or None
+        >>> print(result["answer"])    # final answer
+    """
+    if _THINKING_END_TAG in text:
+        parts = text.split(_THINKING_END_TAG, 1)
+        thinking_block = parts[0]
+        answer = _clean_answer(parts[1])
+
+        # Extract thinking content: strip the start tag if present
+        if _THINKING_START_TAG in thinking_block:
+            thinking = thinking_block.split(_THINKING_START_TAG, 1)[1]
+        else:
+            thinking = thinking_block
+
+        # Strip the "thought\n" channel role label the model emits inside
+        # <|channel>thought\n...<channel|> (analogous to "user\n" in
+        # <|turn>user\n...<turn|>).
+        thinking = _strip_thought_label(thinking.strip())
+        thinking = thinking.strip()
+
+        return {"thinking": thinking, "answer": answer}
+
+    # No thinking delimiters found.
+    # Strip spurious "thought\n" role label that some Gemma4 models sometimes
+    # emit even without thinking mode enabled, then clean trailing tokens.
+    answer = _strip_thought_label(text)
+    answer = _clean_answer(answer)
+    return {"thinking": None, "answer": answer}
+
+
+def _strip_thought_label(text: str) -> str:
+    """Strip the spurious ``thought\\n`` label from the start of text.
+
+    Only strips when ``thought`` appears as the very first word followed by
+    a newline — preserving the word ``thought`` in any other context.
+    """
+    if text.startswith("thought\n"):
+        return text[len("thought\n") :]
+    return text
+
+
+def _clean_answer(text: str) -> str:
+    """Clean trailing sentinel tokens from the answer text.
+
+    Strips ``<turn|>``, ``<eos>``, and surrounding whitespace that the
+    model appends at the end of its response.
+    """
+    text = text.strip()
+    # Strip trailing <turn|> (Gemma4 turn-end marker)
+    if text.endswith(_TURN_END_TAG):
+        text = text[: -len(_TURN_END_TAG)].rstrip()
+    # Strip trailing <eos> if present
+    if text.endswith("<eos>"):
+        text = text[:-5].rstrip()
+    return text
+
+
+# ---- Tool Call Parsing Utility ----
+#
+# NOTE: For the OpenAI-compatible API server tool parser (streaming +
+# non-streaming), see vllm/tool_parsers/gemma4_tool_parser.py.
+# This module provides offline inference utilities for direct user import.
+
+# Tool call delimiter tokens as they appear in decoded text.
+# Standard format: <|tool_call>call:name{args}<tool_call|>
+_TOOL_CALL_START_TAG = "<|tool_call>"
+_TOOL_CALL_END_TAG = "<tool_call|>"
+_TOOL_RESPONSE_START_TAG = "<|tool_response>"
+
+# Gemma4 escape token as it appears in decoded text.
+_ESCAPE_TOKEN = '<|"|>'
+
+
+def _parse_tool_arguments(args_str: str) -> dict[str, str]:
+    """Parse tool call arguments from the Gemma4 compact format.
+
+    Handles the ``key:<|"|>value<|"|>`` format used by Gemma4, with fallback
+    to heuristic key-value extraction. Also tolerates the slightly different
+    ``key: "value"`` format (space + plain quotes) that some chat templates
+    produce.
+
+    Args:
+        args_str: Raw argument string from inside ``call:name{...}``.
+
+    Returns:
+        Dictionary of argument name → value.
+    """
+    if not args_str or not args_str.strip():
+        return {}
+
+    # Replace Gemma4 escape tokens with standard quotes.
+    cleaned = args_str.replace(_ESCAPE_TOKEN, '"')
+
+    # Try JSON parsing first (handles nested values, arrays, etc.).
+    try:
+        parsed = json.loads("{" + cleaned + "}")
+        # Ensure all values are strings for consistency.
+        return {k: str(v) if not isinstance(v, str) else v for k, v in parsed.items()}
+    except (json.JSONDecodeError, ValueError):
+        pass
+
+    # Fallback: extract key:"value" pairs (allow optional space after colon).
+    arguments = {}
+    for key, value in re.findall(r'(\w+):\s*"([^"]*)"', cleaned):
+        arguments[key] = value
+
+    if not arguments:
+        # Last resort: extract key:value pairs (unquoted).
+        for key, value in re.findall(r"(\w+):\s*([^,}]+)", args_str):
+            arguments[key] = value.strip().strip('"').replace(_ESCAPE_TOKEN, "")
+
+    return arguments
+
+
+def parse_tool_calls(text: str, *, strict: bool = False) -> list[dict]:
+    """Parse tool calls from decoded Gemma4 model output.
+
+    Uses a tiered parsing strategy to handle known output variations in
+    Gemma4 models, which may emit
+    non-standard tool call formats.
+
+    Parsing tiers:
+        1. **Standard**: ``<|tool_call>call:name{args}<tool_call|>``
+           (special token IDs 48/49 in decoded text)
+        2. **Fallback** (when ``strict=False``): bare ``call:name{args}``
+           patterns, including ``<call>name{args}`` (fragmented tokens from
+           multimodal inputs)
+
+    Args:
+        text: Decoded model output text (from ``tokenizer.decode(...,
+            skip_special_tokens=False)``).
+        strict: If ``True``, only match the standard ``<|tool_call>`` format.
+            If ``False`` (default), also try fallback patterns for
+            known Gemma4 output variations.
+
+    Returns:
+        A list of dicts, each with keys:
+            - ``"name"``: The tool function name (e.g. ``"get_weather"``).
+            - ``"arguments"``: A dict of argument name → value.
+
+    Example::
+
+        >>> from vllm.model_executor.models.gemma4_utils import (
+        ...     parse_tool_calls
+        ... )
+        >>> output = tokenizer.decode(outputs[0], skip_special_tokens=False)
+        >>> tool_calls = parse_tool_calls(output)
+        >>> for tc in tool_calls:
+        ...     print(f"Call: {tc['name']}({tc['arguments']})")
+    """
+    results = []
+
+    # Tier 1: Standard format with special tokens.
+    # <|tool_call>call:name{args}<tool_call|>
+    # Note: Some Gemma4 models emit <turn|> instead of <tool_call|>.
+    standard_pattern = r"<\|tool_call\>call:(\w+)\{(.*?)\}(?:<tool_call\|>|<turn\|>)"
+    for match in re.finditer(standard_pattern, text, re.DOTALL):
+        name, args_str = match.group(1), match.group(2)
+        results.append(
+            {
+                "name": name,
+                "arguments": _parse_tool_arguments(args_str),
+            }
+        )
+
+    if results or strict:
+        return results
+
+    # Tier 2: Fallback for known Gemma4 output variations.
+    # Matches: <call>name{args}, call:name{args}, or bare call:name{args}<eos>
+    fallback_pattern = r"(?:<call>|(?:^|\s)call:)(\w+)\{(.*?)\}"
+    for match in re.finditer(fallback_pattern, text, re.DOTALL):
+        name, args_str = match.group(1), match.group(2)
+        results.append(
+            {
+                "name": name,
+                "arguments": _parse_tool_arguments(args_str),
+            }
+        )
+
+    return results
+
+
+def has_tool_response_tag(text: str) -> bool:
+    """Check if model output properly ends with a tool response tag.
+
+    Some Gemma4 models sometimes emit ``<eos>`` instead of
+    ``<|tool_response>`` after a tool call. This helper detects
+    whether the model used the proper termination, so callers can
+    decide whether to inject ``<|tool_response>`` into the next prompt.
+
+    Args:
+        text: Decoded model output text.
+
+    Returns:
+        ``True`` if the output ends with ``<|tool_response>``
+        (proper behavior), ``False`` otherwise.
+
+    Example::
+
+        >>> from vllm.model_executor.models.gemma4_utils import (
+        ...     has_tool_response_tag
+        ... )
+        >>> if not has_tool_response_tag(model_output):
+        ...     # Model used <eos> instead — inject <|tool_response> manually
+        ...     next_prompt = "<|tool_response>" + tool_result
+    """
+    stripped = text.rstrip()
+    return stripped.endswith(_TOOL_RESPONSE_START_TAG)

vllm/model_executor/models/kimi_k25.py

@@ -113,7 +113,29 @@ class KimiK25ProcessingInfo(BaseProcessingInfo):
             trust_remote_code=self.ctx.model_config.trust_remote_code,
         )
 
-        self.media_token_id = media_token_id = hf_config.media_placeholder_token_id
+        # Resolve token ID from the tokenizer because transformers v5
+        # may remap token IDs vs config.json.
+        config_token_id = hf_config.media_placeholder_token_id
+        resolved_token_id = tokenizer.convert_tokens_to_ids("<|media_pad|>")
+        is_valid_resolved = isinstance(resolved_token_id, int) and (
+            tokenizer.unk_token_id is None
+            or resolved_token_id != tokenizer.unk_token_id
+        )
+        if is_valid_resolved and resolved_token_id != config_token_id:
+            logger.warning_once(
+                "Kimi-K2.5 config.media_placeholder_token_id (%d) disagrees "
+                "with tokenizer mapping for <|media_pad|> (%d). "
+                "Using tokenizer value.",
+                config_token_id,
+                resolved_token_id,
+            )
+            media_token_id = resolved_token_id
+            # Patch config so downstream code also sees the correct ID.
+            hf_config.media_placeholder_token_id = resolved_token_id
+        else:
+            media_token_id = config_token_id
+
+        self.media_token_id = media_token_id
         self.media_token = tokenizer.decode(media_token_id)
 
         self.image_processor = image_processor
@@ -232,8 +254,7 @@ class KimiK25MultiModalProcessor(BaseMultiModalProcessor[KimiK25ProcessingInfo])
         hf_processor_mm_kwargs: Mapping[str, Any],
         out_mm_kwargs: MultiModalKwargsItems,
     ) -> Sequence[PromptUpdate]:
-        hf_config = self.info.get_hf_config()
-        media_token_id = hf_config.media_placeholder_token_id
+        media_token_id = self.info.media_token_id
 
         def get_replacement(item_idx: int):
             media = mm_items.get_items("vision_chunk", (VisionChunkProcessorItems,))

vllm/model_executor/models/minimax_m2.py

@@ -232,9 +232,7 @@ class MiniMaxM2Attention(nn.Module):
     ) -> torch.Tensor:
         qkv, _ = self.qkv_proj(hidden_states)
         q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
-        q, k = MiniMaxText01RMSNormTP.forward_qk(
-            self.q_norm, self.k_norm, q.contiguous(), k.contiguous()
-        )
+        q, k = MiniMaxText01RMSNormTP.forward_qk(self.q_norm, self.k_norm, q, k)
         q, k = self.rotary_emb(positions, q, k)
         attn_output = self.attn(q, k, v)
         output, _ = self.o_proj(attn_output)

vllm/model_executor/models/musicflamingo.py

@@ -32,9 +32,9 @@ from transformers.models.musicflamingo import (
 
 from vllm.config import VllmConfig
 from vllm.config.multimodal import BaseDummyOptions
+from vllm.inputs import MultiModalDataDict
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.inputs import (
-    MultiModalDataDict,
     MultiModalFieldConfig,
     MultiModalKwargsItems,
 )

vllm/model_executor/models/registry.py

@@ -111,6 +111,7 @@ _TEXT_GENERATION_MODELS = {
     "Gemma2ForCausalLM": ("gemma2", "Gemma2ForCausalLM"),
     "Gemma3ForCausalLM": ("gemma3", "Gemma3ForCausalLM"),
     "Gemma3nForCausalLM": ("gemma3n", "Gemma3nForCausalLM"),
+    "Gemma4ForCausalLM": ("gemma4", "Gemma4ForCausalLM"),
     "Qwen3NextForCausalLM": ("qwen3_next", "Qwen3NextForCausalLM"),
     "GlmForCausalLM": ("glm", "GlmForCausalLM"),
     "Glm4ForCausalLM": ("glm4", "Glm4ForCausalLM"),
@@ -381,6 +382,7 @@ _MULTIMODAL_MODELS = {
         "gemma3n_mm",
         "Gemma3nForConditionalGeneration",
     ),
+    "Gemma4ForConditionalGeneration": ("gemma4_mm", "Gemma4ForConditionalGeneration"),
     "GlmAsrForConditionalGeneration": ("glmasr", "GlmAsrForConditionalGeneration"),
     "GLM4VForCausalLM": ("glm4v", "GLM4VForCausalLM"),
     "Glm4vForConditionalGeneration": ("glm4_1v", "Glm4vForConditionalGeneration"),

vllm/model_executor/models/transformers/__init__.py

@@ -16,13 +16,11 @@
 # limitations under the License.
 """Wrapper around `transformers` models"""
 
-from vllm.compilation.decorators import support_torch_compile
 from vllm.model_executor.models.transformers.base import Base
 from vllm.model_executor.models.transformers.causal import CausalMixin
 from vllm.model_executor.models.transformers.legacy import LegacyMixin
 from vllm.model_executor.models.transformers.moe import MoEMixin
 from vllm.model_executor.models.transformers.multimodal import (
-    DYNAMIC_ARG_DIMS,
     MultiModalDummyInputsBuilder,
     MultiModalMixin,
     MultiModalProcessingInfo,
@@ -32,16 +30,13 @@ from vllm.model_executor.models.transformers.pooling import (
     EmbeddingMixin,
     SequenceClassificationMixin,
 )
-from vllm.model_executor.models.transformers.utils import can_enable_torch_compile
 from vllm.multimodal import MULTIMODAL_REGISTRY
 
 
 # Text only models
-@support_torch_compile(enable_if=can_enable_torch_compile)
 class TransformersForCausalLM(CausalMixin, Base): ...
 
 
-@support_torch_compile(enable_if=can_enable_torch_compile)
 class TransformersMoEForCausalLM(MoEMixin, CausalMixin, Base): ...
 
 
@@ -51,9 +46,6 @@ class TransformersMoEForCausalLM(MoEMixin, CausalMixin, Base): ...
     info=MultiModalProcessingInfo,
     dummy_inputs=MultiModalDummyInputsBuilder,
 )
-@support_torch_compile(
-    dynamic_arg_dims=DYNAMIC_ARG_DIMS, enable_if=can_enable_torch_compile
-)
 class TransformersMultiModalForCausalLM(MultiModalMixin, CausalMixin, Base): ...
 
 
@@ -62,20 +54,15 @@ class TransformersMultiModalForCausalLM(MultiModalMixin, CausalMixin, Base): ...
     info=MultiModalProcessingInfo,
     dummy_inputs=MultiModalDummyInputsBuilder,
 )
-@support_torch_compile(
-    dynamic_arg_dims=DYNAMIC_ARG_DIMS, enable_if=can_enable_torch_compile
-)
 class TransformersMultiModalMoEForCausalLM(
     MoEMixin, MultiModalMixin, CausalMixin, Base
 ): ...
 
 
 # Embedding models
-@support_torch_compile(enable_if=can_enable_torch_compile)
 class TransformersEmbeddingModel(EmbeddingMixin, LegacyMixin, Base): ...
 
 
-@support_torch_compile(enable_if=can_enable_torch_compile)
 class TransformersMoEEmbeddingModel(EmbeddingMixin, MoEMixin, Base): ...
 
 
@@ -84,20 +71,15 @@ class TransformersMoEEmbeddingModel(EmbeddingMixin, MoEMixin, Base): ...
     info=MultiModalProcessingInfo,
     dummy_inputs=MultiModalDummyInputsBuilder,
 )
-@support_torch_compile(
-    dynamic_arg_dims=DYNAMIC_ARG_DIMS, enable_if=can_enable_torch_compile
-)
 class TransformersMultiModalEmbeddingModel(EmbeddingMixin, MultiModalMixin, Base): ...
 
 
 # Sequence classification models
-@support_torch_compile(enable_if=can_enable_torch_compile)
 class TransformersForSequenceClassification(
     SequenceClassificationMixin, LegacyMixin, Base
 ): ...
 
 
-@support_torch_compile(enable_if=can_enable_torch_compile)
 class TransformersMoEForSequenceClassification(
     SequenceClassificationMixin, MoEMixin, Base
 ): ...
@@ -108,9 +90,6 @@ class TransformersMoEForSequenceClassification(
     info=MultiModalProcessingInfo,
     dummy_inputs=MultiModalDummyInputsBuilder,
 )
-@support_torch_compile(
-    dynamic_arg_dims=DYNAMIC_ARG_DIMS, enable_if=can_enable_torch_compile
-)
 class TransformersMultiModalForSequenceClassification(
     SequenceClassificationMixin, MultiModalMixin, Base
 ): ...

vllm/model_executor/models/transformers/base.py

@@ -16,6 +16,7 @@
 # limitations under the License.
 """Transformers modeling backend base class."""
 
+import sys
 from collections.abc import Callable, Iterable
 from itertools import chain
 from operator import attrgetter
@@ -29,6 +30,7 @@ from torch import nn
 from transformers import AutoModel
 from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
 
+from vllm.compilation.decorators import support_torch_compile
 from vllm.config.utils import getattr_iter
 from vllm.distributed import get_pp_group, get_tp_group
 from vllm.distributed.utils import get_pp_indices
@@ -47,6 +49,7 @@ from vllm.model_executor.models.interfaces import (
 )
 from vllm.model_executor.models.interfaces_base import VllmModel
 from vllm.model_executor.models.transformers.utils import (
+    can_enable_torch_compile,
     get_feature_request_tip,
     init_on_device_without_buffers,
     log_replacement,
@@ -117,6 +120,7 @@ class Base(
         self.config = vllm_config.model_config.hf_config
         self.text_config = self.config.get_text_config()
         self.cache_config = vllm_config.cache_config
+        self.compilation_config = vllm_config.compilation_config
         self.device_config = vllm_config.device_config
         self.model_config = vllm_config.model_config
         self.parallel_config = vllm_config.parallel_config
@@ -146,7 +150,7 @@ class Base(
         if self.quant_config:
             quant_method_name = self.quant_config.get_name()
             # Check for unsupported quantization methods.
-            if quant_method_name == "mxfp4":
+            if quant_method_name in ("mxfp4", "gpt_oss_mxfp4"):
                 raise NotImplementedError(
                     "Transformers modeling backend does "
                     "not support MXFP4 quantization yet."
@@ -155,14 +159,16 @@ class Base(
             if "gptq" in quant_method_name:
                 self.ignore_unexpected_suffixes.append(".bias")
 
-        # Patch config and init on "meta" to delay allocating GPU tensors
         self._patch_config()
+        from_config_kwargs = dict(
+            config=self.config,
+            dtype=self.model_config.dtype,
+            trust_remote_code=self.model_config.trust_remote_code,
+        )
+        self._decorate_for_torch_compile(**from_config_kwargs)
+        # Init on "meta" to delay allocating GPU tensors
         with init_on_device_without_buffers("meta"):
-            self.model: PreTrainedModel = AutoModel.from_config(
-                self.config,
-                dtype=self.model_config.dtype,
-                trust_remote_code=self.model_config.trust_remote_code,
-            )
+            self.model: PreTrainedModel = AutoModel.from_config(**from_config_kwargs)
 
         # Create weight name to module qualname mapper
         self._create_hf_to_vllm_mapper()
@@ -218,6 +224,87 @@ class Base(
             if sub_config.dtype != (dtype := self.config.dtype):
                 sub_config.dtype = dtype
 
+    def _get_decoder_cls(self, **kwargs: dict) -> type[PreTrainedModel]:
+        """
+        Get the decoder class from the model.
+
+        Args:
+            kwargs: The kwargs to create the model.
+
+        Returns:
+            The decoder class.
+        """
+        with torch.device("meta"):
+            model: PreTrainedModel = AutoModel.from_config(**kwargs)
+        decoder_cls = type(model.get_decoder())
+        logger.debug("Identified decoder class as: %s", decoder_cls)
+        del model
+        return decoder_cls
+
+    def _decorate_cls_for_torch_compile(
+        self,
+        cls: type[PreTrainedModel],
+        dynamic_arg_dims: dict[str, int] | None,
+        enable_if: Callable[["VllmConfig"], bool],
+        is_encoder: bool,
+    ):
+        """
+        Decorate `cls` to indicate to vLLM that it supports torch compile.
+
+        Args:
+            cls: The PreTrainedModel class to decorate.
+            dynamic_arg_dims: A mapping from argument name to the dynamic dimensions
+                of the argument. If None, default dynamic arg dims will be used. See
+                [`support_torch_compile`][vllm.compilation.decorators.support_torch_compile]
+                for more details.
+            enable_if: A function which takes in the vLLM config and returns whether
+                torch compile should be enabled for this class.
+            is_encoder: Whether the class being decorated is an encoder.
+        """
+        logger.debug(
+            "Decorating `%s` as %s for torch compile with dynamic_arg_dims of %s",
+            cls.__name__,
+            "encoder" if is_encoder else "decoder",
+            dynamic_arg_dims,
+        )
+
+        @support_torch_compile(
+            dynamic_arg_dims=dynamic_arg_dims,
+            enable_if=enable_if,
+            is_encoder=is_encoder,
+        )
+        class SupportTorchCompileWrapper(cls): ...
+
+        # Preserve __module__ so transformers v5's source-file checks
+        # (e.g. _can_set_experts_implementation) read the original
+        # model's module instead of this file.
+        SupportTorchCompileWrapper.__module__ = cls.__module__
+
+        # Patch the class in its module
+        module = sys.modules[cls.__module__]
+        setattr(module, cls.__name__, SupportTorchCompileWrapper)
+
+    def _decorate_for_torch_compile(self, **kwargs: dict):
+        """
+        Decorate the model's decoder class to indicate to vLLM that it supports torch
+        compile if `can_enable_torch_compile` is True.
+
+        Args:
+            kwargs: The kwargs to create the model, which are needed to get the decoder
+                class.
+        """
+        self._decorate_cls_for_torch_compile(
+            cls=self._get_decoder_cls(**kwargs),
+            # Applied to a PreTrainedModel so the batch dimension will exist
+            dynamic_arg_dims=dict[str, int](
+                input_ids=1,  # shape: [1, seq_len]
+                inputs_embeds=1,  # shape: [1, seq_len, hidden_size]
+                position_ids=-1,  # shape: [1, seq_len] or [3, 1, seq_len] for mrope
+            ),
+            enable_if=can_enable_torch_compile,
+            is_encoder=False,
+        )
+
     def _create_hf_to_vllm_mapper(self):
         """
         Create a WeightsMapper to map checkpoint weight names to module qualnames.
@@ -553,11 +640,6 @@ class Base(
             input_ids = None
             inputs_embeds = intermediate_tensors["hidden_states"]
 
-        if input_ids is not None:
-            input_ids = input_ids[None, ...]
-        if inputs_embeds is not None:
-            inputs_embeds = inputs_embeds[None, ...]
-
         # If the model scales embeddings inside the input embedding layer we must
         # ensure they are scaled here since VocabParallelEmbedding will not do it
         if (
@@ -568,22 +650,29 @@ class Base(
             inputs_embeds = self.embed_input_ids(input_ids)
             input_ids = None
 
-        if self.model_config.uses_mrope:
-            position_ids = positions[:, None]
-        else:
-            position_ids = positions[None, ...]
+        # Add batch dimension before entering Transformers model
+        if input_ids is not None and input_ids.ndim == 1:
+            # [seq_len] -> [1, seq_len]
+            input_ids = input_ids[None, ...]
+        if inputs_embeds is not None and inputs_embeds.ndim == 2:
+            # [seq_len, hidden_size] -> [1, seq_len, hidden_size]
+            inputs_embeds = inputs_embeds[None, ...]
+        if positions.ndim == 1:
+            # [seq_len] -> [1, seq_len]
+            positions = positions[None, ...]
 
         outputs = self.model(
             input_ids=input_ids,
             inputs_embeds=inputs_embeds,
             use_cache=False,
-            position_ids=position_ids,
+            position_ids=positions,
             attention_instances=self.attention_instances,
             return_dict=False,
             **self._output_aux_hidden_states_kwargs,
             **kwargs,
         )
-        # We must remove the batch dimension from these outputs
+
+        # Remove batch dimension after exiting Transformers model
         hidden_states = outputs[0][0, ...]
         if self._output_aux_hidden_states_kwargs:
             aux_hidden_states = [x[0][0, ...] for x in outputs[1:]]