forked from Karylab-cklius/vllm
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
b1388b1fbf | ||
|
|
9d780002bd | ||
|
|
39602ebf8a | ||
|
|
7225a69c9c | ||
|
|
9a3a31fd2d | ||
|
|
7bd3f40dda | ||
|
|
c5460385f1 | ||
|
|
4bbb8faa1f | ||
|
|
459d9b38d3 | ||
|
|
b1568cf464 | ||
|
|
a4ac72ceba | ||
|
|
0ab0f70aa9 | ||
|
|
ee642f8753 | ||
|
|
6db56c0997 | ||
|
|
9a234c7adc | ||
|
|
f56ffafdce | ||
|
|
0329e8c9ac | ||
|
|
6f3dc4d0aa | ||
|
|
a6ac49a33b |
@@ -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}"
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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};
|
||||
}
|
||||
@@ -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
@@ -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
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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 -%}
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -9,6 +9,8 @@ pytest-shard
|
||||
# --- Core Tools & Bindings ---
|
||||
absl-py
|
||||
arctic-inference
|
||||
lm_eval[api]
|
||||
modelscope
|
||||
|
||||
# --- Audio Processing ---
|
||||
librosa
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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 = (
|
||||
|
||||
@@ -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
@@ -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
@@ -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
|
||||
|
||||
@@ -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(
|
||||
(
|
||||
|
||||
@@ -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),
|
||||
)
|
||||
|
||||
@@ -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)
|
||||
@@ -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():
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -805,6 +805,8 @@ class SpeculativeConfig:
|
||||
"deepseek_v3",
|
||||
"kimi_k2",
|
||||
"kimi_k25",
|
||||
"minimax_m2",
|
||||
"gemma4",
|
||||
]
|
||||
if (
|
||||
self.method in ("eagle3", "extract_hidden_states")
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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")):
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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,))
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
): ...
|
||||
|
||||
@@ -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
|
||||
)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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]
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -1282,6 +1282,7 @@ class SpecDecodeBaseProposer:
|
||||
"Qwen2_5_VLForConditionalGeneration",
|
||||
"Qwen3VLForConditionalGeneration",
|
||||
"Qwen3VLMoeForConditionalGeneration",
|
||||
"Gemma4ForConditionalGeneration",
|
||||
"HunYuanVLForConditionalGeneration",
|
||||
"GlmOcrForConditionalGeneration",
|
||||
"Qwen3_5ForConditionalGeneration",
|
||||
|
||||
Reference in New Issue
Block a user