Compare commits

...
19 Commits
Author SHA1 Message Date
zhanqiuhuandkhluu b1388b1fbf [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>
2026-04-17 17:57:52 -07:00
wang.yuqiandkhluu 9d780002bd [CI Failure] pin nomic-embed-text-v1 revision (#39292)
Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io>
2026-04-16 18:27:36 -07:00
khluu 39602ebf8a [Bugfix] Fix CPU Dockerfile path after cherry-pick
Cherry-pick 7bd3f40dda (#39557) used the upstream requirements/test/
nested layout, but this branch still has the flat requirements/
layout (cpu-test.in not test/cpu.in). Revert the path rename while
keeping the torch==2.11.0 and sentence-transformers==5.3.0 fixes.

Signed-off-by: khluu <khluu000@gmail.com>
2026-04-16 17:46:27 -07:00
Ilya Boytsovandkhluu 7225a69c9c fix(test): recompute Jina ColBERT rotary inv_freq cleared by transformers v5 weight loader (#39176)
Signed-off-by: Ilya Boytsov <ilyaboytsov1805@gmail.com>
2026-04-16 17:25:30 -07:00
Harry Mellorandkhluu 9a3a31fd2d Don't compile vision encoder for Transformers backend (#30518)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

Signed-off-by: khluu <khluu000@gmail.com>
2026-04-16 17:09:39 -07:00
Li, Jiangandkhluu 7bd3f40dda [CI/Build] Fix sentence-transformers version in CPU test (#39557)
Signed-off-by: jiang1.li <jiang1.li@intel.com>
(cherry picked from commit eb92ba740a)

Signed-off-by: khluu <khluu000@gmail.com>
2026-04-16 17:05:39 -07:00
Jee Jee Liandkhluu c5460385f1 [Kernel] Porting the TRTLLM minimax_allreduce_rms kernels (#37045)
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
(cherry picked from commit ecd1ea1363)

Signed-off-by: khluu <khluu000@gmail.com>
2026-04-16 15:52:31 -07:00
r266-techandkhluu 4bbb8faa1f fix(kimi_k25): resolve media_placeholder_token_id from tokenizer (#39344)
Signed-off-by: r266-tech <r266.tech@gmail.com>
Signed-off-by: Roger Wang <hey@rogerw.io>
Co-authored-by: Roger Wang <hey@rogerw.io>
(cherry picked from commit 17e787a779)
2026-04-16 15:48:34 -07:00
Harry Mellorandkhluu 459d9b38d3 Update to transformers v5 (#30566)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Signed-off-by: khluu <khluu000@gmail.com>
Signed-off-by: Kevin H. Luu <khluu000@gmail.com>
Signed-off-by: Cyrus Leung <cyrus.tl.leung@gmail.com>
Co-authored-by: khluu <khluu000@gmail.com>
Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com>
Co-authored-by: jiang1.li <jiang1.li@intel.com>
(cherry picked from commit 03f8d3a548)
2026-04-16 15:47:49 -07:00
Luciano Martinsandkhluu b1568cf464 [Model] Fix Gemma 4 token repetition by dynamic BOS injection for PT models (#39842)
Signed-off-by: Luciano Martins <lucianommartins@users.noreply.github.com>
Co-authored-by: Luciano Martins <lucianommartins@users.noreply.github.com>
(cherry picked from commit 6dc9491406)
2026-04-15 17:54:37 -07:00
김의진andkhluu a4ac72ceba [Bugfix] Fix Gemma4 tool parser converting bare null to string "null" (#39679)
Signed-off-by: KimuGenie <baby11686@naver.com>
Signed-off-by: Chauncey <chaunceyjiang@gmail.com>
Co-authored-by: Chauncey <chaunceyjiang@gmail.com>
(cherry picked from commit b075604da1)
2026-04-15 17:54:04 -07:00
ShubyMandkhluu 0ab0f70aa9 [Gemma4][Bugfix]: Enable Gemma4ForCasualLM to load lora adapters correctly (#38844)
Signed-off-by: ShubyM <shubymishra20@gmail.com>
(cherry picked from commit 92feb9991d)
2026-04-15 17:53:56 -07:00
Fynn Schmitt-Ulmsandkhluu ee642f8753 Add Gemma4 Eagle3 support (#39450)
Signed-off-by: Rahul-Tuli <rtuli@redhat.com>
Signed-off-by: Fynn Schmitt-Ulms <fschmitt@redhat.com>
Co-authored-by: Rahul-Tuli <rtuli@redhat.com>
Co-authored-by: Claude <noreply@anthropic.com>
(cherry picked from commit e7cfd7c5b9)
2026-04-15 17:53:46 -07:00
Dipika Sikkaandkhluu 6db56c0997 [Gemma4] Support quantized MoE (#39045)
Signed-off-by: Dipika Sikka <dipikasikka1@gmail.com>
(cherry picked from commit 3aecdf08b4)
2026-04-15 17:52:46 -07:00
Ben Browningandkhluu 9a234c7adc [Tool] adjust_request to reasoning parser, and Gemma4 fixes (#39027)
Signed-off-by: Ben Browning <bbrownin@redhat.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
(cherry picked from commit 8477fe427d)
(cherry picked from commit 0d784f4677)
2026-04-15 17:51:11 -07:00
Flora Fengandkhluu f56ffafdce [Bugfix] Fix Gemma4 streaming tool call corruption for split boolean/number values (#39114)
Signed-off-by: sfeng33 <4florafeng@gmail.com>
(cherry picked from commit 13151a4df4)
(cherry picked from commit a7b392731f)
2026-04-15 17:51:03 -07:00
yokeandkhluu 0329e8c9ac [Bugfix][Frontend] Fix Gemma4 streaming HTML duplication after tool calls (#38909)
Signed-off-by: yoke233 <yoke2012@gmail.com>
(cherry picked from commit d734445fcd)
(cherry picked from commit bae948a9a9)
2026-04-15 17:50:50 -07:00
Lucas Wilkinsonandkhluu 6f3dc4d0aa [Gemma4] Enable Fast Prefill Optimization (#38879)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
(cherry picked from commit 47e605092b)
(cherry picked from commit fc29ef13a0)
2026-04-15 17:50:41 -07:00
Greg Pereiraandkhluu a6ac49a33b [Bugfix] Fix invalid JSON in Gemma 4 streaming tool calls by stripping partial delimiters (#38992)
Signed-off-by: greg pereira <grpereir@redhat.com>
Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
(cherry picked from commit f53fa26e05)
(cherry picked from commit 10a26d1d9a)
2026-04-15 17:50:31 -07:00
82 changed files with 5061 additions and 269 deletions
@@ -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}"
+14 -1
View File
@@ -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
+15
View File
@@ -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
+2
View File
@@ -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.
+879
View File
@@ -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};
}
+79
View File
@@ -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
+12
View File
@@ -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
+23
View File
@@ -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
}
+5 -4
View File
@@ -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
+9 -1
View File
@@ -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
+4 -3
View File
@@ -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
+4 -3
View File
@@ -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
@@ -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
+331
View File
@@ -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 -%}
+2 -2
View File
@@ -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
+2
View File
@@ -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
+2 -2
View File
@@ -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
+2 -3
View File
@@ -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
+7 -8
View File
@@ -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
+3 -3
View File
@@ -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
+10 -10
View File
@@ -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
+736
View File
@@ -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
+2
View File
@@ -9,6 +9,8 @@ pytest-shard
# --- Core Tools & Bindings ---
absl-py
arctic-inference
lm_eval[api]
modelscope
# --- Audio Processing ---
librosa
+9
View File
@@ -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,
@@ -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,
)
+13
View File
@@ -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
+23
View File
@@ -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,
)
+11
View File
@@ -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 = (
-18
View File
@@ -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()
@@ -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
@@ -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
@@ -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,
@@ -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:
@@ -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,
@@ -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,
),
]
@@ -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(
@@ -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])
@@ -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",
@@ -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
@@ -397,14 +404,14 @@ VLM_TEST_SETTINGS = {
"gemma4": VLMTestInfo(
models=["google/gemma-4-E2B-it"],
test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE),
prompt_formatter=lambda img_prompt: f"<bos><start_of_turn>user\n{img_prompt}<end_of_turn>\n<start_of_turn>model\n", # noqa: E501
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": "What's the content in the center of the image?",
"cherry_blossom": "What is the season?",
"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="Describe the two images in detail.",
multi_image_prompt="<|image|><|image|>Describe the two images in detail.", # noqa: E501
max_model_len=4096,
max_num_seqs=2,
auto_cls=AutoModelForImageTextToText,
@@ -533,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=[
@@ -545,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"],
@@ -591,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"],
@@ -806,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(
@@ -960,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"],
@@ -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])
@@ -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,
)
@@ -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.
@@ -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,
@@ -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",
@@ -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"]
@@ -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 = {
@@ -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,
+135 -8
View File
@@ -334,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
@@ -469,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",
@@ -509,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",
@@ -544,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),
@@ -754,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"),
@@ -800,9 +835,30 @@ _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
@@ -844,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",
@@ -863,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",
@@ -952,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={
@@ -960,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(
@@ -996,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",
@@ -1014,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",
@@ -1038,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
),
@@ -1133,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",
+11 -1
View File
@@ -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,
@@ -4,6 +4,9 @@
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
@@ -100,6 +103,39 @@ NEW_LINE_STREAMING = {
"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"),
@@ -120,17 +156,22 @@ TEST_CASES = [
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"),
]
@pytest.mark.parametrize("streaming, param_dict", TEST_CASES)
def test_gemma4_reasoning(
streaming: bool,
param_dict: dict,
generic_tokenizer,
):
output = param_dict["output"]
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>"]
@@ -176,6 +217,18 @@ def test_gemma4_reasoning(
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
)
@@ -194,3 +247,29 @@ def test_gemma4_reasoning(
# 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
@@ -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 = {
@@ -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
@@ -85,6 +85,14 @@ class TestParseGemma4Args:
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'
@@ -114,6 +122,19 @@ class TestParseGemma4Args:
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):
@@ -491,6 +512,51 @@ class TestStreamingExtraction:
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 = [
@@ -502,3 +568,119 @@ class TestStreamingExtraction:
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
+5 -1
View File
@@ -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(
(
+35
View File
@@ -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),
)
+17 -2
View File
@@ -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
@@ -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)
+4
View File
@@ -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():
+7 -4
View File
@@ -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
+2
View File
@@ -805,6 +805,8 @@ class SpeculativeConfig:
"deepseek_v3",
"kimi_k2",
"kimi_k25",
"minimax_m2",
"gemma4",
]
if (
self.method in ("eagle3", "extract_hidden_states")
+16
View File
@@ -1577,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)
+2 -1
View File
@@ -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(
+2
View File
@@ -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,
)
@@ -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
+15
View File
@@ -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
@@ -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
@@ -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")):
+442 -68
View File
@@ -19,6 +19,7 @@
"""Gemma 4 model implementation for vLLM."""
from collections.abc import Iterable
from dataclasses import replace
from itertools import islice
import regex as re
@@ -32,6 +33,7 @@ from vllm.distributed import (
get_tensor_model_parallel_rank,
get_tensor_model_parallel_world_size,
)
from vllm.forward_context import get_forward_context
from vllm.logger import init_logger
from vllm.model_executor.layers.activation import GeluAndMul
from vllm.model_executor.layers.attention import Attention
@@ -56,10 +58,18 @@ from vllm.model_executor.model_loader.weight_utils import (
maybe_remap_kv_scale_name,
)
from vllm.sequence import IntermediateTensors
from vllm.v1.attention.backends.utils import KVSharingFastPrefillMetadata
from .interfaces import MixtureOfExperts, SupportsLoRA, SupportsPP
from .interfaces import (
EagleModelMixin,
MixtureOfExperts,
SupportsEagle3,
SupportsLoRA,
SupportsPP,
)
from .utils import (
AutoWeightsLoader,
WeightsMapper,
extract_layer_index,
is_pp_missing_parameter,
make_layers,
@@ -636,8 +646,206 @@ class Gemma4DecoderLayer(nn.Module):
return hidden_states, None
@support_torch_compile
class Gemma4Model(nn.Module):
def _run_decoder_layers(
decoder_layers: list[Gemma4DecoderLayer],
layer_idx_start: int,
positions: torch.Tensor,
hidden_states: torch.Tensor,
per_layer_inputs: torch.Tensor | None = None,
**kwargs,
) -> torch.Tensor:
"""Run a slice of decoder layers with PLE extraction."""
residual = None
for idx, layer in enumerate(decoder_layers):
layer_idx = idx + layer_idx_start
layer_per_input = (
per_layer_inputs[:, layer_idx, :] if per_layer_inputs is not None else None
)
hidden_states, residual = layer(
positions,
hidden_states,
residual,
per_layer_input=layer_per_input,
**kwargs,
)
return hidden_states
@support_torch_compile(
enable_if=lambda vllm_config: vllm_config.cache_config.kv_sharing_fast_prefill
)
class Gemma4SelfDecoderLayers(nn.Module):
"""Compiled wrapper: embedding + non-KV-shared layers (YOCO first half).
Owns the embedding and PLE modules so they are inside the compiled
graph. Gemma4Model delegates embedding methods here.
"""
def __init__(
self,
*,
vllm_config: VllmConfig,
prefix: str = "",
decoder_layers: list[Gemma4DecoderLayer],
layer_idx_start: int,
embed_tokens: VocabParallelEmbedding,
normalizer: torch.Tensor,
embed_tokens_per_layer: VocabParallelEmbedding | None,
embed_scale_per_layer: torch.Tensor | None,
per_layer_model_projection: ColumnParallelLinear | None,
per_layer_projection_norm: RMSNorm | None,
per_layer_input_scale: torch.Tensor | None,
per_layer_projection_scale: torch.Tensor | None,
):
super().__init__()
self.decoder_layers = decoder_layers
self.layer_idx_start = layer_idx_start
config = _get_text_config(vllm_config.model_config.hf_config)
self.config = config
self.hidden_size_per_layer_input = getattr(
config, "hidden_size_per_layer_input", 0
)
self.vocab_size_per_layer_input = getattr(
config, "vocab_size_per_layer_input", config.vocab_size
)
# Shared references to modules owned by Gemma4Model — must be
# inside this nn.Module so torch.compile captures them.
self.embed_tokens = embed_tokens
self.normalizer = normalizer
self.embed_tokens_per_layer = embed_tokens_per_layer
self.embed_scale_per_layer = embed_scale_per_layer
self.per_layer_model_projection = per_layer_model_projection
self.per_layer_projection_norm = per_layer_projection_norm
self.per_layer_input_scale = per_layer_input_scale
self.per_layer_projection_scale = per_layer_projection_scale
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
return self.embed_tokens(input_ids) * self.normalizer
def get_per_layer_inputs(self, input_ids: torch.Tensor) -> torch.Tensor | None:
"""Get per-layer embeddings from embed_tokens_per_layer.
Returns:
Per-layer embeddings (num_tokens, num_layers,
hidden_size_per_layer_input)
"""
if self.embed_tokens_per_layer is None:
return None
per_layer_inputs_mask = torch.logical_and(
input_ids >= 0,
input_ids < self.vocab_size_per_layer_input,
)
per_layer_inputs_tokens = torch.where(
per_layer_inputs_mask, input_ids, torch.zeros_like(input_ids)
)
per_layer_embeds = self.embed_tokens_per_layer(per_layer_inputs_tokens)
per_layer_embeds = per_layer_embeds * self.embed_scale_per_layer
return per_layer_embeds.reshape(
*input_ids.shape,
self.config.num_hidden_layers,
self.hidden_size_per_layer_input,
)
def project_per_layer_inputs(
self,
inputs_embeds: torch.Tensor,
per_layer_inputs: torch.Tensor | None,
) -> torch.Tensor | None:
"""Project inputs_embeds and combine with per_layer_inputs.
Steps:
1. Project inputs_embeds: hidden_size total_ple_dim
2. Scale by hidden_size^{-0.5}
3. Reshape to (num_tokens, num_layers, per_layer_dim)
4. Normalize with per_layer_projection_norm
5. Combine: (projection + per_layer_inputs) * 1/sqrt(2)
"""
if self.per_layer_model_projection is None:
return None
per_layer_projection = self.per_layer_model_projection(inputs_embeds)
per_layer_projection = per_layer_projection * self.per_layer_projection_scale
per_layer_projection = per_layer_projection.reshape(
*inputs_embeds.shape[:-1],
self.config.num_hidden_layers,
self.hidden_size_per_layer_input,
)
per_layer_projection = self.per_layer_projection_norm(per_layer_projection)
if per_layer_inputs is None:
return per_layer_projection
return (per_layer_projection + per_layer_inputs) * self.per_layer_input_scale
def forward(
self,
input_ids: torch.Tensor | None,
positions: torch.Tensor,
inputs_embeds: torch.Tensor | None = None,
per_layer_inputs: torch.Tensor | None = None,
**kwargs,
) -> tuple[torch.Tensor, torch.Tensor | None]:
if inputs_embeds is not None:
hidden_states = inputs_embeds
per_layer_inputs = self.project_per_layer_inputs(
hidden_states, per_layer_inputs
)
else:
hidden_states = self.embed_input_ids(input_ids)
per_layer_embeds = self.get_per_layer_inputs(input_ids)
per_layer_inputs = self.project_per_layer_inputs(
hidden_states, per_layer_embeds
)
hidden_states = _run_decoder_layers(
self.decoder_layers,
self.layer_idx_start,
positions,
hidden_states,
per_layer_inputs,
**kwargs,
)
return hidden_states, per_layer_inputs
@support_torch_compile(
enable_if=lambda vllm_config: vllm_config.cache_config.kv_sharing_fast_prefill
)
class Gemma4CrossDecoderLayers(nn.Module):
"""Cross-decoder layers (YOCO second half, KV-shared)."""
def __init__(
self,
*,
vllm_config: VllmConfig,
prefix: str = "",
decoder_layers: list[Gemma4DecoderLayer],
layer_idx_start: int,
):
super().__init__()
self.decoder_layers = decoder_layers
self.layer_idx_start = layer_idx_start
def forward(
self,
positions: torch.Tensor,
hidden_states: torch.Tensor,
per_layer_inputs: torch.Tensor | None = None,
**kwargs,
) -> torch.Tensor:
return _run_decoder_layers(
self.decoder_layers,
self.layer_idx_start,
positions,
hidden_states,
per_layer_inputs,
**kwargs,
)
@support_torch_compile(
enable_if=lambda vllm_config: not vllm_config.cache_config.kv_sharing_fast_prefill
)
class Gemma4Model(nn.Module, EagleModelMixin):
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
super().__init__()
config = _get_text_config(vllm_config.model_config.hf_config)
@@ -740,6 +948,75 @@ class Gemma4Model(nn.Module):
torch.tensor(config.hidden_size**0.5),
persistent=False,
)
# --- You Only Cache Once (YOCO) split for fast prefill ---
first_kv_shared_layer_idx = config.num_hidden_layers - getattr(
config, "num_kv_shared_layers", 0
)
from vllm.compilation.backends import set_model_tag
# Layers 0..(K-1) are self-decoder layers in YOCO
with set_model_tag("self_decoder"):
self.self_decoder = Gemma4SelfDecoderLayers(
vllm_config=vllm_config,
prefix=f"{prefix}.self_decoder",
decoder_layers=self.layers[:first_kv_shared_layer_idx],
layer_idx_start=0,
embed_tokens=self.embed_tokens,
normalizer=self.normalizer,
embed_tokens_per_layer=getattr(self, "embed_tokens_per_layer", None),
embed_scale_per_layer=getattr(self, "embed_scale_per_layer", None),
per_layer_model_projection=getattr(
self, "per_layer_model_projection", None
),
per_layer_projection_norm=getattr(
self, "per_layer_projection_norm", None
),
per_layer_input_scale=getattr(self, "per_layer_input_scale", None),
per_layer_projection_scale=getattr(
self, "per_layer_projection_scale", None
),
)
# Layers K..(N-1) are cross-decoder layers in YOCO
with set_model_tag("cross_decoder"):
self.cross_decoder = Gemma4CrossDecoderLayers(
vllm_config=vllm_config,
prefix=f"{prefix}.cross_decoder",
decoder_layers=self.layers[first_kv_shared_layer_idx:],
layer_idx_start=first_kv_shared_layer_idx,
)
self.fast_prefill_enabled = cache_config.kv_sharing_fast_prefill
if self.fast_prefill_enabled:
# Allocate static buffers for CUDAGraph
max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens
device = next(self.parameters()).device
self.positions = torch.zeros(
max_num_tokens, dtype=torch.int64, device=device
)
self.hidden_states = torch.zeros(
(max_num_tokens, config.hidden_size),
dtype=self.embed_tokens.weight.dtype,
device=device,
)
if (
self.hidden_size_per_layer_input
and self.hidden_size_per_layer_input > 0
):
self.per_layer_inputs = torch.zeros(
(
max_num_tokens,
config.num_hidden_layers,
self.hidden_size_per_layer_input,
),
dtype=self.embed_tokens.weight.dtype,
device=device,
)
else:
self.per_layer_inputs = None
# Custom factory that includes per_layer_inputs for PLE-enabled PP.
# per_layer_inputs has shape (batch, num_layers, per_layer_dim),
# which differs from the standard (batch, hidden_size) shape,
@@ -776,47 +1053,22 @@ class Gemma4Model(nn.Module):
self.make_empty_intermediate_tensors = _make_empty_intermediate_tensors
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
return self.embed_tokens(input_ids) * self.normalizer
return self.self_decoder.embed_input_ids(input_ids)
def get_per_layer_inputs(self, input_ids: torch.Tensor) -> torch.Tensor:
def get_per_layer_inputs(self, input_ids: torch.Tensor) -> torch.Tensor | None:
"""Get per-layer embeddings from embed_tokens_per_layer.
Returns:
Per-layer embeddings (num_tokens, num_layers,
hidden_size_per_layer_input)
"""
if self.embed_tokens_per_layer is None:
return None
# Handle out-of-vocab tokens for PLE (vocab_size_per_layer_input may
# be smaller than the main vocab_size).
per_layer_inputs_mask = torch.logical_and(
input_ids >= 0,
input_ids < self.vocab_size_per_layer_input,
)
per_layer_inputs_tokens = torch.where(
per_layer_inputs_mask, input_ids, torch.zeros_like(input_ids)
)
# Get packed per-layer embeddings: (num_tokens, total_ple_dim)
per_layer_embeds = self.embed_tokens_per_layer(per_layer_inputs_tokens)
# Apply embed_scale (sqrt of per-layer hidden dim)
per_layer_embeds = per_layer_embeds * self.embed_scale_per_layer
# Reshape to (num_tokens, num_layers, hidden_size_per_layer_input)
per_layer_embeds = per_layer_embeds.reshape(
*input_ids.shape,
self.config.num_hidden_layers,
self.hidden_size_per_layer_input,
)
return per_layer_embeds
return self.self_decoder.get_per_layer_inputs(input_ids)
def project_per_layer_inputs(
self,
inputs_embeds: torch.Tensor,
per_layer_inputs: torch.Tensor | None,
) -> torch.Tensor:
) -> torch.Tensor | None:
"""Project inputs_embeds and combine with per_layer_inputs.
Steps:
@@ -826,29 +1078,94 @@ class Gemma4Model(nn.Module):
4. Normalize with per_layer_projection_norm
5. Combine: (projection + per_layer_inputs) * 1/sqrt(2)
"""
if self.per_layer_model_projection is None:
return None
# Project from hidden_size to total_ple_dim
# Scaled projection: output = linear(input, weight) * scale
per_layer_projection = self.per_layer_model_projection(inputs_embeds)
per_layer_projection = per_layer_projection * self.per_layer_projection_scale
# Reshape to (num_tokens, num_layers, hidden_size_per_layer_input)
per_layer_projection = per_layer_projection.reshape(
*inputs_embeds.shape[:-1],
self.config.num_hidden_layers,
self.hidden_size_per_layer_input,
return self.self_decoder.project_per_layer_inputs(
inputs_embeds, per_layer_inputs
)
# Normalize
per_layer_projection = self.per_layer_projection_norm(per_layer_projection)
def fast_prefill_forward(
self,
input_ids: torch.Tensor | None,
positions: torch.Tensor,
inputs_embeds: torch.Tensor | None = None,
per_layer_inputs: torch.Tensor | None = None,
**kwargs,
) -> torch.Tensor:
logits_indices_padded, num_logits_indices = None, None
attn_metadata = get_forward_context().attn_metadata
if per_layer_inputs is None:
return per_layer_projection
if attn_metadata is not None:
assert isinstance(attn_metadata, dict)
layer_attn_metadata = attn_metadata[
self.layers[-1].self_attn.attn.layer_name
]
if isinstance(layer_attn_metadata, KVSharingFastPrefillMetadata):
logits_indices_padded = layer_attn_metadata.logits_indices_padded
num_logits_indices = layer_attn_metadata.num_logits_indices
# Combine: (projection + per_layer_inputs) * scale
return (per_layer_projection + per_layer_inputs) * self.per_layer_input_scale
batch_size = positions.size(0)
self.positions[:batch_size].copy_(positions)
self_decoder_hidden_states, per_layer_inputs = self.self_decoder(
input_ids=input_ids,
positions=self.positions[:batch_size],
inputs_embeds=inputs_embeds,
per_layer_inputs=per_layer_inputs,
**kwargs,
)
if logits_indices_padded is None:
logits_indices_padded = torch.arange(
batch_size,
dtype=positions.dtype,
device=positions.device,
)
# NOTE: Keep .clone() until fix in
# https://github.com/vllm-project/vllm/pull/22282
hidden_states = self_decoder_hidden_states.clone()
num_padded = logits_indices_padded.size(0)
self.positions[:num_padded].copy_(positions[logits_indices_padded])
self.hidden_states[:num_padded].copy_(
self_decoder_hidden_states[logits_indices_padded]
)
if self.per_layer_inputs is not None and per_layer_inputs is not None:
self.per_layer_inputs[:num_padded].copy_(
per_layer_inputs[logits_indices_padded]
)
# Update batch_descriptor so the cross-decoder's piecewise
# CUDAGraphWrapper dispatches to the correct (reduced) batch size.
forward_context = get_forward_context()
orig_batch_desc = forward_context.batch_descriptor
if orig_batch_desc is not None:
forward_context.batch_descriptor = replace(
orig_batch_desc, num_tokens=num_padded
)
cross_per_layer = (
self.per_layer_inputs[:num_padded]
if self.per_layer_inputs is not None
else None
)
cross_hidden_states = self.cross_decoder(
self.positions[:num_padded],
self.hidden_states[:num_padded],
cross_per_layer,
**kwargs,
)
# Restore the original batch_descriptor
forward_context.batch_descriptor = orig_batch_desc
if num_logits_indices is not None:
assert num_logits_indices > 0
hidden_states[logits_indices_padded[:num_logits_indices]] = (
cross_hidden_states[:num_logits_indices]
)
else:
hidden_states = cross_hidden_states
return hidden_states
def forward(
self,
@@ -858,7 +1175,19 @@ class Gemma4Model(nn.Module):
inputs_embeds: torch.Tensor | None = None,
per_layer_inputs: torch.Tensor | None = None,
**kwargs,
) -> torch.Tensor | IntermediateTensors:
) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
if self.fast_prefill_enabled:
hidden_states = self.fast_prefill_forward(
input_ids,
positions,
inputs_embeds,
per_layer_inputs,
**kwargs,
)
hidden_states = self.norm(hidden_states)
return hidden_states
# Normal (non-fast-prefill) path with PP support
if get_pp_group().is_first_rank:
if inputs_embeds is not None:
hidden_states = inputs_embeds
@@ -882,6 +1211,7 @@ class Gemma4Model(nn.Module):
residual = intermediate_tensors["residual"]
per_layer_inputs = intermediate_tensors.get("per_layer_inputs")
aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual)
for layer_idx, layer in enumerate(
islice(self.layers, self.start_layer, self.end_layer)
):
@@ -900,6 +1230,9 @@ class Gemma4Model(nn.Module):
per_layer_input=layer_per_input,
**kwargs,
)
self._maybe_add_hidden_state(
aux_hidden_states, layer_idx + 1, hidden_states, residual
)
if not get_pp_group().is_last_rank:
return IntermediateTensors(
{
@@ -914,6 +1247,9 @@ class Gemma4Model(nn.Module):
hidden_states = self.norm(hidden_states)
else:
hidden_states, _ = self.norm(hidden_states, residual)
if len(aux_hidden_states) > 0:
return hidden_states, aux_hidden_states
return hidden_states
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
@@ -926,21 +1262,27 @@ class Gemma4Model(nn.Module):
("gate_up_proj", "up_proj", 1),
]
# MoE expert weight mapping: checkpoint 3D packed tensors are
# exploded in _weight_iterator to per-expert 2D weights like:
# MoE expert weight mapping: checkpoint can have either:
# 1. 3D packed tensors (exploded in _weight_iterator to per-expert 2D)
# 2. Already per-expert 2D weights (if quantized)
# Map to FusedMoE parameters:
# moe.experts.{id}.gate_proj → FusedMoE w1 (shard of w13)
# moe.experts.{id}.up_proj → FusedMoE w3 (shard of w13)
# moe.experts.{id}.down_proj → FusedMoE w2
# We build the mapping directly since Gemma4 uses bare param
# names (no .weight suffix) unlike standard MoE checkpoints.
#
# Use prefix matching to handle both weights and
# quantization scale parameters. The param_name is a prefix ending
# in underscore, and weight_name ends with a dot, so that:
# "experts.0.gate_proj.weight_scale" -> "experts.w13_weight_scale"
# "experts.0.gate_proj.weight" -> "experts.w13_weight"
num_experts = getattr(self.config, "num_experts", None) or 0
expert_params_mapping = [
# (param_name, weight_name, expert_id, shard_id)
(
"experts.w13_weight"
"experts.w13_"
if proj_name in ["gate_proj", "up_proj"]
else "experts.w2_weight",
f"experts.{expert_id}.{proj_name}",
else "experts.w2_",
f"experts.{expert_id}.{proj_name}.",
expert_id,
shard_id,
)
@@ -1000,9 +1342,21 @@ class Gemma4Model(nn.Module):
expert_id,
shard_id,
) in expert_params_mapping:
if weight_name not in name:
# Match both:
# - Bare weights: "experts.0.down_proj" (from 3D explosion)
# - With suffix: "experts.0.down_proj.weight_scale" (2D quantized)
# weight_name has trailing dot, so check with and without it
weight_name_base = weight_name.rstrip(".")
if weight_name in name:
# Has suffix (e.g., .weight_scale)
moe_name = name.replace(weight_name, param_name)
elif name.endswith(weight_name_base):
# Bare weight (no suffix)
moe_name = name.replace(
weight_name_base, param_name.rstrip("_") + "_weight"
)
else:
continue
moe_name = name.replace(weight_name, param_name)
if moe_name not in params_dict:
continue
if is_pp_missing_parameter(moe_name, self):
@@ -1012,15 +1366,12 @@ class Gemma4Model(nn.Module):
# orientation for FusedMoE after _weight_iterator:
# gate/up: [I, H] → w1/w3 expects [I, H]
# down: [H, I] → w2 expects [H, I]
assert loaded_weight.dim() == 2, (
f"Expected 2D expert weight for {weight_name}, "
f"got shape {loaded_weight.shape}"
)
# Scales and other quantization params may be 1D or scalar.
weight_loader = param.weight_loader
weight_loader(
param,
loaded_weight,
weight_name + ".weight",
moe_name, # Pass mapped name (handles both weights and scales)
shard_id=shard_id,
expert_id=expert_id,
)
@@ -1044,7 +1395,25 @@ class Gemma4Model(nn.Module):
return loaded_params
class Gemma4ForCausalLM(nn.Module, SupportsLoRA, SupportsPP, MixtureOfExperts):
class Gemma4ForCausalLM(
nn.Module, SupportsLoRA, SupportsPP, MixtureOfExperts, SupportsEagle3
):
hf_to_vllm_mapper = WeightsMapper(
orig_to_new_prefix={
# Gemma4ForConditionalGeneration already loads the text stack
# from `model.language_model.*`. We reuse that same checkpoint
# and adapter naming for the text-only Gemma4ForCausalLM path,
# so LoRA keys from the conditional wrapper map onto `model.*`.
"model.language_model.": "model.",
},
orig_to_new_substr={
# Gemma4ForConditionalGeneration names MoE adapter targets under
# `...moe.experts.*`, while the text-only model exposes them
# under `...moe.*`.
".moe.experts.gate_up_proj": ".moe.gate_up_proj",
".moe.experts.down_proj": ".moe.down_proj",
},
)
# Note: qkv_proj packing applies to non-k_eq_v layers (sliding
# attention and full attention without k_eq_v). k_eq_v layers use
# separate q_proj + k_proj without packing.
@@ -1126,7 +1495,7 @@ class Gemma4ForCausalLM(nn.Module, SupportsLoRA, SupportsPP, MixtureOfExperts):
intermediate_tensors: IntermediateTensors | None = None,
inputs_embeds: torch.Tensor | None = None,
**kwargs,
) -> torch.Tensor | IntermediateTensors:
) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
hidden_states = self.model(
input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs
)
@@ -1177,6 +1546,11 @@ class Gemma4ForCausalLM(nn.Module, SupportsLoRA, SupportsPP, MixtureOfExperts):
".moe.down_proj",
)
# Remap individual 2D expert weights:
# .experts.{id}.{proj} → .moe.experts.{id}.{proj}
# (This handles per-expert 2D quantized weights)
name = re.sub(r"\.experts\.(\d+)\.", r".moe.experts.\1.", name)
# MoE expert weights: checkpoint stores as 3D packed
# tensors. Explode into per-expert 2D weights for
# FusedMoE weight_loader.
+55 -19
View File
@@ -65,7 +65,12 @@ from vllm.multimodal.processing.processor import (
from vllm.sequence import IntermediateTensors
from vllm.utils.tensor_schema import TensorSchema, TensorShape
from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
from .interfaces import (
MultiModalEmbeddings,
SupportsEagle3,
SupportsMultiModal,
SupportsPP,
)
from .utils import (
AutoWeightsLoader,
WeightsMapper,
@@ -121,8 +126,12 @@ class Gemma4AudioInputs(TensorSchema):
"""
type: Literal["audio"] = "audio"
input_features_padded: Annotated[torch.Tensor, TensorShape("bn", "s", "f")]
input_features_mask: Annotated[torch.Tensor, TensorShape("bn", "s")]
input_features_padded: Annotated[
torch.Tensor, TensorShape("bn", "s", "f", dynamic_dims={"s"})
]
input_features_mask: Annotated[
torch.Tensor, TensorShape("bn", "s", dynamic_dims={"s"})
]
Gemma4ImageInputs = Gemma4ImagePixelInputs
@@ -163,10 +172,15 @@ class Gemma4ProcessingInfo(BaseProcessingInfo):
Setting ``add_special_tokens=False`` here prevents the duplicate and
ensures both ``llm.generate()`` and the chat/completions API behave
correctly.
correctly for IT models. For PT models (without chat template), we
keep the default (True) to ensure BOS is added for raw prompts.
"""
tokenizer = self.ctx.get_tokenizer()
has_chat_template = getattr(tokenizer, "chat_template", None) is not None
params = super().get_default_tok_params()
params = params.with_kwargs(add_special_tokens=False)
if has_chat_template:
params = params.with_kwargs(add_special_tokens=False)
return params
def get_hf_processor(self, **kwargs: object) -> Gemma4Processor:
@@ -503,6 +517,8 @@ class Gemma4MultiModalProcessor(BaseMultiModalProcessor[Gemma4ProcessingInfo]):
video_timestamps_per_video: list[list[float]] = []
video_frame_counts: list[int] = []
video_replacements: list[str] = []
for item in videos:
video_array, metadata = item
@@ -555,10 +571,7 @@ class Gemma4MultiModalProcessor(BaseMultiModalProcessor[Gemma4ProcessingInfo]):
video_timestamps_per_video.append(timestamps)
video_frame_counts.append(len(frames))
# Build expanded replacement text and replace the
# <|video|> placeholder in the prompt.
# Use split(token, 1) to avoid collision — the
# replacement text itself contains <|video|> tokens.
# Build expanded replacement text for this video.
ts_strs = [f"{int(s // 60):02d}:{int(s % 60):02d}" for s in timestamps]
replacement = " ".join(
f"{t} {processor.boi_token}"
@@ -566,9 +579,23 @@ class Gemma4MultiModalProcessor(BaseMultiModalProcessor[Gemma4ProcessingInfo]):
f"{processor.eoi_token}"
for t, n in zip(ts_strs, num_soft_per_frame)
)
parts = prompt.split(processor.video_token, 1)
if len(parts) == 2:
prompt = parts[0] + replacement + parts[1]
video_replacements.append(replacement)
# Replace all <|video|> placeholders at once. We split on
# video_token to get N+1 parts, then interleave with the
# N replacement strings. This avoids the iterative
# split-replace bug where replacement text (which itself
# contains <|video|> tokens) collides with later splits.
vt = processor.video_token
parts = prompt.split(vt, len(video_replacements))
# NOTE: len(parts) <= len(video_replacements) + 1
parts_with_repl: list[str] = []
for part, repl in zip(parts, video_replacements):
parts_with_repl.extend([part, repl])
parts_with_repl.extend(parts[len(video_replacements) :])
prompt = "".join(parts_with_repl)
video_outputs = {
"pixel_values_videos": torch.cat(all_video_pixel_values, dim=0),
@@ -631,19 +658,23 @@ class Gemma4MultiModalProcessor(BaseMultiModalProcessor[Gemma4ProcessingInfo]):
)
if "input_features" in processed_outputs:
# Keep padded features for batched audio tower execution.
processed_outputs["input_features_padded"] = processed_outputs[
"input_features"
]
# Unpad per-item so each item's cache entry is self-contained.
# Unpad per-item so each item's cache entry is
# self-contained. The batched() field config in
# _get_mm_fields_config will re-pad all fields to the
# batch's max length at batch time, ensuring consistent
# padding regardless of cache history.
masks = processed_outputs["input_features_mask"]
unpadded_features = [
f[mask]
for f, mask in zip(
processed_outputs["input_features"],
processed_outputs["input_features_mask"],
masks,
)
]
unpadded_masks = [mask[mask] for mask in masks]
processed_outputs["input_features"] = unpadded_features
processed_outputs["input_features_padded"] = unpadded_features
processed_outputs["input_features_mask"] = unpadded_masks
# Merge video outputs into the final result
combined_outputs = dict(processed_outputs, **video_outputs)
@@ -848,7 +879,12 @@ class Gemma4MultimodalEmbedder(nn.Module):
info=Gemma4ProcessingInfo,
dummy_inputs=Gemma4DummyInputsBuilder,
)
class Gemma4ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):
class Gemma4ForConditionalGeneration(
nn.Module,
SupportsMultiModal,
SupportsPP,
SupportsEagle3,
):
packed_modules_mapping = {
"qkv_proj": [
"q_proj",
+24 -3
View File
@@ -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,))
+1 -3
View File
@@ -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)
+1 -1
View File
@@ -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,
)
@@ -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
): ...
+107 -18
View File
@@ -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:]]
@@ -20,7 +20,9 @@ from collections.abc import Mapping
from typing import TYPE_CHECKING
import torch
from transformers import AutoModel
from vllm.compilation.decorators import should_torch_compile_mm_encoder
from vllm.config.utils import getattr_iter
from vllm.inputs import MultiModalDataDict, MultiModalInput, mm_input
from vllm.logger import init_logger
@@ -46,19 +48,11 @@ from vllm.platforms import current_platform
from vllm.sequence import IntermediateTensors
if TYPE_CHECKING:
from transformers import BatchFeature
from transformers import BatchFeature, PreTrainedModel
from vllm.config import VllmConfig
from vllm.config.multimodal import BaseDummyOptions
DYNAMIC_ARG_DIMS = {
"input_ids": 0,
# set `positions` to last dim to support Qwen-mrope
"positions": -1,
"intermediate_tensors": 0,
"inputs_embeds": 0,
}
logger = init_logger(__name__)
@@ -274,6 +268,66 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE):
# Skip SupportsMRoPE.__init__ and call the next class in MRO
super(SupportsMRoPE, self).__init__(vllm_config=vllm_config, prefix=prefix)
def _get_encoder_cls(
self, modality: str = "image", **kwargs: dict
) -> type["PreTrainedModel"]:
"""
Get the encoder class from the model.
Args:
kwargs: The kwargs to create the model.
Returns:
The encoder class.
"""
with torch.device("meta"):
model: PreTrainedModel = AutoModel.from_config(**kwargs)
encoder_cls = type(model.get_encoder(modality=modality))
logger.debug("Identified encoder class as: %s", encoder_cls)
if type(model) is encoder_cls:
raise ValueError(
"Unable to infer vision encoder class from the model. "
"You must either: update the model so that "
"https://huggingface.co/docs/transformers/en/main_classes/model#transformers.PreTrainedModel.get_encoder"
" can detect the vision encoder correctly, or remove "
"'compile_mm_encoder'."
)
del model
return encoder_cls
def _decorate_for_torch_compile(self, **kwargs: dict):
"""
Decorate the model's decoder and encoder classes to indicate to vLLM that they
support torch compile if `can_enable_torch_compile` and
`should_torch_compile_mm_encoder` are True respectively.
Args:
kwargs: The kwargs to create the model, which are needed to get the decoder
and encoder classes.
"""
super()._decorate_for_torch_compile(**kwargs)
# Decorate the vision encoder model class to support torch compile if needed
if self.compilation_config.compile_mm_encoder:
self.check_version("5.0.0", "multimodal encoder compilation support")
logger.warning_once(
"Multimodal encoder compilation with the Transformers modeling backend "
"is an experimental feature. It relies on:\n"
"- The vision encoder being torch compilable.\n"
"- All vision encoder tensor inputs must be type hinted as either "
"`torch.Tensor` or `torch.FloatTensor`.\n"
"- The 0-th dimension of all tensor inputs to the vision encoder being "
"the dynamic dimension (i.e., sequence length or number of patches).\n"
"Please report any issues you encounter to help us improve it."
)
self._decorate_cls_for_torch_compile(
cls=self._get_encoder_cls(**kwargs),
# TODO: properly infer dynamic_arg_dims based on the encoder's forward
# method signature. Currently we assume dim 0 for all tensor inputs.
dynamic_arg_dims=None,
enable_if=should_torch_compile_mm_encoder,
is_encoder=True,
)
def forward(
self,
input_ids: torch.Tensor | None,
@@ -285,6 +339,10 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE):
# Gemma3 and PaliGemma needs `token_type_ids` to work correctly
# Other models will not have `token_type_ids` in kwargs
kwargs = {k: v for k, v in kwargs.items() if k == "token_type_ids"}
# Positions shape handling for MRoPE models
if self.model_config.uses_mrope:
# [3, seq_len] -> [3, 1, seq_len]
positions = positions[:, None]
model_output = super().forward(
input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs
)
+9
View File
@@ -470,6 +470,15 @@ class DelegatingParser(Parser):
# No tool calls
return [], content
def adjust_request(
self, request: ChatCompletionRequest | ResponsesRequest
) -> ChatCompletionRequest | ResponsesRequest:
if self._reasoning_parser is not None:
request = self._reasoning_parser.adjust_request(request)
if self._tool_parser is not None:
request = self._tool_parser.adjust_request(request)
return request
def extract_reasoning_streaming(
self,
previous_text: str,
+8 -2
View File
@@ -6,7 +6,7 @@ import os
from abc import abstractmethod
from collections.abc import Callable, Iterable, Sequence
from functools import cached_property
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, cast
from vllm.entrypoints.mcp.tool_server import ToolServer
from vllm.logger import init_logger
@@ -150,6 +150,12 @@ class ReasoningParser:
previously been parsed and extracted (see constructor)
"""
def adjust_request(
self, request: "ChatCompletionRequest | ResponsesRequest"
) -> "ChatCompletionRequest | ResponsesRequest":
"""Adjust request parameters; override in subclasses as needed."""
return request
def prepare_structured_tag(
self,
original_tag: str | None,
@@ -298,7 +304,7 @@ class ReasoningParserManager:
if isinstance(name, str):
names = [name]
elif is_list_of(name, str):
names = name
names = cast(list[str], name)
else:
names = [class_name]
+35 -3
View File
@@ -52,6 +52,16 @@ class Gemma4ReasoningParser(BaseThinkingReasoningParser):
# skip_special_tokens=True).
self._reasoning_text: str = ""
self._prefix_stripped: bool = False
self.new_turn_token_id = self.vocab["<|turn>"]
self.tool_call_token_id = self.vocab["<|tool_call>"]
self.tool_response_token_id = self.vocab["<|tool_response>"]
def adjust_request(
self, request: "ChatCompletionRequest | ResponsesRequest"
) -> "ChatCompletionRequest | ResponsesRequest":
"""Disable special-token stripping to preserve boundary tokens."""
request.skip_special_tokens = False
return request
@property
def start_token(self) -> str:
@@ -63,6 +73,29 @@ class Gemma4ReasoningParser(BaseThinkingReasoningParser):
"""The token that ends reasoning content."""
return "<channel|>"
def is_reasoning_end(self, input_ids: Sequence[int]) -> bool:
start_token_id = self.start_token_id
end_token_id = self.end_token_id
new_turn_token_id = self.new_turn_token_id
tool_call_token_id = self.tool_call_token_id
tool_response_token_id = self.tool_response_token_id
# Search from the end of input_ids to find the last match.
for i in range(len(input_ids) - 1, -1, -1):
if input_ids[i] == start_token_id:
return False
if input_ids[i] == tool_call_token_id:
# We're generating a tool call, so reasoning must be ended.
return True
if input_ids[i] in (new_turn_token_id, tool_response_token_id):
# We found a new turn or tool response token so don't consider
# reasoning ended yet, since the model starts new reasoning
# after these tokens.
return False
if input_ids[i] == end_token_id:
return True
return False
# ------------------------------------------------------------------
# Non-streaming path
# ------------------------------------------------------------------
@@ -159,11 +192,10 @@ class Gemma4ReasoningParser(BaseThinkingReasoningParser):
result.reasoning = stripped
return result
else:
# This entire delta was prefix — suppress it.
# Don't set _prefix_stripped yet; there may be more
# prefix chars to consume in the next delta.
if len(self._reasoning_text) >= prefix_len:
self._prefix_stripped = True
result.reasoning = ""
return result
return None
# Case 2: Accumulated text is a strict prefix of
+34 -1
View File
@@ -1,5 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import contextlib
from dataclasses import dataclass, field
from functools import lru_cache
from pathlib import Path
@@ -10,6 +11,7 @@ from typing_extensions import TypeVar, assert_never
import vllm.envs as envs
from vllm.logger import init_logger
from vllm.transformers_utils.config import get_config
from vllm.transformers_utils.gguf_utils import (
check_gguf_file,
get_gguf_file_path_from_hf,
@@ -31,6 +33,13 @@ if TYPE_CHECKING:
logger = init_logger(__name__)
# Model types whose hub tokenizer_class is incorrect and should be overridden with
# TokenizersBackend (the generic fast tokenizer). Adding a model type here is always a
# temporary workaround and better long term solutions are:
# - Add model type to MODELS_WITH_INCORRECT_HUB_TOKENIZER_CLASS in transformers (better)
# - Fix tokenizer_class on the hub for the affected models (best)
_MODEL_TYPES_WITH_INCORRECT_TOKENIZER_CLASS: set[str] = {"step3_vl"}
_VLLM_TOKENIZERS = {
"deepseek_v32": ("deepseek_v32", "DeepseekV32Tokenizer"),
"grok2": ("grok2", "Grok2Tokenizer"),
@@ -202,7 +211,31 @@ def get_tokenizer(
**kwargs,
)
if tokenizer_cls == TokenizerLike:
# Ensure that, if the config were to come from vllm.transformers_utils.config, it is
# registered with AutoConfig before the tokenizer is loaded. This is necessary since
# tokenizer_cls_.from_pretrained will call AutoConfig.from_pretrained internally.
# This may fail for paths that don't have a model config (e.g. LoRA adapters),
# which is fine — those don't need custom config registration.
config = None
with contextlib.suppress(ValueError, OSError):
config = get_config(
tokenizer_name,
trust_remote_code=trust_remote_code,
revision=revision,
)
# Some models have an incorrect tokenizer_class on the hub.
# For these model types, bypass AutoTokenizer and use TokenizersBackend directly.
model_type = getattr(config, "model_type", None) if config else None
if model_type in _MODEL_TYPES_WITH_INCORRECT_TOKENIZER_CLASS:
from transformers.tokenization_utils_tokenizers import TokenizersBackend
logger.debug(
"Overriding tokenizer_class to TokenizersBackend for model_type=%r",
model_type,
)
tokenizer_cls_ = TokenizersBackend
elif tokenizer_cls == TokenizerLike:
tokenizer_cls_ = TokenizerRegistry.load_tokenizer_cls(tokenizer_mode)
else:
tokenizer_cls_ = tokenizer_cls
+49 -15
View File
@@ -66,6 +66,10 @@ def _parse_gemma4_value(value_str: str) -> object:
if value_str == "false":
return False
# Null
if value_str.lower() in ("null", "none", "nil"):
return None
# Number (int or float)
try:
if "." in value_str:
@@ -78,7 +82,7 @@ def _parse_gemma4_value(value_str: str) -> object:
return value_str
def _parse_gemma4_args(args_str: str) -> dict:
def _parse_gemma4_args(args_str: str, *, partial: bool = False) -> dict:
"""Parse Gemma4's custom key:value format into a Python dict.
Format examples::
@@ -89,6 +93,12 @@ def _parse_gemma4_args(args_str: str) -> dict:
nested:{inner_key:<|"|>val<|"|>}
items:[<|"|>a<|"|>,<|"|>b<|"|>]
Args:
args_str: The raw Gemma4 argument string.
partial: When True (streaming), bare values at end of string are
omitted because they may be incomplete and type-unstable
(e.g. partial boolean parsed as bare string).
Returns a dict ready for ``json.dumps()``.
"""
if not args_str or not args_str.strip():
@@ -116,14 +126,16 @@ def _parse_gemma4_args(args_str: str) -> dict:
# Parse value
if i >= n:
result[key] = ""
if not partial:
result[key] = ""
break
# Skip whitespace after ':'
while i < n and args_str[i] in (" ", "\n", "\t"):
i += 1
if i >= n:
result[key] = ""
if not partial:
result[key] = ""
break
# String value: <|"|>...<|"|>
@@ -155,7 +167,12 @@ def _parse_gemma4_args(args_str: str) -> dict:
elif args_str[i] == "}":
depth -= 1
i += 1
result[key] = _parse_gemma4_args(args_str[obj_start : i - 1])
if depth > 0:
# Incomplete nested object — use i (not i-1) to avoid
# dropping the last char, and recurse as partial.
result[key] = _parse_gemma4_args(args_str[obj_start:i], partial=True)
else:
result[key] = _parse_gemma4_args(args_str[obj_start : i - 1])
# Array: [...]
elif args_str[i] == "[":
@@ -173,20 +190,26 @@ def _parse_gemma4_args(args_str: str) -> dict:
elif args_str[i] == "]":
depth -= 1
i += 1
arr_content = args_str[arr_start : i - 1]
result[key] = _parse_gemma4_array(arr_content)
if depth > 0:
result[key] = _parse_gemma4_array(args_str[arr_start:i], partial=True)
else:
result[key] = _parse_gemma4_array(args_str[arr_start : i - 1])
# Bare value (number, boolean, etc.)
else:
val_start = i
while i < n and args_str[i] not in (",", "}", "]"):
i += 1
if partial and i >= n:
# Value may be incomplete (e.g. partial boolean) —
# withhold to avoid type instability during streaming.
break
result[key] = _parse_gemma4_value(args_str[val_start:i])
return result
def _parse_gemma4_array(arr_str: str) -> list:
def _parse_gemma4_array(arr_str: str, *, partial: bool = False) -> list:
"""Parse a Gemma4 array content string into a Python list."""
items: list = []
i = 0
@@ -224,7 +247,10 @@ def _parse_gemma4_array(arr_str: str) -> list:
elif arr_str[i] == "}":
depth -= 1
i += 1
items.append(_parse_gemma4_args(arr_str[obj_start : i - 1]))
if depth > 0:
items.append(_parse_gemma4_args(arr_str[obj_start:i], partial=True))
else:
items.append(_parse_gemma4_args(arr_str[obj_start : i - 1]))
# Nested array
elif arr_str[i] == "[":
@@ -237,13 +263,18 @@ def _parse_gemma4_array(arr_str: str) -> list:
elif arr_str[i] == "]":
depth -= 1
i += 1
items.append(_parse_gemma4_array(arr_str[sub_start : i - 1]))
if depth > 0:
items.append(_parse_gemma4_array(arr_str[sub_start:i], partial=True))
else:
items.append(_parse_gemma4_array(arr_str[sub_start : i - 1]))
# Bare value
else:
val_start = i
while i < n and arr_str[i] not in (",", "]"):
i += 1
if partial and i >= n:
break
items.append(_parse_gemma4_value(arr_str[val_start:i]))
return items
@@ -436,8 +467,10 @@ class Gemma4ToolParser(ToolParser):
) -> DeltaMessage | None:
# Buffer delta text to handle multi-token special sequences
delta_text = self._buffer_delta_text(delta_text)
# Reconstruct current_text after buffering to stay in sync
current_text = previous_text + delta_text
# Keep current_text from the upstream stream state. The buffered delta
# is only for emission, and must not be stitched back into the
# accumulated model text or normal content like "<div>" can be
# duplicated into "<<div>" when a tool call just ended.
# If no tool call token seen yet, emit as content
if self.tool_call_start_token not in current_text:
@@ -661,7 +694,7 @@ class Gemma4ToolParser(ToolParser):
DeltaMessage with the argument diff, or None if no new content.
"""
try:
current_args = _parse_gemma4_args(raw_args_str)
current_args = _parse_gemma4_args(raw_args_str, partial=True)
except Exception:
logger.debug(
"Could not parse partial Gemma4 args yet: %s",
@@ -675,10 +708,11 @@ class Gemma4ToolParser(ToolParser):
current_args_json = json.dumps(current_args, ensure_ascii=False)
# Withhold trailing closing characters that may shift as more
# tokens arrive. Strip trailing '}', '"', and ']' sequences
# to get the "safe prefix".
# tokens arrive. Strip trailing '}', '"', ']' and partial
# STRING_DELIM fragments ('<', '|', '\\', '>') to get the
# "safe prefix".
safe_json = current_args_json
while safe_json and safe_json[-1] in ("}", '"', "]"):
while safe_json and safe_json[-1] in ("}", '"', "]", "<", "|", "\\", ">"):
safe_json = safe_json[:-1]
prev_streamed = self.streamed_args_for_tool[self.current_tool_id]
+9 -3
View File
@@ -565,11 +565,17 @@ class SlidingWindowManager(SingleTypeKVCacheManager):
for computed in computed_blocks:
computed.pop()
if use_eagle and computed_blocks[0]:
assert kv_cache_spec.block_size == alignment_tokens, (
"aligned_length is not compatible with eagle now"
)
for computed in computed_blocks:
computed.pop()
# Re-align after eagle pop: the pop may break the alignment
# when block_size != alignment_tokens (hybrid models with
# different page sizes, e.g. Gemma4).
while (
block_size != alignment_tokens
and len(computed_blocks[0]) * block_size % alignment_tokens != 0
):
for computed in computed_blocks:
computed.pop()
return computed_blocks
def get_num_skipped_tokens(self, num_computed_tokens: int) -> int:
+1
View File
@@ -1282,6 +1282,7 @@ class SpecDecodeBaseProposer:
"Qwen2_5_VLForConditionalGeneration",
"Qwen3VLForConditionalGeneration",
"Qwen3VLMoeForConditionalGeneration",
"Gemma4ForConditionalGeneration",
"HunYuanVLForConditionalGeneration",
"GlmOcrForConditionalGeneration",
"Qwen3_5ForConditionalGeneration",