diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index 62cc63833c5..56e528ffb37 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -295,6 +295,20 @@ steps: - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt - bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh +- label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs) + timeout_in_minutes: 30 + working_dir: "/vllm-workspace/tests" + num_devices: 2 + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py + - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py + - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py + - vllm/distributed/kv_transfer/kv_connector/v1/offloading/ + - tests/v1/kv_connector/nixl_integration/ + commands: + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt + - bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh + - label: Pipeline + Context Parallelism (4 GPUs) timeout_in_minutes: 60 working_dir: "/vllm-workspace/tests" diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 362c733c8d7..e7f73b496e1 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -20,7 +20,20 @@ steps: - tests/kernels/core - tests/kernels/test_concat_mla_q.py commands: - - pytest -v -s kernels/core kernels/test_concat_mla_q.py + - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.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 diff --git a/.github/workflows/pre-commit.yml b/.github/workflows/pre-commit.yml index bbc55e9532b..8ab8d3e7035 100644 --- a/.github/workflows/pre-commit.yml +++ b/.github/workflows/pre-commit.yml @@ -2,6 +2,7 @@ name: pre-commit on: pull_request: + types: [opened, synchronize, reopened, labeled] push: branches: [main] @@ -15,7 +16,11 @@ permissions: jobs: pre-run-check: - if: github.event_name == 'pull_request' + if: >- + github.event_name == 'pull_request' && + (github.event.action != 'labeled' || + github.event.label.name == 'ready' || + github.event.label.name == 'verified') runs-on: ubuntu-latest steps: - name: Check PR label and author merge count @@ -44,7 +49,12 @@ jobs: pre-commit: needs: pre-run-check - if: always() && (needs.pre-run-check.result == 'success' || needs.pre-run-check.result == 'skipped') + if: >- + always() && + (github.event.action != 'labeled' || + github.event.label.name == 'ready' || + github.event.label.name == 'verified') && + (needs.pre-run-check.result == 'success' || needs.pre-run-check.result == 'skipped') runs-on: ubuntu-latest steps: - uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1 diff --git a/CMakeLists.txt b/CMakeLists.txt index 346f13d1459..f24c12eff83 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -307,6 +307,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. diff --git a/csrc/minimax_reduce_rms_kernel.cu b/csrc/minimax_reduce_rms_kernel.cu new file mode 100644 index 00000000000..6245b02d6e9 --- /dev/null +++ b/csrc/minimax_reduce_rms_kernel.cu @@ -0,0 +1,879 @@ + +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include + +#include +#include +#include + +#include "cuda_compat.h" +#include "cuda_utils.h" +#include "core/registration.h" +#include "minimax_reduce_rms_kernel.h" + +#include + +#define FINAL_MASK 0xffffffff +#define MINIMAX_REDUCE_RMS_WARP_SIZE 32 + +namespace vllm { +namespace tensorrt_llm { + +template +struct LamportComm { + __device__ __forceinline__ LamportComm(void** workspace, int rank) { + counter_ptr = &reinterpret_cast(workspace[NRanks * 3])[0]; + flag_ptr = &reinterpret_cast(workspace[NRanks * 3])[2]; + clear_ptr = &reinterpret_cast(workspace[NRanks * 3 + 1])[0]; + flag_value = *flag_ptr; + auto comm_size = reinterpret_cast(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(workspace[2 * NRanks + r]) + + data_offset * comm_size; + } + clear_buf = reinterpret_cast(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(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(&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(&vec)[i] = 0x80000000; + } + return vec; +} + +template +__device__ __forceinline__ float rms_rsqrt(float& v, float eps) { + constexpr float kInvDim = 1.0F / static_cast(Dim); + v = rsqrtf((v * kInvDim) + eps); + return v; +} + +template +__device__ __forceinline__ float4 rms_rsqrt(float4& v, float eps) { + constexpr float kInvDim = 1.0F / static_cast(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 +__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 +__inline__ __device__ T blockReduceSumV2(T* val) { + static __shared__ T shared[NUM][33]; + int lane = threadIdx.x & 0x1f; + int wid = threadIdx.x >> 5; + + warpReduceSumV2(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(val); + return (T)0.0f; +} + +// for float4 version +template +__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 +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 + + access_id_in_token; + access_stride = token_stride * params.hidden_dim / kElemsPerAccess; + tot_access = params.size_q / kElemsPerAccess; + } + + 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 +__global__ void __launch_bounds__(1024) + minimax_reduce_rms_kernel_lamport(MiniMaxReduceRMSParams params) { + IndexHelper 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 comm(params.workspace, params.rank); + int clear_access = comm.clear_size / kElemsPerAccess; +#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]; + float sum_variance = 0.F; + *reinterpret_cast(vals) = + reinterpret_cast(params.allreduce_in)[idx]; +#pragma unroll + for (int i = 0; i < kElemsPerAccess; ++i) { + sum_variance += static_cast(vals[i]) * static_cast(vals[i]); + } + blockReduceSumV2(&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( + 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( + 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]; + *reinterpret_cast::vec_type*>(norm_weight) = + reinterpret_cast::vec_type*>( + params.rms_gamma)[access_id_in_token]; + +#pragma unroll + for (int i = 0; i < kElemsPerAccess; ++i) { + vals[i] = static_cast( + static_cast(vals[i]) * + rsqrtf( + (sum_variance / static_cast(params.hidden_dim) / NRanks) + + params.rms_eps) * + static_cast(norm_weight[i])); + } + + reinterpret_cast(params.rms_norm_out)[idx] = + *reinterpret_cast(vals); + } + for (int idx = access_id; idx < clear_access; idx += access_stride) { + reinterpret_cast(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 +__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; + constexpr int ThreadsPerRowK = RankKDim / kElemsPerAccess; + // 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; + int access_stride_k = (params.stride_k > 0 ? params.stride_k : RankKDim) / + kElemsPerAccess; + // 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; + int access_stride_k_out = + (params.stride_k_out > 0 ? params.stride_k_out : params.hidden_dim_k) / + kElemsPerAccess; + +#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 comm(params.workspace, params.rank); + + DType norm_weight[kElemsPerAccess]{}; +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.wait;"); +#endif + if (is_q) { + if (is_valid_q) { + *reinterpret_cast::vec_type*>( + norm_weight) = + reinterpret_cast::vec_type const*>( + params.rms_gamma)[access_id_in_token]; + } + } else { + if (is_valid_k) { + *reinterpret_cast::vec_type*>( + norm_weight) = + reinterpret_cast::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]{}; + 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(&vals[row][0]) = + reinterpret_cast(params.allreduce_in)[idx_r]; +#pragma unroll + for (int i = 0; i < kElemsPerAccess; ++i) { + float x = static_cast(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(&vals[row][0]) = + reinterpret_cast(params.allreduce_in_k)[idx_r]; +#pragma unroll + for (int i = 0; i < kElemsPerAccess; ++i) { + float x = static_cast(vals[row][i]); + warp_sum_variance[row] += x * x; + } + } + } + + local_warp_reduce_sum_array( + 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(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( + comm.data_bufs[tid])[(params.rank * tot_groups * 2) + (2 * g)] = + *reinterpret_cast(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( + 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( + reinterpret_cast(&var_all_ranks), kQActiveMask); + + // Thread 0 computes rsqrt with compile-time Dim and writes to smem + if (tid == 0) { + *reinterpret_cast(global_scale_q) = + rms_rsqrt(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(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( + comm.data_bufs[k_thread_idx])[(params.rank * tot_groups * 2) + + (2 * g + 1)] = + *reinterpret_cast(local_sum); + + bool done = false; + float4 var_all_ranks; + while (!done) { + done = true; + var_all_ranks = ld_global_volatile(&reinterpret_cast( + 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( + reinterpret_cast(&var_all_ranks), kKActiveMask); + + if (k_thread_idx == 0) { + *reinterpret_cast(global_scale_k) = + rms_rsqrt(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; ++i) { + vals[r][i] = static_cast(static_cast(vals[r][i]) * + warp_sum_variance[r] * + static_cast(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(params.rms_norm_out)[idx_out] = + *reinterpret_cast(&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; ++i) { + vals[r][i] = static_cast(static_cast(vals[r][i]) * + warp_sum_variance[r] * + static_cast(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(params.rms_norm_out_k)[idx_out] = + *reinterpret_cast(&vals[r][0]); + } + } + } // end group loop +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.launch_dependents;"); +#endif + + int clear_access = static_cast(comm.clear_size / kElemsPerAccess); + int clear_stride = group_stride * blockDim.x; + for (int idx = group_id * blockDim.x + threadIdx.x; idx < clear_access; + idx += clear_stride) { + reinterpret_cast(comm.clear_buf)[idx] = clear_vec; + } + + comm.update(static_cast(2) * tot_groups * kElemsPerAccess * + 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 +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 +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; + int block_size = threads_per_token; + + int max_blocks_per_sm = get_max_active_blocks( + minimax_reduce_rms_kernel_lamport, 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, params)); +} + +template +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 == 0); + if (params.stride_q > 0) { + TORCH_CHECK(params.stride_q % kElemsPerAccess == 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 == 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 == 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; + int access_per_row_k = params.hidden_dim_k / kElemsPerAccess; + + // 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; + + 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 +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( + params); + } else { + minimax_reduce_rms_kernel_launcher(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( + params); + } else { + minimax_reduce_rms_kernel_launcher(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(nranks); + allreduce_params.rank = static_cast(rank); + allreduce_params.dtype = input.scalar_type(); + allreduce_params.size_q = static_cast(input.numel()); + allreduce_params.hidden_dim = static_cast(input.size(-1)); + allreduce_params.stride_q = allreduce_params.hidden_dim; + allreduce_params.workspace = + reinterpret_cast(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(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 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(nranks); + params.rank = static_cast(rank); + params.dtype = qkv.scalar_type(); + params.size_q = static_cast(num_tokens * q_size); + params.hidden_dim = static_cast(q_size); + params.size_k = static_cast(num_tokens * kv_size); + params.hidden_dim_k = static_cast(kv_size); + params.stride_q = static_cast(qkv_dim); + params.stride_k = static_cast(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(workspace.mutable_data_ptr()); + + uint8_t* base = static_cast(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(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}; +} diff --git a/csrc/minimax_reduce_rms_kernel.h b/csrc/minimax_reduce_rms_kernel.h new file mode 100644 index 00000000000..e8c2d012247 --- /dev/null +++ b/csrc/minimax_reduce_rms_kernel.h @@ -0,0 +1,79 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include + +#include + +namespace vllm { +namespace tensorrt_llm { + +template +struct ElemsPerAccess; + +template <> +struct ElemsPerAccess { + static constexpr int value = 8; + using vec_type = float4; +}; + +template <> +struct ElemsPerAccess { + static constexpr int value = 8; + using vec_type = float4; +}; + +template <> +struct ElemsPerAccess { + static constexpr int value = 4; + using vec_type = float4; +}; + +template +static constexpr int kElemsPerAccess = ElemsPerAccess::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 diff --git a/csrc/ops.h b/csrc/ops.h index b5e9ff36a59..da066512b7c 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -308,4 +308,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 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 \ No newline at end of file diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index 78156cfced9..27966e6808e 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -496,6 +496,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 } diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base index 51ce663a1b2..1ab2d7229bc 100644 --- a/docker/Dockerfile.rocm_base +++ b/docker/Dockerfile.rocm_base @@ -9,7 +9,7 @@ ARG PYTORCH_AUDIO_BRANCH="v2.9.0" ARG PYTORCH_AUDIO_REPO="https://github.com/pytorch/audio.git" ARG FA_BRANCH="0e60e394" ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git" -ARG AITER_BRANCH="v0.1.12" +ARG AITER_BRANCH="v0.1.10.post3" ARG AITER_REPO="https://github.com/ROCm/aiter.git" ARG MORI_BRANCH="2d02c6a9" ARG MORI_REPO="https://github.com/ROCm/mori.git" diff --git a/docs/models/pooling_models/classify.md b/docs/models/pooling_models/classify.md index 69a6fe75d37..a565553e5e1 100644 --- a/docs/models/pooling_models/classify.md +++ b/docs/models/pooling_models/classify.md @@ -267,9 +267,34 @@ You can modify the `problem_type` via problem_type in the Hugging Face config. T Implement alignment with transformers [ForSequenceClassificationLoss](https://github.com/huggingface/transformers/blob/57bb6db6ee4cfaccc45b8d474dfad5a17811ca60/src/transformers/loss/loss_utils.py#L92). -### Logit bias +### Affine Score Calibration -You can modify the `logit_bias` (aka `sigmoid_normalize`) through the logit_bias parameter in `vllm.config.PoolerConfig`. +Affine Score Calibration, also known as [Platt Scaling](https://en.wikipedia.org/wiki/Platt_scaling) (Platt, 1999), is the most widely used method for calibrating classifier outputs into well-calibrated probabilities. + +The calibration follows the transformation: + +`activation(logit_scale * (logit - logit_bias))` + +| Parameter | Default | Description | +| --------- | ------- | ----------- | +| `logit_bias` | `None` | Bias subtracted from logits before activation | +| `logit_scale` | `None` | Scale factor applied to logits after bias subtraction | + +Note: `logit_bias` is **subtracted** from the logits (not added), consistent with the `sigmoid_normalize` convention where `sigmoid(x - bias)` centers the sigmoid around the bias value. + +The computation order is as follows: + +```python +logits -= logit_bias # subtract bias (center scores) +logits *= logit_scale # scale logits +logits = activation(logits) # e.g. sigmoid +``` + +Example configuration: + +```bash +--pooler-config '{"use_activation": true, "logit_bias": 4.5, "logit_scale": 1.0}' +``` ## Removed Features diff --git a/docs/models/pooling_models/token_embed.md b/docs/models/pooling_models/token_embed.md index 3396f4eac2d..b0e094267db 100644 --- a/docs/models/pooling_models/token_embed.md +++ b/docs/models/pooling_models/token_embed.md @@ -71,6 +71,14 @@ Models of any architecture can be converted into embedding models using `--conve If your model is not in the above list, we will try to automatically convert the model using [as_embedding_model][vllm.model_executor.models.adapters.as_embedding_model]. +### Special models + +| Architecture | Models | Example HF Models | [LoRA](../../features/lora.md) | [PP](../../serving/parallelism_scaling.md) | +| ------------ | ------ | ----------------- | -------------------- | ------------------------- | +| `JinaForRanking` | Qwen3-based | `jinaai/jina-reranker-v3` | | | + +jina-reranker-v3 is a listwise document reranker model with a novel `last but not late interaction` architecture. More information can be found at: [examples/pooling/token_embed/jina_reranker_v3_offline.py](../../../examples/pooling/token_embed/jina_reranker_v3_offline.py) + --8<-- [end:supported-token-embed-models] ## Offline Inference diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 68e30d708e4..ef1f5901ed0 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -666,7 +666,7 @@ Speech2Text models trained specifically for Automatic Speech Recognition. | `Gemma3nForConditionalGeneration` | Gemma3n | `google/gemma-3n-E2B-it`, `google/gemma-3n-E4B-it`, etc. | | | | `GlmAsrForConditionalGeneration` | GLM-ASR | `zai-org/GLM-ASR-Nano-2512` | ✅︎ | ✅︎ | | `GraniteSpeechForConditionalGeneration` | Granite Speech | `ibm-granite/granite-4.0-1b-speech`, `ibm-granite/granite-speech-3.3-2b`, etc. | ✅︎ | ✅︎ | -| `Qwen3ASRForConditionalGeneration` | Qwen3-ASR | `Qwen/Qwen3-ASR-1.7B`, etc. | | ✅︎ | +| `Qwen3ASRForConditionalGeneration` | Qwen3-ASR | `Qwen/Qwen3-ASR-1.7B`, etc. | ✅︎ | ✅︎ | | `Qwen3OmniMoeThinkerForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct`, etc. | | ✅︎ | | `VoxtralForConditionalGeneration` | Voxtral (Mistral format) | `mistralai/Voxtral-Mini-3B-2507`, `mistralai/Voxtral-Small-24B-2507`, etc. | ✅︎ | ✅︎ | | `WhisperForConditionalGeneration` | Whisper | `openai/whisper-small`, `openai/whisper-large-v3-turbo`, etc. | | | diff --git a/examples/pooling/token_embed/jina_reranker_v3_offline.py b/examples/pooling/token_embed/jina_reranker_v3_offline.py new file mode 100644 index 00000000000..c250eccc62a --- /dev/null +++ b/examples/pooling/token_embed/jina_reranker_v3_offline.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# ruff: noqa: E501 + +import torch.nn.functional as F + +from vllm import LLM + +query = "What are the health benefits of green tea?" +documents = [ + "Green tea contains antioxidants called catechins that may help reduce inflammation and protect cells from damage.", + "El precio del café ha aumentado un 20% este año debido a problemas en la cadena de suministro.", + "Studies show that drinking green tea regularly can improve brain function and boost metabolism.", + "Basketball is one of the most popular sports in the United States.", + "绿茶富含儿茶素等抗氧化剂，可以降低心脏病风险，还有助于控制体重。", + "Le thé vert est riche en antioxydants et peut améliorer la fonction cérébrale.", +] + + +def main(): + # Initialize model + llm = LLM( + model="jinaai/jina-reranker-v3", + runner="pooling", + ) + + # Generate scores. + outputs = llm.score(query, documents) + + # Print the outputs. + print("\nGenerated Outputs:\n" + "-" * 60) + for document, output in zip(documents, outputs): + score = output.outputs.score + print(f"Pair: {[query, document]!r} \nScore: {score}") + print("-" * 60) + + # Generate embeddings. + # The JinaForRanking model concatenates docs first, then query. + # Let's stay consistent with this novel design. + outputs = llm.encode(documents + [query], pooling_task="token_embed") + embeds = outputs[0].outputs.data.float() + + doc_embeds = embeds[:-1] + query_embeds = embeds[-1] + + scores = F.cosine_similarity(query_embeds, doc_embeds) + + # Print the outputs. + print("\nGenerated Outputs:\n" + "-" * 60) + for document, score in zip(documents, scores): + print(f"Pair: {[query, document]!r} \nScore: {score}") + print("-" * 60) + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index 2758c3e0ac1..660b94fac77 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -120,7 +120,8 @@ python = "./.venv" [tool.typos.files] # these files may be written in non english words extend-exclude = ["tests/models/fixtures/*", "tests/prompts/*", "tests/tokenizers_/*", - "benchmarks/sonnet.txt", "tests/lora/data/*", "examples/pooling/token_embed/*", "build/*", + "benchmarks/sonnet.txt", "tests/lora/data/*", "build/*", + "examples/pooling/token_embed/*", "tests/models/language/pooling/*", "vllm/third_party/*", "vllm/entrypoints/serve/instrumentator/static/*", "tests/entrypoints/openai/speech_to_text/test_transcription_validation.py", "docs/governance/process.md", "tests/v1/engine/test_fast_incdec_prefix_err.py", ".git/*"] ignore-hidden = false diff --git a/tests/compile/passes/test_fusion_attn.py b/tests/compile/passes/test_fusion_attn.py index 2bbf0bda626..40904f2db11 100644 --- a/tests/compile/passes/test_fusion_attn.py +++ b/tests/compile/passes/test_fusion_attn.py @@ -39,7 +39,7 @@ from vllm.platforms import current_platform from vllm.utils.flashinfer import has_flashinfer from vllm.v1.attention.backend import AttentionMetadata from vllm.v1.attention.backends.registry import AttentionBackendEnum -from vllm.v1.kv_cache_interface import AttentionSpec +from vllm.v1.kv_cache_interface import AttentionSpec, get_kv_quant_mode FP8_DTYPE = current_platform.fp8_dtype() FP4_DTYPE = torch.uint8 @@ -53,7 +53,6 @@ class AttentionQuantPatternModel(torch.nn.Module): num_qo_heads: int, num_kv_heads: int, head_size: int, - kv_cache_dtype: torch.dtype, device: torch.device, vllm_config: VllmConfig, block_size: int, @@ -63,7 +62,6 @@ class AttentionQuantPatternModel(torch.nn.Module): self.num_qo_heads = num_qo_heads self.num_kv_heads = num_kv_heads self.head_size = head_size - self.kv_cache_dtype = kv_cache_dtype self.device = device self.vllm_config = vllm_config self.dtype = vllm_config.model_config.dtype @@ -81,13 +79,14 @@ class AttentionQuantPatternModel(torch.nn.Module): self.block_size = block_size - # Initialize attn MetadataBuilder + # Initialize attn MetadataBuilder (match Attention.get_kv_cache_spec) self.builder = self.attn.attn_backend.get_builder_cls()( kv_cache_spec=AttentionSpec( block_size=self.block_size, num_kv_heads=self.num_kv_heads, head_size=self.head_size, - dtype=self.kv_cache_dtype, + dtype=self.attn.kv_cache_torch_dtype, + kv_quant_mode=get_kv_quant_mode(self.attn.kv_cache_dtype), ), layer_names=[self.attn.layer_name], vllm_config=self.vllm_config, @@ -126,7 +125,7 @@ class AttentionQuantPatternModel(torch.nn.Module): # Create dummy KV cache raw_tensor = torch.zeros( 2 * num_blocks * self.block_size * self.num_kv_heads * self.head_size, - dtype=self.kv_cache_dtype, + dtype=self.attn.kv_cache_torch_dtype, device=self.device, ) raw_tensor = raw_tensor.view(kv_cache_shape) @@ -348,7 +347,6 @@ def test_attention_quant_pattern( num_qo_heads=num_qo_heads, num_kv_heads=num_kv_heads, head_size=head_size, - kv_cache_dtype=FP8_DTYPE, device=device, vllm_config=vllm_config_unfused, block_size=block_size, @@ -376,7 +374,6 @@ def test_attention_quant_pattern( num_qo_heads=num_qo_heads, num_kv_heads=num_kv_heads, head_size=head_size, - kv_cache_dtype=FP8_DTYPE, device=device, vllm_config=vllm_config, w=model_unfused.w, diff --git a/tests/kernels/core/test_minimax_reduce_rms.py b/tests/kernels/core/test_minimax_reduce_rms.py new file mode 100644 index 00000000000..d17a448bd97 --- /dev/null +++ b/tests/kernels/core/test_minimax_reduce_rms.py @@ -0,0 +1,152 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for MiniMax QK RMS-norm: NCCL reference vs Lamport fused kernel.""" + +import pytest +import torch +import torch.nn as nn +from torch.multiprocessing import spawn + +from tests.kernels.utils import opcheck +from tests.utils import ensure_current_vllm_config, init_test_distributed_environment +from vllm.distributed import cleanup_dist_env_and_memory +from vllm.model_executor.layers.mamba.linear_attn import MiniMaxText01RMSNormTP +from vllm.platforms import current_platform +from vllm.utils.network_utils import get_open_port +from vllm.utils.torch_utils import set_random_seed + + +@ensure_current_vllm_config() +def _worker_forward_qk( + local_rank, + world_size, + port, + num_tokens, + hidden_q_full, + hidden_k_full, + dtype, + seed, + eps, +): + """Per-rank worker: compare NCCL allreduce path vs Lamport fused kernel.""" + + if not hasattr(torch.ops._C, "minimax_allreduce_rms_qk"): + cleanup_dist_env_and_memory() + return + device = torch.device(f"cuda:{local_rank}") + torch.accelerator.set_device_index(device) + init_test_distributed_environment( + world_size, 1, local_rank, port, local_rank=local_rank + ) + + hq = hidden_q_full // world_size + hk = hidden_k_full // world_size + + q_norm = MiniMaxText01RMSNormTP(hidden_q_full, eps=eps).cuda() + k_norm = MiniMaxText01RMSNormTP(hidden_k_full, eps=eps).cuda() + + set_random_seed(seed) + qw = torch.randn(hidden_q_full, dtype=dtype, device="cuda") + kw = torch.randn(hidden_k_full, dtype=dtype, device="cuda") + q_norm.weight = nn.Parameter(qw[local_rank * hq : (local_rank + 1) * hq]) + k_norm.weight = nn.Parameter(kw[local_rank * hk : (local_rank + 1) * hk]) + + torch.manual_seed(seed + 1000 + local_rank) + qkv = torch.randn(num_tokens, hq + hk + hk, dtype=dtype, device="cuda") + + q_ref, k_ref, v_ref = qkv.clone().split([hq, hk, hk], dim=-1) + ref_q, ref_k = MiniMaxText01RMSNormTP.forward_qk(q_norm, k_norm, q_ref, k_ref) + + # Set up Lamport workspace. + from vllm.distributed.parallel_state import get_tp_group + from vllm.model_executor.layers.mamba.lamport_workspace import ( + get_allreduce_workspace, + ) + + workspace = get_allreduce_workspace( + rank=local_rank, + world_size=world_size, + max_tokens=num_tokens, + process_group=get_tp_group().cpu_group, + ) + + opcheck( + torch.ops._C.minimax_allreduce_rms_qk, + ( + qkv.clone(), + q_norm.weight, + k_norm.weight, + workspace, + hq, + hk, + local_rank, + world_size, + eps, + ), + ) + fused_q, fused_k = torch.ops._C.minimax_allreduce_rms_qk( + qkv.clone(), + q_norm.weight, + k_norm.weight, + workspace, + hq, + hk, + local_rank, + world_size, + eps, + ) + _, _, fused_v = qkv.split([hq, hk, hk], dim=-1) + torch.accelerator.synchronize() + + torch.testing.assert_close( + fused_q, + ref_q, + atol=3e-2, + rtol=3e-2, + ) + torch.testing.assert_close(fused_k, ref_k, atol=3e-2, rtol=3e-2) + + cleanup_dist_env_and_memory() + + +@pytest.mark.skipif( + not current_platform.is_cuda(), + reason="CUDA required", +) +@pytest.mark.parametrize("world_size", [2, 4, 8]) +@pytest.mark.parametrize("num_tokens", [1, 128, 333]) +@pytest.mark.parametrize( + "hidden_dims", + [(6144, 1024)], +) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("eps", [1e-6]) +@pytest.mark.parametrize("seed", [42]) +def test_minimax_reduce_rms_qk( + world_size, + num_tokens, + hidden_dims, + dtype, + eps, + seed, +): + num_gpus = current_platform.device_count() + if num_gpus < world_size: + pytest.skip(f"Need >= {world_size} GPUs, have {num_gpus}") + hidden_q_full, hidden_k_full = hidden_dims + port = str(get_open_port()) + spawn( + _worker_forward_qk, + args=( + world_size, + port, + num_tokens, + hidden_q_full, + hidden_k_full, + dtype, + seed, + eps, + ), + nprocs=world_size, + join=True, + ) diff --git a/tests/models/language/pooling/test_jina_reranker_v3.py b/tests/models/language/pooling/test_jina_reranker_v3.py new file mode 100644 index 00000000000..dcce6d5bd4a --- /dev/null +++ b/tests/models/language/pooling/test_jina_reranker_v3.py @@ -0,0 +1,275 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# ruff: noqa: E501 +import pytest +import requests +import torch +import torch.nn.functional as F + +from tests.utils import RemoteOpenAIServer +from vllm.entrypoints.pooling.pooling.protocol import PoolingResponse +from vllm.entrypoints.pooling.scoring.protocol import ScoreResponse + +model_name = "jinaai/jina-reranker-v3" +query = "What are the health benefits of green tea?" +documents = [ + "Green tea contains antioxidants called catechins that may help reduce inflammation and protect cells from damage.", + "El precio del café ha aumentado un 20% este año debido a problemas en la cadena de suministro.", + "Studies show that drinking green tea regularly can improve brain function and boost metabolism.", + "Basketball is one of the most popular sports in the United States.", + "绿茶富含儿茶素等抗氧化剂，可以降低心脏病风险，还有助于控制体重。", + "Le thé vert est riche en antioxydants et peut améliorer la fonction cérébrale.", +] + +EMBEDDING_SIZE = 512 +REFERENCE_1_VS_1 = [ + 0.345703125, + -0.10498046, + 0.314453125, + -0.1376953125, + 0.3398437500, + 0.2539062, +] +REFERENCE_1_VS_N = [ + 0.294921875, + -0.16015625, + 0.189453125, + -0.1708984375, + 0.2255859375, + 0.1640625, +] +TOL = 0.01 + + +def test_offline(vllm_runner): + with vllm_runner(model_name, runner="pooling") as llm_runner: + llm = llm_runner.get_llm() + _test_offline_1_v_1(llm) + _test_offline_1_v_n(llm) + _test_offline_n_v_n(llm) + _test_offline_token_embed_illegal_inputs(llm) + assert llm.model_config.embedding_size == EMBEDDING_SIZE + + +def test_online(): + with RemoteOpenAIServer(model_name, ["--runner", "pooling"]) as server: + _test_online_1_v_1(server) + _test_online_1_v_n(server) + _test_online_n_v_n(server) + _test_online_token_embed_illegal_inputs(server) + + +def _test_offline_1_v_1(llm): + # test llm.score + outputs = llm.score(query, documents[0]) + assert len(outputs) == 1 + assert outputs[0].outputs.score == pytest.approx(REFERENCE_1_VS_1[0], abs=TOL) + + # test llm.encode + outputs = llm.encode(documents[:1] + [query], pooling_task="token_embed") + embeds = outputs[0].outputs.data.float() + assert embeds.shape[0] == 2 + assert embeds.shape[-1] == EMBEDDING_SIZE + + doc_embeds = embeds[:-1] + query_embeds = embeds[-1] + + scores = F.cosine_similarity(query_embeds, doc_embeds) + assert scores[0] == pytest.approx(REFERENCE_1_VS_1[0], abs=TOL) + + +def _test_offline_1_v_n(llm): + # test llm.score + outputs = llm.score(query, documents) + assert len(outputs) == len(documents) + + for expected, output in zip(REFERENCE_1_VS_N, outputs): + actual = output.outputs.score + assert actual == pytest.approx(expected, abs=TOL) + + # test llm.encode + outputs = llm.encode(documents + [query], pooling_task="token_embed") + embeds = outputs[0].outputs.data.float() + assert embeds.shape[0] == len(documents) + 1 + + doc_embeds = embeds[:-1] + query_embeds = embeds[-1] + + scores = F.cosine_similarity(query_embeds, doc_embeds) + + assert len(scores) == len(documents) + for expected, actual in zip(REFERENCE_1_VS_N, scores): + assert actual == pytest.approx(expected, abs=TOL) + + +def _test_offline_n_v_n(llm): + # test llm.score + outputs = llm.score([query] * len(documents), documents) + assert len(outputs) == len(documents) + + for expected, output in zip(REFERENCE_1_VS_1, outputs): + actual = output.outputs.score + assert actual == pytest.approx(expected, abs=TOL) + + # test llm.encode + for doc, expected in zip(documents, REFERENCE_1_VS_1): + outputs = llm.encode([doc, query], pooling_task="token_embed") + embeds = outputs[0].outputs.data.float() + assert embeds.shape[0] == 2 + + doc_embeds = embeds[:-1] + query_embeds = embeds[-1] + + scores = F.cosine_similarity(query_embeds, doc_embeds) + assert scores[0] == pytest.approx(expected, abs=TOL) + + +def _test_offline_token_embed_illegal_inputs(llm): + with pytest.raises( + ValueError, match="The JinaForRanking model requires at least 2 inputs." + ): + llm.encode([query], pooling_task="token_embed") + + with pytest.raises( + ValueError, match="The JinaForRanking model only supports text as input." + ): + llm.encode([1, 2, 3], pooling_task="token_embed") + + +def _get_scores(server, query, document): + score_response = requests.post( + server.url_for("score"), + json={ + "model": model_name, + "queries": query, + "documents": document, + }, + ) + + score_response.raise_for_status() + score = ScoreResponse.model_validate(score_response.json()) + + return [d.score for d in score.data] + + +def _get_embeds(server, prompts: list[str]): + response = requests.post( + server.url_for("pooling"), + json={ + "model": model_name, + "task": "token_embed", + "input": prompts, + "encoding_format": "float", + }, + ) + response.raise_for_status() + poolings = PoolingResponse.model_validate(response.json()) + + return torch.as_tensor([d.data for d in poolings.data][0]).float() + + +def _test_online_1_v_1(server): + # test scoring api + scores = _get_scores(server, query, documents[0]) + assert len(scores) == 1 + assert scores[0] == pytest.approx(REFERENCE_1_VS_1[0], abs=TOL) + + # test pooling api + embeds = _get_embeds(server, [documents[0], query]) + assert embeds.shape[0] == 2 + assert embeds.shape[-1] == EMBEDDING_SIZE + + doc_embeds = embeds[:-1] + query_embeds = embeds[-1] + + scores = F.cosine_similarity(query_embeds, doc_embeds) + assert scores[0] == pytest.approx(REFERENCE_1_VS_1[0], abs=TOL) + + +def _test_online_1_v_n(server): + # test scoring api + scores = _get_scores(server, query, documents) + assert len(scores) == len(documents) + + for expected, actual in zip(REFERENCE_1_VS_N, scores): + assert actual == pytest.approx(expected, abs=TOL) + + # test pooling api + embeds = _get_embeds(server, documents + [query]) + assert embeds.shape[0] == len(documents) + 1 + + doc_embeds = embeds[:-1] + query_embeds = embeds[-1] + + scores = F.cosine_similarity(query_embeds, doc_embeds) + + assert len(scores) == len(documents) + for expected, actual in zip(REFERENCE_1_VS_N, scores): + assert actual == pytest.approx(expected, abs=TOL) + + +def _test_online_n_v_n(server): + # test scoring api + scores = _get_scores(server, [query] * len(documents), documents) + assert len(scores) == len(documents) + + for expected, actual in zip(REFERENCE_1_VS_1, scores): + assert actual == pytest.approx(expected, abs=TOL) + + # test pooling api + for doc, expected in zip(documents, REFERENCE_1_VS_1): + embeds = _get_embeds(server, [doc, query]) + assert embeds.shape[0] == 2 + + doc_embeds = embeds[:-1] + query_embeds = embeds[-1] + + scores = F.cosine_similarity(query_embeds, doc_embeds) + assert len(scores) == 1 + assert scores[0] == pytest.approx(expected, abs=TOL) + + +def _test_online_token_embed_illegal_inputs(server): + response = requests.post( + server.url_for("pooling"), + json={ + "model": model_name, + "task": "token_embed", + "input": [query], + "encoding_format": "float", + }, + ) + assert response.json()["error"]["message"].startswith( + "The JinaForRanking model requires at least 2 inputs." + ) + + response = requests.post( + server.url_for("pooling"), + json={ + "model": model_name, + "task": "token_embed", + "input": [1, 2, 3], + "encoding_format": "float", + }, + ) + assert response.json()["error"]["message"].startswith( + "The JinaForRanking model only supports text as input." + ) + + response = requests.post( + server.url_for("pooling"), + json={ + "model": model_name, + "task": "token_embed", + "messages": [ + { + "role": "user", + "content": "The cat sat on the mat.", + } + ], + "encoding_format": "float", + }, + ) + assert response.json()["error"]["message"].startswith( + "The JinaForRanking does not support chat Request." + ) diff --git a/tests/models/registry.py b/tests/models/registry.py index b9ea5103594..45733f7d435 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -645,6 +645,7 @@ _LATE_INTERACTION_EXAMPLE_MODELS = { trust_remote_code=True, hf_overrides={"architectures": ["ColBERTLfm2Model"]}, ), + "JinaForRanking": _HfExamplesInfo("jinaai/jina-reranker-v3"), # [Multimodal] "ColModernVBertForRetrieval": _HfExamplesInfo( "ModernVBERT/colmodernvbert-merged", diff --git a/tests/v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh b/tests/v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh new file mode 100755 index 00000000000..a80950b3413 --- /dev/null +++ b/tests/v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh @@ -0,0 +1,174 @@ +#!/bin/bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# Integration edge-case tests for MultiConnector (NixlConnector + OffloadingConnector). +# +# Launches a P/D setup where both prefill and decode instances use MultiConnector +# wrapping NixlConnector and OffloadingConnector, then runs scenario-based edge +# case tests including Prometheus metrics validation. +# +# Tests cover: block-size boundaries, decode-side cache-hit scenarios +# (cold / full / partial), direct decode (control), and prefill-side CPU +# offload recovery after GPU eviction. +# +# Usage: +# bash tests/v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh +# +# Environment variables: +# MODEL_NAMES - model to test (default: Qwen/Qwen3-0.6B) +# KV_CACHE_MEMORY_BYTES - GPU KV cache size in bytes (default: 268435456 = 256 MiB) +# BLOCK_SIZE - KV cache block size (default: 128) +# VLLM_SERVE_EXTRA_ARGS - comma-separated extra args for vllm serve +set -xe + +# ── Configuration ──────────────────────────────────────────────────────── + +MODEL_NAMES=${MODEL_NAMES:-} +if [[ -n "$MODEL_NAMES" ]]; then + MODELS=("$MODEL_NAMES") +else + MODELS=("Qwen/Qwen3-0.6B") +fi + +KV_CACHE_MEMORY_BYTES=${KV_CACHE_MEMORY_BYTES:-268435456} # 256 MiB +MAX_MODEL_LEN=${MAX_MODEL_LEN:-2048} +BLOCK_SIZE=${BLOCK_SIZE:-128} +VLLM_SERVE_EXTRA_ARGS=${VLLM_SERVE_EXTRA_ARGS:-} + +GIT_ROOT=$(git rev-parse --show-toplevel) + +# ── KV transfer config ────────────────────────────────────────────────── + +KV_CONFIG='{ + "kv_connector":"MultiConnector", + "kv_role":"kv_both", + "kv_connector_extra_config":{ + "connectors":[ + {"kv_connector":"NixlConnector","kv_role":"kv_both"}, + {"kv_connector":"OffloadingConnector","kv_role":"kv_both", + "kv_connector_extra_config":{"cpu_bytes_to_use":2147483648}} + ] + } +}' +KV_CONFIG=$(echo "$KV_CONFIG" | tr -d '[:space:]') + +# ── Helpers ────────────────────────────────────────────────────────────── + +trap 'kill $(jobs -pr) 2>/dev/null || true' SIGINT SIGTERM EXIT + +wait_for_server() { + local port=$1 + timeout 1200 bash -c " + until curl -s localhost:${port}/v1/completions > /dev/null; do + sleep 1 + done" && return 0 || return 1 +} + +cleanup_instances() { + echo "Cleaning up any running vLLM instances and proxy..." + pkill -f "vllm serve" || true + pkill -f "toy_proxy_server.py" || true + sleep 2 +} + +# ── Run tests for one model ────────────────────────────────────────────── + +run_tests_for_model() { + local model_name=$1 + + echo "================================================================" + echo "Testing model: $model_name (MultiConnector edge cases)" + echo "================================================================" + + local PREFILL_PORT=8100 + local DECODE_PORT=8200 + local PROXY_PORT=8192 + local PREFILL_GPU=0 + local DECODE_GPU=1 + local PREFILL_SIDE_CHANNEL_PORT=5559 + local DECODE_SIDE_CHANNEL_PORT=5659 + + # ── Start prefill instance ── + echo "Starting prefill instance on GPU $PREFILL_GPU, port $PREFILL_PORT" + BASE_CMD="CUDA_VISIBLE_DEVICES=$PREFILL_GPU \ + VLLM_KV_CACHE_LAYOUT='HND' \ + UCX_NET_DEVICES=all \ + VLLM_NIXL_SIDE_CHANNEL_PORT=$PREFILL_SIDE_CHANNEL_PORT \ + vllm serve \"$model_name\" \ + --port $PREFILL_PORT \ + --enforce-eager \ + --block-size ${BLOCK_SIZE} \ + --max-model-len $MAX_MODEL_LEN \ + --kv-cache-memory-bytes $KV_CACHE_MEMORY_BYTES \ + --tensor-parallel-size 1 \ + --kv-transfer-config '$KV_CONFIG'" + + if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then + IFS=',' read -r -a extra_args <<< "$VLLM_SERVE_EXTRA_ARGS" + for arg in "${extra_args[@]}"; do + BASE_CMD="${BASE_CMD} $arg" + done + fi + eval "$BASE_CMD &" + + # ── Start decode instance ── + echo "Starting decode instance on GPU $DECODE_GPU, port $DECODE_PORT" + BASE_CMD="CUDA_VISIBLE_DEVICES=$DECODE_GPU \ + VLLM_KV_CACHE_LAYOUT='HND' \ + UCX_NET_DEVICES=all \ + VLLM_NIXL_SIDE_CHANNEL_PORT=$DECODE_SIDE_CHANNEL_PORT \ + vllm serve \"$model_name\" \ + --port $DECODE_PORT \ + --enforce-eager \ + --block-size ${BLOCK_SIZE} \ + --max-model-len $MAX_MODEL_LEN \ + --kv-cache-memory-bytes $KV_CACHE_MEMORY_BYTES \ + --tensor-parallel-size 1 \ + --kv-transfer-config '$KV_CONFIG'" + + if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then + IFS=',' read -r -a extra_args <<< "$VLLM_SERVE_EXTRA_ARGS" + for arg in "${extra_args[@]}"; do + BASE_CMD="${BASE_CMD} $arg" + done + fi + eval "$BASE_CMD &" + + # ── Wait for servers ── + echo "Waiting for prefill instance on port $PREFILL_PORT to start..." + wait_for_server "$PREFILL_PORT" + echo "Waiting for decode instance on port $DECODE_PORT to start..." + wait_for_server "$DECODE_PORT" + + # ── Start proxy ── + echo "Starting proxy server on port $PROXY_PORT" + python3 "${GIT_ROOT}/tests/v1/kv_connector/nixl_integration/toy_proxy_server.py" \ + --port "$PROXY_PORT" \ + --prefiller-hosts localhost \ + --prefiller-ports "$PREFILL_PORT" \ + --decoder-hosts localhost \ + --decoder-ports "$DECODE_PORT" & + sleep 5 + + # ── Run edge case tests ── + echo "Running MultiConnector edge case tests for $model_name" + PREFILL_PORT=$PREFILL_PORT \ + DECODE_PORT=$DECODE_PORT \ + PROXY_PORT=$PROXY_PORT \ + BLOCK_SIZE=$BLOCK_SIZE \ + python3 -m pytest -s -x \ + "${GIT_ROOT}/tests/v1/kv_connector/nixl_integration/test_multi_connector_edge_cases.py" + + # ── Cleanup ── + cleanup_instances + sleep 3 +} + +# ── Main ───────────────────────────────────────────────────────────────── + +for model in "${MODELS[@]}"; do + run_tests_for_model "$model" +done + +echo "All MultiConnector edge case tests passed!" diff --git a/tests/v1/kv_connector/nixl_integration/test_multi_connector_edge_cases.py b/tests/v1/kv_connector/nixl_integration/test_multi_connector_edge_cases.py new file mode 100644 index 00000000000..1886860035d --- /dev/null +++ b/tests/v1/kv_connector/nixl_integration/test_multi_connector_edge_cases.py @@ -0,0 +1,477 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Integration edge-case tests for MultiConnector (NixlConnector + OffloadingConnector). + +Tests cover: + - Output correctness across block-size boundaries (proxy vs direct prefill). + - Decode-side Prometheus metrics validation (local_cache_hit, + external_kv_transfer, local_compute) for cold/warm/partial cache scenarios. + - Prefill-side CPU offload recovery after GPU cache eviction. + +Requires running servers started by run_multi_connector_edge_case_test.sh. +""" + +import os +import time +import urllib.request + +import openai +import regex as re + +# ── Server configuration from environment ───────────────────────────────── + +PREFILL_HOST = os.getenv("PREFILL_HOST", "localhost") +PREFILL_PORT = os.environ["PREFILL_PORT"] +DECODE_HOST = os.getenv("DECODE_HOST", "localhost") +DECODE_PORT = os.environ["DECODE_PORT"] +PROXY_HOST = os.getenv("PROXY_HOST", "localhost") +PROXY_PORT = os.environ["PROXY_PORT"] +BLOCK_SIZE = int(os.getenv("BLOCK_SIZE", "128")) + +# ── OpenAI clients ──────────────────────────────────────────────────────── + +decode_client = openai.OpenAI( + api_key="EMPTY", + base_url=f"http://{DECODE_HOST}:{DECODE_PORT}/v1", +) +prefill_client = openai.OpenAI( + api_key="EMPTY", + base_url=f"http://{PREFILL_HOST}:{PREFILL_PORT}/v1", +) +proxy_client = openai.OpenAI( + api_key="EMPTY", + base_url=f"http://{PROXY_HOST}:{PROXY_PORT}/v1", +) + +_MODEL = None + + +def _get_model() -> str: + global _MODEL + if _MODEL is None: + models = decode_client.models.list() + _MODEL = models.data[0].id + return _MODEL + + +def _complete(client: openai.OpenAI, prompt: str, max_tokens: int = 20): + """Send a completion request and return (text, prompt_tokens).""" + resp = client.completions.create( + model=_get_model(), + prompt=prompt, + max_tokens=max_tokens, + temperature=0, + ) + return resp.choices[0].text, resp.usage.prompt_tokens + + +# ── Prometheus metrics helpers ──────────────────────────────────────────── + +_METRIC_RE = re.compile( + r'vllm:prompt_tokens_by_source_total\{.*?source="([^"]+)".*?\}\s+' + r"([\d.eE+\-]+)" +) + + +def _fetch_metrics(host: str, port: str) -> dict[str, float]: + """Scrape prompt_tokens_by_source counters from a vLLM server.""" + body = urllib.request.urlopen(f"http://{host}:{port}/metrics").read().decode() + result = { + "local_compute": 0.0, + "local_cache_hit": 0.0, + "external_kv_transfer": 0.0, + } + for m in _METRIC_RE.finditer(body): + source, val = m.group(1), float(m.group(2)) + if source in result: + result[source] += val + return result + + +def _fetch_decode_metrics() -> dict[str, float]: + return _fetch_metrics(DECODE_HOST, DECODE_PORT) + + +def _fetch_prefill_metrics() -> dict[str, float]: + return _fetch_metrics(PREFILL_HOST, PREFILL_PORT) + + +_NIXL_BYTES_RE = re.compile(r"vllm:nixl_bytes_transferred_sum\b.*?\s+([\d.eE+\-]+)") + + +def _fetch_nixl_bytes(host: str, port: str) -> float: + """Scrape total NIXL bytes transferred from a vLLM server.""" + body = urllib.request.urlopen(f"http://{host}:{port}/metrics").read().decode() + total = 0.0 + for m in _NIXL_BYTES_RE.finditer(body): + total += float(m.group(1)) + return total + + +_OFFLOAD_BYTES_RE = re.compile( + r'vllm:kv_offload_total_bytes_total\{.*?transfer_type="([^"]+)".*?\}\s+' + r"([\d.eE+\-]+)" +) + + +def _fetch_offload_bytes(host: str, port: str) -> dict[str, float]: + """Scrape kv_offload_total_bytes counters (CPU_to_GPU / GPU_to_CPU).""" + body = urllib.request.urlopen(f"http://{host}:{port}/metrics").read().decode() + result = {"CPU_to_GPU": 0.0, "GPU_to_CPU": 0.0} + for m in _OFFLOAD_BYTES_RE.finditer(body): + transfer_type, val = m.group(1), float(m.group(2)) + if transfer_type in result: + result[transfer_type] += val + return result + + +def _metrics_delta(before: dict, after: dict) -> dict[str, float]: + return {k: after.get(k, 0) - before.get(k, 0) for k in before} + + +# ── Prompts (unique per test to avoid cross-test cache interference) ────── + +SHORT_PROMPT = "Red Hat is " + +MEDIUM_PROMPT = ( + "Red Hat is the best company in the world to work for because it works " + "on open source software, which means that all the contributions are " + "delivered to the community. As a result," +) + + +def _make_prompt(n_tokens: int) -> str: + """Build a prompt of ~n_tokens tokens (1 word ~ 1 token).""" + return "word " * n_tokens + + +BLOCK_BOUNDARY_PROMPT = _make_prompt(BLOCK_SIZE) +ABOVE_BOUNDARY_PROMPT = _make_prompt(BLOCK_SIZE + 2) +MULTI_BLOCK_PROMPT = _make_prompt(BLOCK_SIZE * 4) + +FULL_CACHE_HIT_PROMPT = ( # noqa: E501 + "The history of computing begins with Charles Babbage who designed the " + "Analytical Engine in the 1830s which is considered the first general " + "purpose computer design in history. Ada Lovelace is widely regarded as " + "the first computer programmer for her work on the Analytical Engine. " + "The modern era of computing began with Alan Turing who formalized the " + "concept of computation with his Turing machine in 1936. During World " + "War Two Turing worked at Bletchley Park to break the Enigma cipher. " + "After the war the first electronic computers were built including ENIAC " + "at the University of Pennsylvania and Colossus at Bletchley Park. " + "These early machines filled entire rooms and used vacuum tubes for logic. " + "The invention of the transistor at Bell Labs in 1947 revolutionized " + "computing by making smaller and more reliable machines possible. " + "The integrated circuit followed in the late 1950s combining multiple " + "transistors on a single chip. This led to the microprocessor in the 1970s " + "and eventually to the personal computer revolution of the 1980s." +) + +PARTIAL_CACHE_PREFIX = ( # noqa: E501 + "Machine learning has transformed the field of artificial intelligence " + "by enabling computers to learn patterns from data without being " + "explicitly programmed for every task. The field has evolved dramatically " + "since its inception in the 1950s when Arthur Samuel coined the term while " + "working at IBM. Early approaches focused on symbolic reasoning and expert " + "systems that encoded human knowledge as rules. The statistical revolution " + "of the 1990s shifted the paradigm toward data driven methods. Support " + "vector machines and random forests became popular for classification tasks. " + "The breakthrough of deep learning in 2012 with AlexNet winning ImageNet " + "changed everything. Neural networks with many layers could automatically " + "learn hierarchical feature representations from raw data." +) +PARTIAL_CACHE_EXTENDED = PARTIAL_CACHE_PREFIX + ( + " Transformers have become the dominant architecture for natural language " + "processing tasks including translation, summarization, and generation. " + "The attention mechanism allows models to weigh the importance of different " + "parts of the input sequence. Large language models like GPT and BERT " + "demonstrated that pre-training on massive text corpora followed by fine " + "tuning on specific tasks could achieve state of the art results across " + "a wide range of benchmarks. Scaling laws suggest that larger models " + "trained on more data continue to improve in capability." +) + +# ═══════════════════════════════════════════════════════════════════════════ +# Output correctness across block-size boundaries (decode-side metrics) +# +# Each test sends via proxy, verifies output matches prefill_direct at +# temperature=0, and checks decode-side metrics for NIXL transfer. +# ═══════════════════════════════════════════════════════════════════════════ + + +def test_short_prompt_correctness(): + """Short prompt (< block_size): output matches prefill, NIXL used.""" + n0 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT) + m0 = _fetch_decode_metrics() + proxy_text, _ = _complete(proxy_client, SHORT_PROMPT) + time.sleep(1) + m1 = _fetch_decode_metrics() + n1 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT) + d = _metrics_delta(m0, m1) + + prefill_text, _ = _complete(prefill_client, SHORT_PROMPT) + print(f"SHORT PROMPT: {proxy_text=}, nixl_bytes_delta={n1 - n0}") + assert proxy_text == prefill_text + assert d["external_kv_transfer"] > 0, ( + "NIXL transfer did not occur — decode may have silently fallen back " + "to local compute" + ) + assert n1 - n0 > 0, ( + f"expected nixl_bytes_transferred to increase, got delta={n1 - n0}" + ) + + +def test_block_boundary_correctness(): + """Exactly block_size tokens: output matches prefill, NIXL used.""" + n0 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT) + m0 = _fetch_decode_metrics() + proxy_text, pt = _complete(proxy_client, BLOCK_BOUNDARY_PROMPT) + time.sleep(1) + m1 = _fetch_decode_metrics() + n1 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT) + d = _metrics_delta(m0, m1) + + prefill_text, _ = _complete(prefill_client, BLOCK_BOUNDARY_PROMPT) + print(f"BLOCK BOUNDARY: {pt} prompt tokens, nixl_bytes_delta={n1 - n0}") + assert proxy_text == prefill_text + assert d["external_kv_transfer"] > 0, ( + "NIXL transfer did not occur — decode may have silently fallen back " + "to local compute" + ) + assert n1 - n0 > 0, ( + f"expected nixl_bytes_transferred to increase, got delta={n1 - n0}" + ) + + +def test_above_block_boundary_correctness(): + """Just above block_size (partial second block): output matches prefill.""" + n0 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT) + m0 = _fetch_decode_metrics() + proxy_text, pt = _complete(proxy_client, ABOVE_BOUNDARY_PROMPT) + time.sleep(1) + m1 = _fetch_decode_metrics() + n1 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT) + d = _metrics_delta(m0, m1) + + prefill_text, _ = _complete(prefill_client, ABOVE_BOUNDARY_PROMPT) + print(f"ABOVE BOUNDARY: {pt} prompt tokens, nixl_bytes_delta={n1 - n0}") + assert proxy_text == prefill_text + assert d["external_kv_transfer"] > 0, ( + "NIXL transfer did not occur — decode may have silently fallen back " + "to local compute" + ) + assert n1 - n0 > 0, ( + f"expected nixl_bytes_transferred to increase, got delta={n1 - n0}" + ) + + +def test_multi_block_correctness(): + """Multi-block prompt (~4x block_size): output matches prefill.""" + n0 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT) + m0 = _fetch_decode_metrics() + proxy_text, pt = _complete(proxy_client, MULTI_BLOCK_PROMPT) + time.sleep(1) + m1 = _fetch_decode_metrics() + n1 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT) + d = _metrics_delta(m0, m1) + + prefill_text, _ = _complete(prefill_client, MULTI_BLOCK_PROMPT) + print(f"MULTI BLOCK: {pt} prompt tokens, nixl_bytes_delta={n1 - n0}") + assert proxy_text == prefill_text + assert d["external_kv_transfer"] > 0, ( + "NIXL transfer did not occur — decode may have silently fallen back " + "to local compute" + ) + assert n1 - n0 > 0, ( + f"expected nixl_bytes_transferred to increase, got delta={n1 - n0}" + ) + + +# ═══════════════════════════════════════════════════════════════════════════ +# Decode-side KV source validation via Prometheus metrics +# +# Scrape vllm:prompt_tokens_by_source_total from the DECODE server to +# verify which code path (GPU prefix, NIXL, local compute) was exercised. +# ═══════════════════════════════════════════════════════════════════════════ + + +def test_cold_decode_no_cache_hit_metrics(): + """Cold decode: external_kv_transfer==P, local_cache_hit==0.""" + n0 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT) + m0 = _fetch_decode_metrics() + proxy_text, P = _complete(proxy_client, MEDIUM_PROMPT) + time.sleep(1) + m1 = _fetch_decode_metrics() + n1 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT) + d = _metrics_delta(m0, m1) + + print(f"COLD DECODE: {P} prompt tokens, metrics delta: {d}") + print(f" nixl_bytes_delta={n1 - n0}") + assert len(proxy_text) > 0, "proxy returned empty response" + assert d["external_kv_transfer"] == P, ( + f"expected external_kv_transfer={P}, got {d['external_kv_transfer']}" + ) + assert d["local_compute"] == 1, ( + f"expected local_compute=1, got {d['local_compute']}" + ) + assert d["local_cache_hit"] == 0, ( + f"expected local_cache_hit=0, got {d['local_cache_hit']}" + ) + assert n1 - n0 > 0, ( + f"expected nixl_bytes_transferred to increase, got delta={n1 - n0}" + ) + + +def test_full_decode_gpu_cache_hit_metrics(): + """Prime decode, resend via proxy: local_cache_hit==cached blocks.""" + decode_text, _ = _complete(decode_client, FULL_CACHE_HIT_PROMPT) + + n0 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT) + m0 = _fetch_decode_metrics() + proxy_text, P = _complete(proxy_client, FULL_CACHE_HIT_PROMPT) + time.sleep(1) + m1 = _fetch_decode_metrics() + n1 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT) + d = _metrics_delta(m0, m1) + + cached = (P // BLOCK_SIZE) * BLOCK_SIZE + expected_nixl = P - cached + + print(f"FULL CACHE HIT: {P} tokens, cached={cached}, nixl={expected_nixl}") + print(f" metrics delta: {d}, nixl_bytes_delta={n1 - n0}") + assert len(proxy_text) > 0, "proxy returned empty response" + assert d["local_cache_hit"] == cached, ( + f"expected local_cache_hit={cached}, got {d['local_cache_hit']}" + ) + assert d["external_kv_transfer"] == expected_nixl, ( + f"expected external_kv_transfer={expected_nixl}, " + f"got {d['external_kv_transfer']}" + ) + assert d["local_compute"] == 1, ( + f"expected local_compute=1 (recomputed last token), got {d['local_compute']}" + ) + assert n1 - n0 > 0, ( + f"expected nixl_bytes_transferred to increase (partial NIXL for " + f"uncached tail), got delta={n1 - n0}" + ) + + +def test_partial_decode_gpu_cache_hit_metrics(): + """Prime with prefix, extend via proxy: partial local_cache_hit.""" + _, prefix_tokens = _complete(decode_client, PARTIAL_CACHE_PREFIX) + cached = (prefix_tokens // BLOCK_SIZE) * BLOCK_SIZE + assert cached >= BLOCK_SIZE, ( + f"PARTIAL_CACHE_PREFIX too short ({prefix_tokens} tokens) for partial " + f"cache hit test with block_size={BLOCK_SIZE}" + ) + + n0 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT) + m0 = _fetch_decode_metrics() + proxy_text, P = _complete(proxy_client, PARTIAL_CACHE_EXTENDED) + time.sleep(1) + m1 = _fetch_decode_metrics() + n1 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT) + d = _metrics_delta(m0, m1) + + expected_nixl = P - cached + + print(f"PARTIAL CACHE HIT: {P} tokens, cached={cached}, nixl={expected_nixl}") + print(f" metrics delta: {d}, nixl_bytes_delta={n1 - n0}") + assert len(proxy_text) > 0, "proxy returned empty response" + assert d["external_kv_transfer"] == expected_nixl, ( + f"expected external_kv_transfer={expected_nixl}, " + f"got {d['external_kv_transfer']}" + ) + assert d["local_cache_hit"] == cached, ( + f"expected local_cache_hit={cached}, got {d['local_cache_hit']}" + ) + assert d["local_compute"] == 1, ( + f"expected local_compute=1 (recomputed last token), got {d['local_compute']}" + ) + assert n1 - n0 > 0, ( + f"expected nixl_bytes_transferred to increase (NIXL for uncached " + f"tail), got delta={n1 - n0}" + ) + + +def test_decode_direct_all_local_compute(): + """Direct decode (no proxy): local_compute==P, no transfers.""" + prompt = "The speed of light is approximately" + n0 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT) + m0 = _fetch_decode_metrics() + text, P = _complete(decode_client, prompt) + time.sleep(1) + m1 = _fetch_decode_metrics() + n1 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT) + d = _metrics_delta(m0, m1) + + print(f"DIRECT DECODE: {text!r} ({P} tokens), metrics delta: {d}") + print(f" nixl_bytes_delta={n1 - n0}") + assert len(text.strip()) > 0, "empty output from direct decode" + assert d["local_compute"] == P, ( + f"expected local_compute={P}, got {d['local_compute']}" + ) + assert d["external_kv_transfer"] == 0, ( + f"expected external_kv_transfer=0, got {d['external_kv_transfer']}" + ) + assert n1 - n0 == 0, ( + f"expected no nixl_bytes_transferred for direct decode, got delta={n1 - n0}" + ) + + +# ═══════════════════════════════════════════════════════════════════════════ +# Prefill-side CPU offload validation via Prometheus metrics +# +# Scrape vllm:prompt_tokens_by_source_total from the PREFILL server. +# Exercises the OffloadingConnector read path: after GPU cache eviction, +# the OffloadingConnector restores KV from CPU (NixlConnector cannot help +# for direct requests without kv_transfer_params). +# ═══════════════════════════════════════════════════════════════════════════ + +EVICTION_PROMPT = ( # noqa: E501 + "Quantum computing leverages quantum mechanical phenomena like " + "superposition and entanglement to perform computations that would be " + "intractable for classical computers. This has implications for " + "cryptography, drug discovery, and optimization problems. Richard Feynman " + "first proposed the idea of quantum computing in 1982 when he observed " + "that simulating quantum systems on classical computers was exponentially " + "hard. Peter Shor developed a quantum algorithm for factoring large " + "numbers in polynomial time which threatens RSA encryption. Grover search " + "algorithm provides a quadratic speedup for unstructured search problems. " + "Companies like IBM Google and Rigetti are building quantum processors " + "with increasing numbers of qubits. Error correction remains a major " + "challenge as quantum states are extremely fragile and prone to decoherence." +) + + +def test_prefill_cpu_offload_after_gpu_eviction(): + """Prefill-side: evict GPU, re-request directly, CPU offload restores KV.""" + text1, P = _complete(prefill_client, EVICTION_PROMPT, max_tokens=30) + + for i in range(100): + _complete(prefill_client, f"Eviction prompt number {i}: " + _make_prompt(200)) + + ob0 = _fetch_offload_bytes(PREFILL_HOST, PREFILL_PORT) + m0 = _fetch_prefill_metrics() + text2, _ = _complete(prefill_client, EVICTION_PROMPT, max_tokens=30) + + cpu_to_gpu_delta = 0.0 + for _ in range(10): + time.sleep(1) + ob1 = _fetch_offload_bytes(PREFILL_HOST, PREFILL_PORT) + cpu_to_gpu_delta = ob1["CPU_to_GPU"] - ob0["CPU_to_GPU"] + if cpu_to_gpu_delta > 0: + break + + m1 = _fetch_prefill_metrics() + d = _metrics_delta(m0, m1) + + print(f"PREFILL CPU OFFLOAD: run1={text1[:60]!r}, run2={text2[:60]!r}") + print(f" prefill metrics delta: {d}") + print(f" cpu_to_gpu bytes delta: {cpu_to_gpu_delta}") + assert text1 == text2, f"inconsistent after eviction: {text1=!r}, {text2=!r}" + assert cpu_to_gpu_delta > 0, ( + f"expected cpu_to_gpu bytes > 0 (OffloadingConnector should restore " + f"KV from CPU to GPU), got {cpu_to_gpu_delta}" + ) diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 0c2a53ec02e..01ac03f27a3 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -3491,3 +3491,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), + ) diff --git a/vllm/compilation/passes/fusion/minimax_qk_norm_fusion.py b/vllm/compilation/passes/fusion/minimax_qk_norm_fusion.py new file mode 100644 index 00000000000..7445028da63 --- /dev/null +++ b/vllm/compilation/passes/fusion/minimax_qk_norm_fusion.py @@ -0,0 +1,340 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +Fusion pass: replace MiniMax QK allreduce + RMS norm with the Lamport +fused kernel (minimax_allreduce_rms_qk) for decode-size batches. + +Pattern (inlined forward_qk in compiled graph): + q, k, v = qkv.split([q_size, kv_size, kv_size], -1) + q_fp32 = q.to(float32); k_fp32 = k.to(float32) + q_var = q_fp32.pow(2).mean(-1, keepdim=True) + k_var = k_fp32.pow(2).mean(-1, keepdim=True) + qk_var = cat([q_var, k_var], -1) + qk_var = allreduce(qk_var) / tp_world + q_var, k_var = qk_var.chunk(2, -1) + q_out = (q_fp32 * rsqrt(q_var + eps) * q_weight).to(orig_dtype) + k_out = (k_fp32 * rsqrt(k_var + eps) * k_weight).to(orig_dtype) + return q_out, k_out, v + +Replacement (pure, no in-place on qkv/q/k): + q_out, k_out = minimax_qk_norm_fused(qkv, q_weight, k_weight, workspace, ...) + v = qkv.split([q_size, kv_size, kv_size], -1)[2] + return q_out, k_out, v + +is_applicable_for_range: only fires for compile_range.end <= max_decode_tokens +so that large prefill batches fall through to the original forward_qk (= main). +""" + +import torch +import torch._inductor.pattern_matcher as pm +import torch.fx as fx +from torch._inductor.pattern_matcher import PatternMatcherPass + +from vllm.config import VllmConfig +from vllm.config.utils import Range +from vllm.distributed import tensor_model_parallel_all_reduce +from vllm.distributed.parallel_state import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) +from vllm.logger import init_logger +from vllm.utils.torch_utils import direct_register_custom_op + +from ..inductor_pass import enable_fake_mode +from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass + +logger = init_logger(__name__) + +MAX_TOKEN_NUM = 2048 + +_MINIMAX_QK_NORM_FUSED_OP = None +if hasattr(torch.ops._C, "minimax_allreduce_rms_qk"): + + def _minimax_qk_norm_fused( + qkv: torch.Tensor, + norm_weight_q: torch.Tensor, + norm_weight_k: torch.Tensor, + q_size: int, + kv_size: int, + rank: int, + nranks: int, + eps: float, + max_tokens: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + from vllm.distributed.parallel_state import get_tp_group + from vllm.model_executor.layers.mamba.lamport_workspace import ( + get_allreduce_workspace, + ) + + workspace = get_allreduce_workspace( + rank=rank, + world_size=nranks, + max_tokens=max_tokens, + process_group=get_tp_group().cpu_group, + ) + return torch.ops._C.minimax_allreduce_rms_qk( + qkv, + norm_weight_q, + norm_weight_k, + workspace, + q_size, + kv_size, + rank, + nranks, + eps, + ) + + def _minimax_qk_norm_fused_fake( + qkv: torch.Tensor, + norm_weight_q: torch.Tensor, + norm_weight_k: torch.Tensor, + q_size: int, + kv_size: int, + rank: int, + nranks: int, + eps: float, + max_tokens: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + T = qkv.shape[0] + return ( + torch.empty([T, q_size], dtype=qkv.dtype, device=qkv.device), + torch.empty([T, kv_size], dtype=qkv.dtype, device=qkv.device), + ) + + direct_register_custom_op( + op_name="minimax_qk_norm_fused", + op_func=_minimax_qk_norm_fused, + fake_impl=_minimax_qk_norm_fused_fake, + mutates_args=[], + ) + _MINIMAX_QK_NORM_FUSED_OP = torch.ops.vllm.minimax_qk_norm_fused.default + + +class MiniMaxQKNormPattern: + """ + Match the forward_qk allreduce+rms pattern and replace with Lamport kernel. + """ + + def __init__( + self, + q_size: int, + kv_size: int, + eps: float, + tp_world: int, + tp_rank: int, + max_tokens: int, + dtype: torch.dtype, + device: str | None, + ) -> None: + self.q_size = q_size + self.kv_size = kv_size + self.eps = eps + self.tp_world = tp_world + self.tp_rank = tp_rank + self.max_tokens = max_tokens + self.dtype = dtype + self.device = device + + def get_inputs(self) -> list[torch.Tensor]: + T = 4 + qkv = torch.empty( + [T, self.q_size + 2 * self.kv_size], + device=self.device, + dtype=self.dtype, + ) + q_weight = torch.empty([self.q_size], device=self.device, dtype=self.dtype) + k_weight = torch.empty([self.kv_size], device=self.device, dtype=self.dtype) + return [qkv, q_weight, k_weight] + + def register(self, pm_pass: PatternMatcherPass) -> None: + q_size = self.q_size + kv_size = self.kv_size + eps = self.eps + tp_world = self.tp_world + max_tokens = self.max_tokens + tp_rank = self.tp_rank + dtype = self.dtype + + def pattern( + qkv: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + q, k, v = qkv.split([q_size, kv_size, kv_size], dim=-1) + q_fp32 = q.to(torch.float32) + k_fp32 = k.to(torch.float32) + q_var = q_fp32.pow(2).mean(dim=-1, keepdim=True) + k_var = k_fp32.pow(2).mean(dim=-1, keepdim=True) + qk_var = torch.cat([q_var, k_var], dim=-1) + qk_var = tensor_model_parallel_all_reduce(qk_var) / tp_world + q_var, k_var = qk_var.chunk(2, dim=-1) + q_out = (q_fp32 * torch.rsqrt(q_var + eps) * q_weight).to(dtype) + k_out = (k_fp32 * torch.rsqrt(k_var + eps) * k_weight).to(dtype) + return q_out, k_out, v + + def replacement( + qkv: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + assert _MINIMAX_QK_NORM_FUSED_OP is not None + q_out, k_out = torch.ops.vllm.minimax_qk_norm_fused( + qkv, + q_weight, + k_weight, + q_size, + kv_size, + tp_rank, + tp_world, + eps, + max_tokens, + ) + _, _, v = qkv.split([q_size, kv_size, kv_size], dim=-1) + return q_out, k_out, v + + pm.register_replacement( + pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass + ) + + # Second pattern: three separate split_with_sizes nodes (one per output), + # each with _users=1. This occurs when the QKV projection uses a + # functional GEMM kernel (e.g. cutlass_scaled_mm via auto_functionalized), + # which causes inductor to generate one split per consumer. + def pattern_split3( + qkv: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + q = qkv.split([q_size, kv_size, kv_size], dim=-1)[0] + k = qkv.split([q_size, kv_size, kv_size], dim=-1)[1] + v = qkv.split([q_size, kv_size, kv_size], dim=-1)[2] + q_fp32 = q.to(torch.float32) + k_fp32 = k.to(torch.float32) + q_var = q_fp32.pow(2).mean(dim=-1, keepdim=True) + k_var = k_fp32.pow(2).mean(dim=-1, keepdim=True) + qk_var = torch.cat([q_var, k_var], dim=-1) + qk_var = tensor_model_parallel_all_reduce(qk_var) / tp_world + q_var, k_var = qk_var.chunk(2, dim=-1) + q_out = (q_fp32 * torch.rsqrt(q_var + eps) * q_weight).to(dtype) + k_out = (k_fp32 * torch.rsqrt(k_var + eps) * k_weight).to(dtype) + return q_out, k_out, v + + pm.register_replacement( + pattern_split3, replacement, self.get_inputs(), pm.fwd_only, pm_pass + ) + + +class MiniMaxQKNormPass(VllmPatternMatcherPass): + """ + Replace forward_qk allreduce+norm with the Lamport fused kernel. + Only applied for decode-size compile ranges (small token counts). + """ + + def __init__(self, config: VllmConfig) -> None: + super().__init__(config) + self.disabled = True + + if _MINIMAX_QK_NORM_FUSED_OP is None: + logger.warning_once( + "minimax_allreduce_rms_qk op not found, MiniMaxQKNormPass disabled." + ) + return + + tp_world = get_tensor_model_parallel_world_size() + if tp_world <= 1: + logger.warning_once("MiniMaxQKNormPass disabled: tp_size <= 1.") + return + + if config.model_config is None: + logger.warning_once("MiniMaxQKNormPass disabled: no model_config.") + return + + hf_cfg = config.model_config.hf_config + + model_name = getattr(hf_cfg, "architectures", "")[0] + if model_name != "MiniMaxM2ForCausalLM": + return + + num_attention_heads = getattr(hf_cfg, "num_attention_heads", 0) + num_key_value_heads = getattr(hf_cfg, "num_key_value_heads", 0) + hidden_size = getattr(hf_cfg, "hidden_size", 0) + head_dim = getattr(hf_cfg, "head_dim", 0) + eps: float = getattr(hf_cfg, "rms_norm_eps", 1e-6) + + if ( + num_attention_heads != 48 + or num_key_value_heads != 8 + or hidden_size != 3072 + or head_dim != 128 + ): + logger.warning_once( + "MiniMaxQKNormPass disabled: cannot infer model info from hf_config." + ) + return + + num_heads_per_rank = num_attention_heads // tp_world + num_kv_heads_per_rank = max(1, num_key_value_heads // tp_world) + q_size = num_heads_per_rank * head_dim + kv_size = num_kv_heads_per_rank * head_dim + + self.max_token_num = min( + MAX_TOKEN_NUM, config.scheduler_config.max_num_batched_tokens + ) + + tp_rank = get_tensor_model_parallel_rank() + # Allocate Lamport workspace first. + from vllm.distributed.parallel_state import get_tp_group + from vllm.model_executor.layers.mamba.lamport_workspace import ( + get_allreduce_workspace, + ) + + get_allreduce_workspace( + rank=tp_rank, + world_size=tp_world, + max_tokens=self.max_token_num, + process_group=get_tp_group().cpu_group, + ) + + self.patterns: PatternMatcherPass = PatternMatcherPass( + pass_name="minimax_qk_norm_pass" + ) + self._register_patterns(q_size, kv_size, eps, tp_world, tp_rank) + self.dump_patterns(config, self.patterns) + self.disabled = False + + @enable_fake_mode + def _register_patterns( + self, + q_size: int, + kv_size: int, + eps: float, + tp_world: int, + tp_rank: int, + ) -> None: + MiniMaxQKNormPattern( + q_size=q_size, + kv_size=kv_size, + eps=eps, + tp_world=tp_world, + tp_rank=tp_rank, + max_tokens=self.max_token_num, + dtype=self.model_dtype, + device=self.device, + ).register(self.patterns) + + def is_applicable_for_range(self, compile_range: Range) -> bool: + if self.disabled: + return False + + return bool(compile_range.end <= self.max_token_num) + + @VllmInductorPass.time_and_log + def __call__(self, graph: fx.Graph) -> None: + if self.disabled: + return + self.matched_count = self.patterns.apply(graph) + logger.debug("MiniMaxQKNormPass replaced %s patterns", self.matched_count) + + def uuid(self) -> str: + return VllmInductorPass.hash_source(self, MiniMaxQKNormPattern) diff --git a/vllm/compilation/passes/pass_manager.py b/vllm/compilation/passes/pass_manager.py index b4823a0afde..91e10145607 100644 --- a/vllm/compilation/passes/pass_manager.py +++ b/vllm/compilation/passes/pass_manager.py @@ -38,6 +38,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, @@ -137,6 +138,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(): diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index ef2a4bf5a4f..11c933fc72f 100644 --- a/vllm/config/compilation.py +++ b/vllm/config/compilation.py @@ -26,6 +26,8 @@ from vllm.utils.torch_utils import is_torch_equal_or_newer if TYPE_CHECKING: from vllm.config import VllmConfig + from vllm.v1.attention.backend import AttentionCGSupport + from vllm.v1.kv_cache_interface import KVCacheConfig else: VllmConfig = object @@ -132,6 +134,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.""" @@ -1241,6 +1245,152 @@ class CompilationConfig: assert "none" in self.custom_ops return f"+{op}" in self.custom_ops + def resolve_cudagraph_mode_and_sizes( + self, + min_cg_support: "AttentionCGSupport", + min_cg_attn_backend: str | None, + uniform_decode_query_len: int = 1, + tensor_parallel_size: int = 1, + kv_cache_config: "KVCacheConfig | None" = None, + max_num_reqs: int | None = None, + is_profiling: bool = False, + ) -> CUDAGraphMode: + from vllm.v1.attention.backend import AttentionCGSupport + + cudagraph_mode = self.cudagraph_mode + if cudagraph_mode is None or cudagraph_mode == CUDAGraphMode.NONE: + self.cudagraph_mode = CUDAGraphMode.NONE + return CUDAGraphMode.NONE + + # Check cudagraph for mixed batch is supported + if ( + cudagraph_mode.mixed_mode() == CUDAGraphMode.FULL + and min_cg_support != AttentionCGSupport.ALWAYS + ): + msg = ( + f"CUDAGraphMode.{cudagraph_mode.name} is not supported " + f"with {min_cg_attn_backend} backend (support: " + f"{min_cg_support})" + ) + if min_cg_support == AttentionCGSupport.NEVER: + # if not supported any full cudagraphs, just raise it. + msg += ( + "; please try cudagraph_mode=PIECEWISE, and " + "make sure compilation mode is VLLM_COMPILE" + ) + raise ValueError(msg) + + # attempt to resolve the full cudagraph related mode + if self.splitting_ops_contain_attention(): + msg += "; setting cudagraph_mode=FULL_AND_PIECEWISE" + cudagraph_mode = CUDAGraphMode.FULL_AND_PIECEWISE + else: + msg += "; setting cudagraph_mode=FULL_DECODE_ONLY" + cudagraph_mode = CUDAGraphMode.FULL_DECODE_ONLY + logger.warning(msg) + + # check that if we are doing decode full-cudagraphs it is supported + if ( + cudagraph_mode.decode_mode() == CUDAGraphMode.FULL + and min_cg_support == AttentionCGSupport.NEVER + ): + msg = ( + f"CUDAGraphMode.{cudagraph_mode.name} is not supported " + f"with {min_cg_attn_backend} backend (support: " + f"{min_cg_support})" + ) + if self.mode == CompilationMode.VLLM_COMPILE and ( + self.splitting_ops_contain_attention() + or self.use_inductor_graph_partition + ): + msg += ( + "; setting cudagraph_mode=PIECEWISE because " + "attention is compiled piecewise" + ) + cudagraph_mode = CUDAGraphMode.PIECEWISE + else: + msg += ( + "; setting cudagraph_mode=NONE because " + "attention is not compiled piecewise" + ) + cudagraph_mode = CUDAGraphMode.NONE + logger.warning(msg) + + # check that if we are doing spec-decode + decode full-cudagraphs it is + # supported + if ( + cudagraph_mode.decode_mode() == CUDAGraphMode.FULL + and uniform_decode_query_len > 1 + and min_cg_support.value < AttentionCGSupport.UNIFORM_BATCH.value + ): + msg = ( + f"CUDAGraphMode.{cudagraph_mode.name} is not supported" + f" with spec-decode for attention backend " + f"{min_cg_attn_backend} (support: {min_cg_support})" + ) + if self.splitting_ops_contain_attention(): + msg += "; setting cudagraph_mode=PIECEWISE" + cudagraph_mode = CUDAGraphMode.PIECEWISE + else: + msg += "; setting cudagraph_mode=NONE" + cudagraph_mode = CUDAGraphMode.NONE + logger.warning(msg) + + # double check that we can support full cudagraph if they are requested + # even after automatic downgrades + if ( + cudagraph_mode.has_full_cudagraphs() + and min_cg_support == AttentionCGSupport.NEVER + ): + raise ValueError( + f"CUDAGraphMode.{cudagraph_mode.name} is not " + f"supported with {min_cg_attn_backend} backend (" + f"support:{min_cg_support}) " + "; please try cudagraph_mode=PIECEWISE, " + "and make sure compilation mode is VLLM_COMPILE" + ) + + # Adjust cudagraph sizes to be a multiple of uniform_decode_query_len + # to avoid: https://github.com/vllm-project/vllm/issues/28207 and temp-fix: + # https://github.com/vllm-project/vllm/issues/28207#issuecomment-3504004536 + # Will be removed in the near future when we have separate cudagraph capture + # sizes for decode and mixed prefill-decode. + if ( + cudagraph_mode.decode_mode() == CUDAGraphMode.FULL + and uniform_decode_query_len > 1 + ): + self.adjust_cudagraph_sizes_for_spec_decode( + uniform_decode_query_len, + tensor_parallel_size, + ) + + # For Mamba models with FULL decode cudagraphs, each decode + # sequence needs one Mamba cache block. The decode cudagraph + # dispatcher already caps batch sizes at max_num_seqs, so we just + # need to verify that enough blocks exist. Raising here instead + # of silently capping cudagraph_capture_sizes avoids unintended + # restrictions on PIECEWISE (prefill) cudagraphs. + # See: https://github.com/vllm-project/vllm/issues/34094 + if ( + kv_cache_config is not None + and max_num_reqs is not None + and cudagraph_mode.has_full_cudagraphs() + and not is_profiling + and kv_cache_config.has_mamba_layers + and max_num_reqs > kv_cache_config.num_blocks + ): + raise ValueError( + f"max_num_seqs ({max_num_reqs}) exceeds available Mamba cache " + f"blocks ({kv_cache_config.num_blocks}). Each decode sequence " + "requires one Mamba cache block, so CUDA graph capture cannot " + "proceed. Please lower max_num_seqs to at most " + f"{kv_cache_config.num_blocks} or increase " + "gpu_memory_utilization." + ) + + self.cudagraph_mode = cudagraph_mode + return cudagraph_mode + def adjust_cudagraph_sizes_for_spec_decode( self, uniform_decode_query_len: int, tensor_parallel_size: int ): diff --git a/vllm/config/pooler.py b/vllm/config/pooler.py index 24368c3494e..cfe9eee49d6 100644 --- a/vllm/config/pooler.py +++ b/vllm/config/pooler.py @@ -83,6 +83,13 @@ class PoolerConfig: If provided, apply classification logit biases. Defaults to None. """ + logit_scale: float | None = None + """ + If provided, scale the classification logits by this factor before + activation. Combined with logit_bias, enables affine score calibration: + activation(logit_scale * (score - logit_bias)). Defaults to None. + """ + ## for reward models step_tag_id: int | None = None """ diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 8f9b9dcc811..61f311c8429 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -825,6 +825,7 @@ class SpeculativeConfig: "kimi_k2", "kimi_k25", "minimax_m2", + "gemma4", ] if ( self.method in ("eagle3", "extract_hidden_states", "dflash") diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 3b8431e9530..348ea40884a 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1627,6 +1627,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) diff --git a/vllm/entrypoints/pooling/embed/io_processor.py b/vllm/entrypoints/pooling/embed/io_processor.py index 623fee4fd3b..d7ee6ae5813 100644 --- a/vllm/entrypoints/pooling/embed/io_processor.py +++ b/vllm/entrypoints/pooling/embed/io_processor.py @@ -24,7 +24,13 @@ from vllm.entrypoints.pooling.embed.protocol import ( EmbeddingChatRequest, EmbeddingCompletionRequest, ) -from vllm.entrypoints.pooling.typing import PoolingServeContext +from vllm.entrypoints.pooling.scoring.io_processor import JinaRankingIOProcessorMixin +from vllm.entrypoints.pooling.typing import ( + OfflineInputsContext, + PoolingChatLikeRequest, + PoolingCompletionLikeRequest, + PoolingServeContext, +) from vllm.inputs import EngineInput, tokens_input from vllm.logger import init_logger from vllm.outputs import PoolingOutput, PoolingRequestOutput @@ -553,3 +559,48 @@ class EmbedIOProcessor(PoolingIOProcessor): class TokenEmbedIOProcessor(PoolingIOProcessor): name = "token_embed" + + +class JinaRankingTokenEmbedIOProcessor( + TokenEmbedIOProcessor, JinaRankingIOProcessorMixin +): + def pre_process_online(self, ctx: PoolingServeContext): + request = ctx.request + if isinstance(request, PoolingCompletionLikeRequest): + prompts = request.input + if not isinstance(prompts, Sequence) or len(prompts) < 2: + raise ValueError("The JinaForRanking model requires at least 2 inputs.") + + text_prompts = self.ensure_str(prompts) + + # The JinaForRanking model concatenates docs first, then query. + # Let's stay consistent with this novel design. + prompt_input = self.format_docs_prompts_func( + query=text_prompts[-1], docs=text_prompts[:-1] + ) + + engine_inputs = self._preprocess_completion_online( + request, + prompt_input=prompt_input, + prompt_embeds=None, + ) + elif isinstance(request, PoolingChatLikeRequest): + raise ValueError("The JinaForRanking does not support chat Request.") + else: + raise ValueError(f"Invalid {self.name} request type") + + ctx.engine_inputs = engine_inputs + + def pre_process_offline(self, ctx: OfflineInputsContext) -> Sequence[EngineInput]: + if not isinstance(ctx.prompts, Sequence) or len(ctx.prompts) < 2: + raise ValueError("The JinaForRanking model requires at least 2 inputs.") + + text_prompts = self.ensure_str(ctx.prompts) + + # The JinaForRanking model concatenates docs first, then query. + # Let's stay consistent with this novel design. + ctx.prompts = self.format_docs_prompts_func( + query=text_prompts[-1], docs=text_prompts[:-1] + ) + + return super().pre_process_offline(ctx) diff --git a/vllm/entrypoints/pooling/io_processor_factories.py b/vllm/entrypoints/pooling/io_processor_factories.py index de60a746a35..5e67d069da7 100644 --- a/vllm/entrypoints/pooling/io_processor_factories.py +++ b/vllm/entrypoints/pooling/io_processor_factories.py @@ -59,6 +59,13 @@ def init_pooling_io_processors( if score_type is not None and score_type in ScoringIOProcessors: processors[score_type] = ScoringIOProcessors[score_type] + if model_config.architecture == "JinaForRanking": + from .embed.io_processor import JinaRankingTokenEmbedIOProcessor + from .scoring.io_processor import ScoringIOProcessors + + processors["token_embed"] = JinaRankingTokenEmbedIOProcessor + processors["late-interaction"] = ScoringIOProcessors["jina-reranking-scoring"] + return { task: processor_cls( vllm_config=vllm_config, diff --git a/vllm/entrypoints/pooling/scoring/io_processor.py b/vllm/entrypoints/pooling/scoring/io_processor.py index dd505c79cf1..994c3c453bf 100644 --- a/vllm/entrypoints/pooling/scoring/io_processor.py +++ b/vllm/entrypoints/pooling/scoring/io_processor.py @@ -416,11 +416,137 @@ class CrossEncoderIOProcessor(ScoringIOProcessor): return full_prompt, engine_prompt +class JinaRankingIOProcessorMixin: + @staticmethod + def sanitize_input(text: str, special_tokens: dict[str, str]) -> str: + for token in special_tokens.values(): + text = text.replace(token, "") + return text + + @staticmethod + def format_docs_prompts_func( + query: str, + docs: list[str], + special_tokens: dict[str, str] | None = None, + instruction: str | None = None, + no_thinking: bool = True, + ) -> str: + # TODO: Try converting the code below into a chat template. + + default_special_tokens = { + "query_embed_token": "<|rerank_token|>", + "doc_embed_token": "<|embed_token|>", + } + if special_tokens is None: + special_tokens = default_special_tokens + + query = JinaRankingIOProcessorMixin.sanitize_input(query, special_tokens) + docs = [ + JinaRankingIOProcessorMixin.sanitize_input(doc, special_tokens) + for doc in docs + ] + + prefix = ( + "<|im_start|>system\n" + "You are a search relevance expert who can determine a ranking of the passages based on how relevant they are to the query. " # noqa: E501 + "If the query is a question, how relevant a passage is depends on how well it answers the question. " # noqa: E501 + "If not, try to analyze the intent of the query and assess how well each passage satisfies the intent. " # noqa: E501 + "If an instruction is provided, you should follow the instruction when determining the ranking." # noqa: E501 + "<|im_end|>\n<|im_start|>user\n" + ) + suffix = "<|im_end|>\n<|im_start|>assistant\n" + if no_thinking: + suffix += "\n\n\n\n" + + doc_emb_token = special_tokens["doc_embed_token"] + query_emb_token = special_tokens["query_embed_token"] + + prompt = ( + f"I will provide you with {len(docs)} passages, each indicated by a numerical identifier. " # noqa: E501 + f"Rank the passages based on their relevance to query: {query}\n" + ) + + if instruction: + prompt += f"\n{instruction}\n\n" + + doc_prompts = [ + f'\n{doc}{doc_emb_token}\n' + for i, doc in enumerate(docs) + ] + prompt += "\n".join(doc_prompts) + "\n" + prompt += f"\n{query}{query_emb_token}\n" + + return prefix + prompt + suffix + + @staticmethod + def ensure_str(data: Sequence[Any]) -> list[str]: + text: list[str] = [] + for prompt in data: + if not isinstance(prompt, str): + raise ValueError( + "The JinaForRanking model only supports text as input." + ) + text.append(prompt) + return text + + +class JinaRankingIOProcessor(LateInteractionIOProcessor, JinaRankingIOProcessorMixin): + name = "jina-reranking-scoring" + pooling_task: PoolingTask = "token_embed" + + def _pre_process( + self, + scoring_data: ScoringData, + tok_params: TokenizeParams, + prompt_extras: dict[str, Any] | None = None, + ) -> Sequence[EngineInput]: + queries = self.ensure_str(scoring_data.data_1) + docs = self.ensure_str(scoring_data.data_2) + + if len(queries) == 1: + prompts = [self.format_docs_prompts_func(query=queries[0], docs=docs)] + else: + prompts = [ + self.format_docs_prompts_func(query=q, docs=[d]) + for q, d in zip(queries, docs) + ] + + return self._preprocess_completion_offline( + prompts=prompts, tok_params=tok_params, prompt_extras=prompt_extras + ) + + def _post_process(self, outputs: list[PoolingRequestOutput], n_queries: int): + final_res_batch: list[PoolingRequestOutput] = [] + + for i in range(len(outputs)): + embeds = outputs[i].outputs.data.float() + + # The JinaForRanking model concatenates docs first, then query. + # Let's stay consistent with this novel design. + query_embeds = embeds[-1] + doc_embeds = embeds[:-1] + + scores = F.cosine_similarity(query_embeds, doc_embeds) + + for score in scores: + final_res_batch.append( + PoolingRequestOutput( + request_id=outputs[i].request_id, + outputs=score, + prompt_token_ids=outputs[i].prompt_token_ids, + num_cached_tokens=outputs[i].num_cached_tokens, + finished=True, + ) + ) + return final_res_batch + + ScoringIOProcessors: dict[str, type[ScoringIOProcessor]] = { p.name: p for p in [ BiEncoderIOProcessor, LateInteractionIOProcessor, + JinaRankingIOProcessor, FlashLateInteractionIOProcessor, CrossEncoderIOProcessor, ] diff --git a/vllm/entrypoints/pooling/scoring/serving.py b/vllm/entrypoints/pooling/scoring/serving.py index fc52075611e..571fd981f3b 100644 --- a/vllm/entrypoints/pooling/scoring/serving.py +++ b/vllm/entrypoints/pooling/scoring/serving.py @@ -4,7 +4,6 @@ from fastapi.responses import JSONResponse, Response from vllm import PoolingParams -from vllm.config import VllmConfig from vllm.engine.protocol import EngineClient from vllm.entrypoints.openai.engine.protocol import UsageInfo from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor @@ -42,25 +41,23 @@ class ServingScores(PoolingServing): enable_flash_late_interaction: bool = True, **kwargs, ): - self.score_type = engine_client.model_config.score_type + self.io_processor_name: str = engine_client.model_config.score_type self.enable_flash_late_interaction = ( - self.score_type == "late-interaction" and enable_flash_late_interaction + self.io_processor_name == "late-interaction" + and enable_flash_late_interaction ) + if self.enable_flash_late_interaction: + self.io_processor_name = "flash-late-interaction" + + if engine_client.model_config.architecture == "JinaForRanking": + self.io_processor_name = "jina-reranking-scoring" + self.enable_flash_late_interaction = False + super().__init__(engine_client, *args, **kwargs) - def init_io_processor( - self, vllm_config: VllmConfig, *args, **kwargs - ) -> PoolingIOProcessor: - model_config = vllm_config.model_config - - score_type: str = model_config.score_type - if self.enable_flash_late_interaction: - score_type = "flash-late-interaction" - - assert score_type in ScoringIOProcessors - processor_cls = ScoringIOProcessors[score_type] - return processor_cls(vllm_config, *args, **kwargs) + def init_io_processor(self, *args, **kwargs) -> PoolingIOProcessor: + return ScoringIOProcessors[self.io_processor_name](*args, **kwargs) async def __call__(self, *args, **kwargs) -> Response: if not self.enable_flash_late_interaction: diff --git a/vllm/model_executor/layers/fused_moe/configs/E=256,N=384,device_name=NVIDIA_RTX_PRO_6000_Blackwell_Server_Edition,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/fused_moe/configs/E=256,N=384,device_name=NVIDIA_RTX_PRO_6000_Blackwell_Server_Edition,dtype=fp8_w8a8,block_shape=[128,128].json new file mode 100644 index 00000000000..bcec61632e3 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/configs/E=256,N=384,device_name=NVIDIA_RTX_PRO_6000_Blackwell_Server_Edition,dtype=fp8_w8a8,block_shape=[128,128].json @@ -0,0 +1,147 @@ +{ + "triton_version": "3.6.0", + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/vllm/model_executor/layers/fused_moe/configs/E=256,N=512,device_name=NVIDIA_RTX_PRO_6000_Blackwell_Server_Edition,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/fused_moe/configs/E=256,N=512,device_name=NVIDIA_RTX_PRO_6000_Blackwell_Server_Edition,dtype=fp8_w8a8,block_shape=[128,128].json new file mode 100644 index 00000000000..705ca33d594 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/configs/E=256,N=512,device_name=NVIDIA_RTX_PRO_6000_Blackwell_Server_Edition,dtype=fp8_w8a8,block_shape=[128,128].json @@ -0,0 +1,147 @@ +{ + "triton_version": "3.6.0", + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 2 + }, + "8": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 2 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/vllm/model_executor/layers/fused_moe/configs/E=64,N=1536,device_name=NVIDIA_RTX_PRO_6000_Blackwell_Server_Edition,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/fused_moe/configs/E=64,N=1536,device_name=NVIDIA_RTX_PRO_6000_Blackwell_Server_Edition,dtype=fp8_w8a8,block_shape=[128,128].json new file mode 100644 index 00000000000..9c2ebaddd83 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/configs/E=64,N=1536,device_name=NVIDIA_RTX_PRO_6000_Blackwell_Server_Edition,dtype=fp8_w8a8,block_shape=[128,128].json @@ -0,0 +1,147 @@ +{ + "triton_version": "3.6.0", + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 2 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/vllm/model_executor/layers/mamba/lamport_workspace.py b/vllm/model_executor/layers/mamba/lamport_workspace.py new file mode 100644 index 00000000000..afae19c75fe --- /dev/null +++ b/vllm/model_executor/layers/mamba/lamport_workspace.py @@ -0,0 +1,302 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + + +import array +import contextlib +import struct +import sys +import threading + +import torch + +try: + from cuda.bindings import runtime as cudart +except ImportError: + from cuda import cudart + +_ALIGN = 1 << 21 # 2 MiB — CUDA IPC allocation alignment + + +# --------------------------------------------------------------------------- +# CUDA helpers +# --------------------------------------------------------------------------- + + +def _check(error): + """Raise on CUDA runtime error.""" + success = getattr(cudart.cudaError_t, "cudaSuccess", None) or cudart.cudaError_t(0) + if error != success: + raise RuntimeError(f"CUDA runtime error: {error}") + + +def _cuda_malloc(size: int): + aligned = ((size + _ALIGN - 1) >> 21) << 21 + err, ptr = cudart.cudaMalloc(aligned) + _check(err) + return ptr, aligned + + +def _cuda_free(ptr: int): + if ptr: + _check(cudart.cudaFree(ptr)[0]) + + +def _cuda_memset_zero(ptr: int, size: int): + _check(cudart.cudaMemset(ptr, 0, size)[0]) + + +def _cuda_memcpy_d2d(dst: int, src: int, size: int): + _check( + cudart.cudaMemcpy( + dst, src, size, cudart.cudaMemcpyKind.cudaMemcpyDeviceToDevice + )[0] + ) + + +# --------------------------------------------------------------------------- +# IPC buffer +# --------------------------------------------------------------------------- + + +class IpcBuffer: + """ + Allocates CUDA device memory and exchanges IPC handles with all ranks + so that every rank holds a valid device pointer to every other rank's buffer. + """ + + def __init__(self, rank: int, world_size: int, size: int, process_group=None): + self.rank = rank + self.world_size = world_size + self.peer_ptrs: list[int] = [0] * world_size + self.local_ptr: int = 0 + self._alive = False + + if size <= 0: + return + + self.local_ptr, _ = _cuda_malloc(size) + _cuda_memset_zero(self.local_ptr, size) + self._alive = True + + # --- exchange IPC handles via torch.distributed --- + err, local_handle = cudart.cudaIpcGetMemHandle(self.local_ptr) + _check(err) + + all_handles: list[bytes | None] = [None] * world_size + torch.distributed.all_gather_object( + all_handles, bytes(local_handle.reserved), group=process_group + ) + + for r in range(world_size): + if r == rank: + self.peer_ptrs[r] = self.local_ptr + else: + handle = cudart.cudaIpcMemHandle_t() + handle.reserved = all_handles[r] + err, ptr = cudart.cudaIpcOpenMemHandle( + handle, cudart.cudaIpcMemLazyEnablePeerAccess + ) + _check(err) + self.peer_ptrs[r] = ptr + + def serialize(self) -> list[int]: + """Return peer pointers as a list of int64 values (one per rank).""" + raw = b"" + for ptr in self.peer_ptrs: + raw += struct.pack("P", ptr) + return array.array("Q", raw).tolist() + + def cleanup(self): + if not self._alive: + return + self._alive = False + for r in range(self.world_size): + if self.peer_ptrs[r] == 0: + continue + if r == self.rank: + _cuda_free(self.peer_ptrs[r]) + else: + with contextlib.suppress(RuntimeError): + _check(cudart.cudaIpcCloseMemHandle(self.peer_ptrs[r])[0]) + self.peer_ptrs[r] = 0 + self.local_ptr = 0 + + def __del__(self): + if not sys.is_finalizing(): + self.cleanup() + + +# --------------------------------------------------------------------------- +# Lamport negative-zero initialization +# --------------------------------------------------------------------------- + + +def _lamport_fill_neg_zero(device_ptr: int, size_bytes: int): + """ + Fill device memory with IEEE-754 negative zero (-0.0f = 0x80000000). + This is the "slot empty" sentinel for the Lamport protocol: the kernel + spin-waits until a value is *not* negative zero. + """ + if size_bytes == 0 or device_ptr == 0: + return + n_floats = size_bytes // 4 + # torch preserves -0.0 in IEEE-754 + fill = torch.full((n_floats,), -0.0, dtype=torch.float32, device="cuda") + _cuda_memcpy_d2d(device_ptr, fill.data_ptr(), size_bytes) + del fill + + +# --------------------------------------------------------------------------- +# LamportWorkspace — the main class +# --------------------------------------------------------------------------- + + +class LamportWorkspace: + """ + Self-contained workspace for Lamport-based cross-GPU AllReduce. + + Parameters + ---------- + rank : int + Local rank (0-based). + world_size : int + Total number of ranks in the TP group. + comm_size : int + Size in bytes of *one* Lamport buffer slot. The total IPC allocation + per rank is ``3 * comm_size`` (triple-buffering). Must be large enough + to hold the per-slot data written by the kernel. Use + ``compute_comm_size_for_minimax()`` for a safe default. + process_group : optional + ``torch.distributed`` process group for IPC handle exchange. + ``None`` uses the default group. + """ + + def __init__(self, rank: int, world_size: int, comm_size: int, process_group=None): + assert world_size >= 2, "Lamport workspace requires at least 2 ranks" + assert comm_size > 0, "comm_size must be positive" + + self.rank = rank + self.world_size = world_size + self.comm_size = comm_size + + # 1) Lamport triple-buffer (the only IPC memory the kernel reads/writes) + lamport_total = 3 * comm_size + self._lamport = IpcBuffer(rank, world_size, lamport_total, process_group) + _lamport_fill_neg_zero(self._lamport.local_ptr, lamport_total) + + # 2) flag_buffer on device: int32[3] = {counter, unused, lamport_flag} + # counter — used for block-level sync inside the kernel + # unused — reserved (index 1) + # lamport_flag — triple-buffer rotation index (0 → 1 → 2 → 0 …) + self._flag_buf = torch.zeros(3, dtype=torch.int32, device="cuda") + + # 3) layout_buffer on device: int64[2] = {clear_size, comm_size} + # clear_size — bytes to clear from *previous* slot (set by kernel) + # comm_size — size of one triple-buffer slot + self._layout_buf = torch.tensor( + [0, comm_size], dtype=torch.int64, device="cuda" + ) + + # 4) Assemble device-side void* pointer array + N = world_size + ptrs: list[int] = [] + ptrs += [0] * N # [0 .. N-1] ipc_buffers (placeholder) + ptrs += [0] * N # [N .. 2N-1] ipc_barriers (placeholder) + ptrs += self._lamport.serialize() # [2N .. 3N-1] lamport peer ptrs + ptrs.append(self._flag_buf.data_ptr()) # [3N] flag_buffer + ptrs.append(self._layout_buf.data_ptr()) # [3N+1] layout_buffer + + self._workspace = torch.tensor(ptrs, dtype=torch.int64, device="cuda") + + @property + def workspace(self) -> torch.Tensor: + """Device tensor (int64) that can be passed to the kernel + as ``void** workspace``.""" + return self._workspace + + # ------------------------------------------------------------------ + # Helpers + # ------------------------------------------------------------------ + + @staticmethod + def compute_comm_size_for_minimax( + max_tokens: int, + world_size: int, + fused_qk: bool = True, + ) -> int: + """ + Return a safe ``comm_size`` (in bytes) for MiniMaxReduceRMSKernel. + + The kernel stores per-token variance scalars in the Lamport buffer: + - single-matrix path: ``world_size × max_tokens × 4`` bytes per slot + - fused Q+K path: ``world_size × 2 × ceil(max_tokens/4) × 16`` bytes per slot + + The returned value is rounded up to 2 MiB alignment. + """ + if fused_qk: + groups = (max_tokens + 3) // 4 + slot_bytes = world_size * 2 * groups * 16 # 16 = sizeof(float4) + else: + slot_bytes = world_size * max_tokens * 4 # 4 = sizeof(float) + return ((slot_bytes + _ALIGN - 1) >> 21) << 21 + + def cleanup(self): + if hasattr(self, "_lamport"): + self._lamport.cleanup() + + def __del__(self): + if not sys.is_finalizing(): + self.cleanup() + + def __repr__(self): + return ( + f"LamportWorkspace(rank={self.rank}, world_size={self.world_size}, " + f"comm_size={self.comm_size})" + ) + + +# --------------------------------------------------------------------------- +# Cached convenience function (mirrors TRT-LLM's get_allreduce_workspace) +# --------------------------------------------------------------------------- + +_cache_lock = threading.Lock() +_workspace_cache: dict = {} + + +def get_allreduce_workspace( + rank: int, + world_size: int, + comm_size: int | None = None, + max_tokens: int = 16384, + process_group=None, +) -> torch.Tensor: + """ + Return a cached workspace tensor for the given (rank, world_size) pair. + + On first call the workspace is allocated and IPC handles are exchanged; + subsequent calls with the same arguments return the cached tensor. + + Parameters + ---------- + rank, world_size : int + TP rank and TP size. + comm_size : int, optional + Explicit slot size in bytes. If ``None``, computed automatically + from ``max_tokens`` and ``world_size`` (fused Q+K path). + max_tokens : int + Maximum number of tokens per batch (used when ``comm_size is None``). + process_group : optional + ``torch.distributed`` process group. + """ + if comm_size is None: + comm_size = LamportWorkspace.compute_comm_size_for_minimax( + max_tokens, world_size, fused_qk=True + ) + pg_id = id(process_group) if process_group is not None else 0 + key = (rank, world_size, comm_size, pg_id) + with _cache_lock: + if key not in _workspace_cache: + ws = LamportWorkspace(rank, world_size, comm_size, process_group) + _workspace_cache[key] = ws + return _workspace_cache[key].workspace diff --git a/vllm/model_executor/layers/pooler/seqwise/heads.py b/vllm/model_executor/layers/pooler/seqwise/heads.py index 31a96122392..4c0c16c3475 100644 --- a/vllm/model_executor/layers/pooler/seqwise/heads.py +++ b/vllm/model_executor/layers/pooler/seqwise/heads.py @@ -104,6 +104,7 @@ class ClassifierPoolerHead(SequencePoolerHead): self, classifier: ClassifierFn | None = None, logit_bias: float | None = None, + logit_scale: float | None = None, head_dtype: torch.dtype | str | None = None, activation: ActivationFn | None = None, ) -> None: @@ -111,6 +112,7 @@ class ClassifierPoolerHead(SequencePoolerHead): self.classifier = classifier self.logit_bias = logit_bias + self.logit_scale = logit_scale self.head_dtype = head_dtype self.activation = activation @@ -140,6 +142,8 @@ class ClassifierPoolerHead(SequencePoolerHead): # logits shape: [batchsize, num_labels] if self.logit_bias is not None: logits -= self.logit_bias + if self.logit_scale is not None: + logits *= self.logit_scale if self.activation is not None: flags = [p.use_activation for p in pooling_params] diff --git a/vllm/model_executor/layers/pooler/seqwise/poolers.py b/vllm/model_executor/layers/pooler/seqwise/poolers.py index f46834a7c3f..ba851fed645 100644 --- a/vllm/model_executor/layers/pooler/seqwise/poolers.py +++ b/vllm/model_executor/layers/pooler/seqwise/poolers.py @@ -119,6 +119,7 @@ def pooler_for_classify( head_dtype=model_config.head_dtype, classifier=classifier, logit_bias=model_config.pooler_config.logit_bias, + logit_scale=model_config.pooler_config.logit_scale, activation=resolve_classifier_act_fn( model_config, static_num_labels=True, act_fn=act_fn ), diff --git a/vllm/model_executor/layers/pooler/tokwise/heads.py b/vllm/model_executor/layers/pooler/tokwise/heads.py index bd676f09a07..0d8d0926ad9 100644 --- a/vllm/model_executor/layers/pooler/tokwise/heads.py +++ b/vllm/model_executor/layers/pooler/tokwise/heads.py @@ -138,6 +138,7 @@ class TokenClassifierPoolerHead(TokenPoolerHead): self, classifier: ClassifierFn | None = None, logit_bias: float | None = None, + logit_scale: float | None = None, head_dtype: torch.dtype | str | None = None, activation: ActivationFn | None = None, ) -> None: @@ -145,6 +146,7 @@ class TokenClassifierPoolerHead(TokenPoolerHead): self.classifier = classifier self.logit_bias = logit_bias + self.logit_scale = logit_scale self.head_dtype = head_dtype self.activation = activation @@ -172,6 +174,8 @@ class TokenClassifierPoolerHead(TokenPoolerHead): if self.logit_bias is not None: logits -= self.logit_bias + if self.logit_scale is not None: + logits *= self.logit_scale if self.activation is not None and pooling_param.use_activation: logits = self.activation(logits) diff --git a/vllm/model_executor/layers/pooler/tokwise/methods.py b/vllm/model_executor/layers/pooler/tokwise/methods.py index 27f622e9d71..b3f3fe07880 100644 --- a/vllm/model_executor/layers/pooler/tokwise/methods.py +++ b/vllm/model_executor/layers/pooler/tokwise/methods.py @@ -162,7 +162,7 @@ class StepPool(AllPool): ): # for unfinished chunked prefill if data is None: - pass + pooled_data.append(None) else: step_tag_id = pooling_param.step_tag_id returned_token_ids = pooling_param.returned_token_ids diff --git a/vllm/model_executor/layers/pooler/tokwise/poolers.py b/vllm/model_executor/layers/pooler/tokwise/poolers.py index 6a477f75d5c..5a62c149f34 100644 --- a/vllm/model_executor/layers/pooler/tokwise/poolers.py +++ b/vllm/model_executor/layers/pooler/tokwise/poolers.py @@ -60,7 +60,7 @@ class TokenPooler(Pooler): def __init__( self, pooling: TokenPoolingMethod | TokenPoolingFn, - head: TokenPoolerHead | TokenPoolingHeadFn, + head: TokenPoolerHead | TokenPoolingHeadFn | None = None, ) -> None: super().__init__() @@ -96,9 +96,9 @@ class TokenPooler(Pooler): return self.head.forward_ragged( pooled_data, pooling_metadata.pooling_params ) - else: - pooled_data = pooled_data.split() - pooled_data = self.head(pooled_data, pooling_metadata) + pooled_data = pooled_data.split() + if self.head is not None: + pooled_data = self.head(pooled_data, pooling_metadata) return pooled_data @@ -136,6 +136,7 @@ def pooler_for_token_classify( head_dtype=model_config.head_dtype, classifier=classifier, logit_bias=model_config.pooler_config.logit_bias, + logit_scale=model_config.pooler_config.logit_scale, activation=resolve_classifier_act_fn( model_config, static_num_labels=False, act_fn=act_fn ), diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index 22d300a7ebf..6e511f4f0a3 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -192,6 +192,12 @@ class JambaForSequenceClassificationConfig(VerifyAndUpdateConfig): pooler_config.use_activation = False +class JinaForRankingConfig(VerifyAndUpdateConfig): + @staticmethod + def verify_and_update_model_config(model_config: "ModelConfig") -> None: + model_config.hf_config.embedding_size = 512 + + class JinaRobertaModelConfig(VerifyAndUpdateConfig): @staticmethod def verify_and_update_model_config(model_config: "ModelConfig") -> None: @@ -612,6 +618,7 @@ MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = { "GteNewForSequenceClassification": GteNewModelConfig, "GteNewModel": GteNewModelConfig, "JambaForSequenceClassification": JambaForSequenceClassificationConfig, + "JinaForRanking": JinaForRankingConfig, "JinaVLForRanking": JinaVLForSequenceClassificationConfig, "LlamaBidirectionalForSequenceClassification": LlamaBidirectionalConfig, "LlamaBidirectionalModel": LlamaBidirectionalConfig, diff --git a/vllm/model_executor/models/gemma4.py b/vllm/model_executor/models/gemma4.py index c41c0de52b0..debf8ccded4 100644 --- a/vllm/model_executor/models/gemma4.py +++ b/vllm/model_executor/models/gemma4.py @@ -60,7 +60,13 @@ from vllm.model_executor.model_loader.weight_utils import ( 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, extract_layer_index, @@ -838,7 +844,7 @@ class Gemma4CrossDecoderLayers(nn.Module): @support_torch_compile( enable_if=lambda vllm_config: not vllm_config.cache_config.kv_sharing_fast_prefill ) -class Gemma4Model(nn.Module): +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) @@ -1168,7 +1174,7 @@ 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, @@ -1204,6 +1210,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) ): @@ -1222,6 +1229,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( { @@ -1236,6 +1246,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]: @@ -1381,7 +1394,9 @@ class Gemma4Model(nn.Module): return loaded_params -class Gemma4ForCausalLM(nn.Module, SupportsLoRA, SupportsPP, MixtureOfExperts): +class Gemma4ForCausalLM( + nn.Module, SupportsLoRA, SupportsPP, MixtureOfExperts, SupportsEagle3 +): # 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. @@ -1463,7 +1478,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 ) diff --git a/vllm/model_executor/models/gemma4_mm.py b/vllm/model_executor/models/gemma4_mm.py index fa597fe96a0..e22f23c5c8b 100644 --- a/vllm/model_executor/models/gemma4_mm.py +++ b/vllm/model_executor/models/gemma4_mm.py @@ -64,7 +64,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, @@ -845,7 +850,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", diff --git a/vllm/model_executor/models/jina.py b/vllm/model_executor/models/jina.py new file mode 100644 index 00000000000..980502191dd --- /dev/null +++ b/vllm/model_executor/models/jina.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from https://huggingface.co/jinaai/jina-reranker-v3/blob/main/modeling.py +from collections.abc import Iterable + +import torch +from torch import nn + +from vllm.config import VllmConfig +from vllm.sequence import IntermediateTensors +from vllm.tasks import PoolingTask +from vllm.v1.pool.metadata import PoolingMetadata + +from ..layers.pooler import DispatchPooler +from ..layers.pooler.tokwise import ( + StepPool, + TokenPooler, + TokenPoolingMethodOutputItem, +) +from .interfaces import SupportsLateInteraction +from .qwen3 import Qwen3Model +from .utils import AutoWeightsLoader, maybe_prefix + + +class JinaForRanking(nn.Module, SupportsLateInteraction): + is_pooling_model = True + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + + self.config = config + self.projector_dim: int = config.embedding_size + + self.vllm_config = vllm_config + self.quant_config = quant_config + self.model = Qwen3Model( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + + self.projector = nn.Sequential( + nn.Linear(config.hidden_size, config.hidden_size // 2, bias=False), + nn.ReLU(), + nn.Linear(config.hidden_size // 2, self.projector_dim, bias=False), + ) + + self.pooler = DispatchPooler( + { + "token_embed": TokenPooler( + pooling=JinaForRankingPool(self.projector), + ) + } + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | IntermediateTensors: + hidden_states = self.model( + input_ids, positions, intermediate_tensors, inputs_embeds + ) + return hidden_states + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader(self, skip_prefixes=(["lm_head."])) + return loader.load_weights(weights) + + +class JinaForRankingPool(StepPool): + def __init__(self, projector: nn.Sequential): + super().__init__() + + self.doc_token_id = 151670 + self.query_token_id = 151671 + self.projector = projector + + def get_supported_tasks(self) -> set[PoolingTask]: + return {"token_embed"} + + def forward( + self, + hidden_states: torch.Tensor, + pooling_metadata: PoolingMetadata, + ) -> list[TokenPoolingMethodOutputItem]: + pooled_data_lst = super().forward(hidden_states, pooling_metadata) + prompt_token_ids = pooling_metadata.get_prompt_token_ids() + + embeds_list = list[torch.Tensor | None]() + for data, token_ids in zip(pooled_data_lst, prompt_token_ids): + # for unfinished chunked prefill + if data is None: + embeds_list.append(None) + else: + docs_indexes = torch.where(torch.eq(token_ids, self.doc_token_id))[0] + query_indexes = torch.where(torch.eq(token_ids, self.query_token_id))[0] + + # The JinaForRanking model concatenates docs first, then query. + # Let's stay consistent with this novel design. + indexes = torch.cat([docs_indexes, query_indexes]) + embeds = self.projector(data[indexes]) + embeds_list.append(embeds) + + return embeds_list diff --git a/vllm/model_executor/models/minimax_m2.py b/vllm/model_executor/models/minimax_m2.py index f10452c5738..2a9e5f4e08a 100644 --- a/vllm/model_executor/models/minimax_m2.py +++ b/vllm/model_executor/models/minimax_m2.py @@ -233,9 +233,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) diff --git a/vllm/model_executor/models/qwen3_asr.py b/vllm/model_executor/models/qwen3_asr.py index 3015ae03132..1d8e1bbaa4c 100644 --- a/vllm/model_executor/models/qwen3_asr.py +++ b/vllm/model_executor/models/qwen3_asr.py @@ -37,6 +37,7 @@ from vllm.inputs import ModalityData, MultiModalDataDict, PromptType, TokensProm from vllm.logger import init_logger from vllm.model_executor.models.interfaces import ( MultiModalEmbeddings, + SupportsLoRA, SupportsMRoPE, SupportsMultiModal, SupportsPP, @@ -266,7 +267,21 @@ class Qwen3ASRForConditionalGeneration( SupportsPP, SupportsMRoPE, SupportsTranscription, + SupportsLoRA, ): + # LoRA support + packed_modules_mapping = { + "qkv_proj": [ + "q_proj", + "k_proj", + "v_proj", + ], + "gate_up_proj": [ + "gate_proj", + "up_proj", + ], + } + supported_languages = ISO639_1_SUPPORTED_LANGS hf_to_vllm_mapper = WeightsMapper( @@ -513,6 +528,17 @@ class Qwen3ASRForConditionalGeneration( tower_model=["audio_tower."], ) + def get_num_mm_encoder_tokens(self, num_audio_tokens: int) -> int: + """Return the number of tokens processed by the audio tower encoder. + + Required for LoRA support on the tower module. + """ + # For Qwen3-ASR, the audio tower produces one embedding per audio + # placeholder token inserted into the prompt (no additional + # merge/downsample step like vision towers). Therefore, the encoder + # token budget is identity. + return num_audio_tokens + @classmethod def get_speech_to_text_config( cls, model_config: ModelConfig, task_type: str diff --git a/vllm/model_executor/models/qwen3_omni_moe_thinker.py b/vllm/model_executor/models/qwen3_omni_moe_thinker.py index 5e113dac5cc..8cee51a1269 100755 --- a/vllm/model_executor/models/qwen3_omni_moe_thinker.py +++ b/vllm/model_executor/models/qwen3_omni_moe_thinker.py @@ -57,6 +57,7 @@ from vllm.model_executor.layers.conv import Conv3dLayer from vllm.model_executor.layers.linear import ( ColumnParallelLinear, QKVParallelLinear, + ReplicatedLinear, RowParallelLinear, ) from vllm.model_executor.layers.logits_processor import LogitsProcessor @@ -357,7 +358,13 @@ class Qwen3OmniMoeAudioEncoder(nn.Module): conv_out_dim = config.downsample_hidden_size * ( (((config.num_mel_bins + 1) // 2 + 1) // 2 + 1) // 2 ) - self.conv_out = nn.Linear(conv_out_dim, config.d_model, bias=False) + self.conv_out = ReplicatedLinear( + conv_out_dim, + config.d_model, + bias=False, + return_bias=False, + prefix=f"{prefix}.conv_out", + ) # Transformer encoder layers self.layers = nn.ModuleList( @@ -372,9 +379,21 @@ class Qwen3OmniMoeAudioEncoder(nn.Module): # Output layers self.ln_post = nn.LayerNorm(config.d_model) - self.proj1 = nn.Linear(config.d_model, config.d_model) + self.proj1 = ReplicatedLinear( + config.d_model, + config.d_model, + bias=True, + return_bias=False, + prefix=f"{prefix}.proj1", + ) self.act = _ACTIVATION_REGISTRY[config.activation_function] - self.proj2 = nn.Linear(config.d_model, config.output_dim) + self.proj2 = ReplicatedLinear( + config.d_model, + config.output_dim, + bias=True, + return_bias=False, + prefix=f"{prefix}.proj2", + ) # Get attention backend self.attn_backend = get_vit_attn_backend( diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 66ded7e554b..948618d82a0 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -273,6 +273,7 @@ _LATE_INTERACTION_MODELS = { "ColBERTModernBertModel": ("colbert", "ColBERTModernBertModel"), "ColBERTJinaRobertaModel": ("colbert", "ColBERTJinaRobertaModel"), "ColBERTLfm2Model": ("colbert", "ColBERTLfm2Model"), + "JinaForRanking": ("jina", "JinaForRanking"), # [Multimodal] "ColModernVBertForRetrieval": ("colmodernvbert", "ColModernVBertForRetrieval"), "ColPaliForRetrieval": ("colpali", "ColPaliModel"), diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index 3543bdabb8c..4adad61c2fa 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -63,7 +63,11 @@ from vllm.v1.attention.backends.utils import ( from vllm.v1.attention.ops.common import cp_lse_ag_out_rs from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce from vllm.v1.attention.ops.merge_attn_states import merge_attn_states -from vllm.v1.kv_cache_interface import AttentionSpec, UniformTypeKVCacheSpecs +from vllm.v1.kv_cache_interface import ( + AttentionSpec, + KVQuantMode, + UniformTypeKVCacheSpecs, +) from vllm.v1.utils import CpuGpuBuffer FLASHINFER_WORKSPACE_BUFFER_SIZE_BATCH_INVARIANT = 2048 * 1024 * 1024 @@ -600,12 +604,15 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]): self.head_dim = self.kv_cache_spec.head_size self.page_size = self.kv_cache_spec.block_size - self.cache_dtype = self.cache_config.cache_dtype - if is_quantized_kv_cache(self.cache_dtype): + if self.kv_cache_spec.kv_quant_mode != KVQuantMode.NONE: + self.cache_dtype = self.cache_config.cache_dtype + # Cannot use self.kv_cache_spec.dtype here because kv_cache_spec + # storage dtype may not be the same as the op dtype (uint8 vs fp8_e4m3) self.kv_cache_dtype = FlashInferBackend.get_fp8_dtype_for_flashinfer( self.cache_dtype ) else: + self.cache_dtype = "auto" assert self.kv_cache_spec.dtype == self.model_config.dtype self.kv_cache_dtype = self.kv_cache_spec.dtype diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 62bdb8113a3..fa5395685e6 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -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: diff --git a/vllm/v1/spec_decode/eagle.py b/vllm/v1/spec_decode/eagle.py index 323b16931fb..f22e15b79f6 100644 --- a/vllm/v1/spec_decode/eagle.py +++ b/vllm/v1/spec_decode/eagle.py @@ -1329,6 +1329,7 @@ class SpecDecodeBaseProposer: "Qwen3_5MoeForConditionalGeneration", "Qwen3VLForConditionalGeneration", "Qwen3VLMoeForConditionalGeneration", + "Gemma4ForConditionalGeneration", ]: self.model.config.image_token_index = target_model.config.image_token_id elif self.get_model_name(target_model) == "PixtralForConditionalGeneration": diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index 34089a67b3b..ee6244c42a0 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Sequence +from dataclasses import dataclass from typing import Any, cast import numpy as np @@ -8,7 +9,11 @@ import torch from vllm.config import VllmConfig, get_layers_from_vllm_config from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase -from vllm.v1.attention.backend import AttentionBackend, CommonAttentionMetadata +from vllm.v1.attention.backend import ( + AttentionBackend, + AttentionCGSupport, + CommonAttentionMetadata, +) from vllm.v1.kv_cache_interface import ( AttentionSpec, KVCacheConfig, @@ -18,6 +23,12 @@ from vllm.v1.kv_cache_interface import ( from vllm.v1.worker.utils import AttentionGroup, bind_kv_cache +@dataclass(frozen=True) +class AttentionCGSupportInfo: + min_cg_support: AttentionCGSupport = AttentionCGSupport.ALWAYS + min_cg_attn_backend: str | None = None + + def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]: kv_cache_spec: dict[str, KVCacheSpec] = {} layer_type = cast(type[Any], AttentionLayerBase) @@ -34,10 +45,17 @@ def init_attn_backend( vllm_config: VllmConfig, device: torch.device, active_layer_names: set[str] | None = None, -): +) -> tuple[ + dict[str, type[AttentionBackend]], + list[list[AttentionGroup]], + AttentionCGSupportInfo, +]: attn_backends: dict[str, type[AttentionBackend]] = {} attn_groups: list[list[AttentionGroup]] = [] attn_backend_workspace: torch.Tensor | None = None + # Find minimum cudagraph support across all attention backends + min_cg_support = AttentionCGSupport.ALWAYS + min_cg_attn_backend = None for kv_cache_group_id, kv_cache_group_spec in enumerate( kv_cache_config.kv_cache_groups ): @@ -86,8 +104,24 @@ def init_attn_backend( else: if hasattr(builder, "set_workspace_buffer"): builder.set_workspace_buffer(attn_backend_workspace) + # Check cudagraph support for the attention backend + cg_support = builder.get_cudagraph_support( + vllm_config, + cast(AttentionSpec, kv_cache_group_spec.kv_cache_spec), + ) + if cg_support.value < min_cg_support.value: + min_cg_support = cg_support + min_cg_attn_backend = attn_backend.__name__ attn_groups.append(groups) - return attn_backends, attn_groups + + return ( + attn_backends, + attn_groups, + AttentionCGSupportInfo( + min_cg_support=min_cg_support, + min_cg_attn_backend=min_cg_attn_backend, + ), + ) def _allocate_kv_cache(kv_cache_config: KVCacheConfig, device: torch.device): @@ -110,7 +144,7 @@ def _allocate_kv_cache(kv_cache_config: KVCacheConfig, device: torch.device): def _reshape_kv_cache( kv_cache_config: KVCacheConfig, kv_cache_raw_tensors: dict[str, torch.Tensor], - attn_backends: dict[str, AttentionBackend], + attn_backends: dict[str, type[AttentionBackend]], cache_dtype: str, ) -> dict[str, torch.Tensor]: kv_caches: dict[str, torch.Tensor] = {} @@ -158,7 +192,7 @@ def init_kv_cache( runner_kv_caches: list[torch.Tensor], forward_context: dict[str, Any], kv_cache_config: KVCacheConfig, - attn_backends: dict[str, AttentionBackend], + attn_backends: dict[str, type[AttentionBackend]], device: torch.device, cache_dtype: str, ) -> dict[str, torch.Tensor]: diff --git a/vllm/v1/worker/gpu/dp_utils.py b/vllm/v1/worker/gpu/dp_utils.py index f0e2bfcf54b..09ac5b5af64 100644 --- a/vllm/v1/worker/gpu/dp_utils.py +++ b/vllm/v1/worker/gpu/dp_utils.py @@ -20,7 +20,7 @@ def make_num_tokens_across_dp(dp_size: int, num_tokens: int) -> torch.Tensor | N def sync_cudagraph_and_dp_padding( - cudagraph_manager: CudaGraphManager, + cudagraph_manager: CudaGraphManager | None, desired_batch_desc: BatchExecutionDescriptor, num_tokens: int, num_reqs: int, @@ -61,6 +61,10 @@ def sync_cudagraph_and_dp_padding( num_reqs=num_reqs, ), num_tokens_across_dp + assert cudagraph_manager is not None, ( + "cudagraph_manager should only be None during profile run, " + "where synced_cg_mode must be NONE across all DP ranks" + ) synced_num_tokens = int(num_tokens_across_dp.max().item()) synced_uniform_token_count = uniform_token_counts_across_dp[0] # If ranks disagree on the uniform token count, or its 0 (means None) set to None @@ -79,3 +83,41 @@ def sync_cudagraph_and_dp_padding( num_tokens_across_dp[:] = synced_desc.num_tokens return synced_desc, num_tokens_across_dp + + +def dispatch_cg_and_sync_dp( + cudagraph_manager: CudaGraphManager | None, + num_reqs: int, + num_tokens: int, + uniform_token_count: int | None, + dp_size: int, + dp_rank: int, + need_eager: bool = False, +) -> tuple[BatchExecutionDescriptor, torch.Tensor | None]: + if need_eager: + batch_desc = BatchExecutionDescriptor( + cg_mode=CUDAGraphMode.NONE, + num_tokens=num_tokens, + num_reqs=num_reqs, + ) + else: + assert cudagraph_manager is not None, ( + "cudagraph_manager should only be None during profile run, " + "where need_eager must be True" + ) + batch_desc = cudagraph_manager.dispatch( + num_reqs, num_tokens, uniform_token_count + ) + + if dp_size == 1: + return batch_desc, None + + return sync_cudagraph_and_dp_padding( + cudagraph_manager, + batch_desc, + num_tokens, + num_reqs, + uniform_token_count, + dp_size, + dp_rank, + ) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index f188b061a6c..ca52b8b66d0 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -61,7 +61,7 @@ from vllm.v1.worker.gpu.cudagraph_utils import ( ModelCudaGraphManager, get_uniform_token_count, ) -from vllm.v1.worker.gpu.dp_utils import sync_cudagraph_and_dp_padding +from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp from vllm.v1.worker.gpu.eplb_utils import EPLBController, step_eplb_after from vllm.v1.worker.gpu.input_batch import ( InputBatch, @@ -176,6 +176,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): # Draft tokens propagation - for spec-dec + struct outputs. self.draft_tokens_handler = DraftTokensHandler(self.device) + self.uniform_decode_query_len = 1 + self.num_speculative_steps # Pooling models. self.is_pooling_model = self.model_config.runner_type == "pooling" @@ -224,14 +225,9 @@ class GPUModelRunner(LoRAModelRunnerMixin): device=self.device, ) - # CUDA graphs. + # For CUDA graphs, and will init cudagraph_manager after init_attn_backend. self.decode_query_len = self.num_speculative_steps + 1 - self.cudagraph_manager = ModelCudaGraphManager( - self.vllm_config, - self.device, - self.compilation_config.cudagraph_mode, - decode_query_len=self.decode_query_len, - ) + self.cudagraph_manager: ModelCudaGraphManager | None = None # LoRA-related workers. self.lora_state = LoraState(max_num_reqs=self.max_num_reqs) # KV Connector if configured. @@ -361,9 +357,26 @@ class GPUModelRunner(LoRAModelRunnerMixin): cp_interleave=self.cp_interleave, ) - self.attn_backends, self.attn_groups = init_attn_backend( + self.attn_backends, self.attn_groups, attn_cg_support = init_attn_backend( self.kv_cache_config, self.vllm_config, self.device ) + cudagraph_mode = self.compilation_config.resolve_cudagraph_mode_and_sizes( + attn_cg_support.min_cg_support, + attn_cg_support.min_cg_attn_backend, + self.uniform_decode_query_len, + self.parallel_config.tensor_parallel_size, + self.kv_cache_config, + self.max_num_reqs, + ) + self.cudagraph_manager = ModelCudaGraphManager( + self.vllm_config, + self.device, + cudagraph_mode, + decode_query_len=self.decode_query_len, + ) + if self.speculator is not None: + self.speculator.init_cudagraph_manager(cudagraph_mode) + check_attention_cp_compatibility(self.vllm_config) if self.speculator is not None: # HACK(woosuk) @@ -437,6 +450,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): intermediate_tensors=intermediate_tensors, dummy_run=True, skip_attn_for_dummy_run=skip_attn, + is_profile=is_profile, ) self.kv_connector.set_disabled(False) @@ -486,6 +500,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): dummy_run=True, skip_attn_for_dummy_run=skip_attn, mm_inputs=mm_inputs, + is_profile=is_profile, ) assert hidden_states is not None # Last PP rank always has hidden_states @@ -547,6 +562,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): @torch.inference_mode() def capture_model(self) -> int: + assert self.cudagraph_manager is not None if not self.cudagraph_manager.needs_capture(): logger.warning( "Skipping CUDA graph capture. To turn on CUDA graph capture, " @@ -915,6 +931,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): intermediate_tensors: IntermediateTensors | None = None, dummy_run: bool = False, skip_attn_for_dummy_run: bool = False, + is_profile: bool = False, ) -> ModelRunnerOutput | IntermediateTensors | None: if not dummy_run: # Update the request states. @@ -934,34 +951,22 @@ class GPUModelRunner(LoRAModelRunnerMixin): max_query_len = max(scheduler_output.num_scheduled_tokens.values()) uniform_tok_count = get_uniform_token_count(num_reqs, num_toks, max_query_len) - batch_desc = self.cudagraph_manager.dispatch( - num_reqs, num_toks, uniform_tok_count - ) - num_tokens_across_dp = None - skip_compiled = False if self.is_encoder_decoder and scheduler_output.scheduled_encoder_inputs: # Encoder-decoder models such as Whisper should run eager/non-compiled # when encoder inputs are scheduled, because this step updates # cross-attention cache with dynamic encoder outputs. - # Override batch_desc to NONE. skip_compiled = True - batch_desc = BatchExecutionDescriptor( - cg_mode=CUDAGraphMode.NONE, - num_tokens=num_toks, - num_reqs=num_reqs, - ) - if self.dp_size > 1: - batch_desc, num_tokens_across_dp = sync_cudagraph_and_dp_padding( - self.cudagraph_manager, - batch_desc, - num_toks, - num_reqs, - uniform_tok_count, - self.dp_size, - self.dp_rank, - ) + batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp( + self.cudagraph_manager, + num_reqs, + num_toks, + uniform_tok_count, + self.dp_size, + self.dp_rank, + need_eager=is_profile or skip_compiled, + ) if batch_desc.num_tokens == 0: # All DP ranks have zero tokens to run. @@ -1059,6 +1064,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): # Use explicit cudagraph replay for FULL mode. # NOTE(woosuk): Here, we don't need to pass the input tensors, # because they are already copied to the CUDA graph input buffers. + assert self.cudagraph_manager is not None self.kv_connector.pre_forward(scheduler_output) model_output = self.cudagraph_manager.run_fullgraph(batch_desc) else: diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py index 887fd52794c..c311057dc28 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py @@ -19,10 +19,7 @@ from vllm.v1.worker.gpu.attn_utils import ( init_attn_backend, ) from vllm.v1.worker.gpu.block_table import BlockTables -from vllm.v1.worker.gpu.cudagraph_utils import ( - BatchExecutionDescriptor, -) -from vllm.v1.worker.gpu.dp_utils import sync_cudagraph_and_dp_padding +from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers from vllm.v1.worker.gpu.model_states.interface import ModelState from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample @@ -98,15 +95,19 @@ class EagleSpeculator: device=device, ) - # currently we don't support PIECEWISE for Eagle. - cudagraph_mode = vllm_config.compilation_config.cudagraph_mode + self.cudagraph_manager: EagleCudaGraphManager | None = None + + def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None: if cudagraph_mode.decode_mode() == CUDAGraphMode.FULL: cudagraph_mode = CUDAGraphMode.FULL_DECODE_ONLY else: cudagraph_mode = CUDAGraphMode.NONE self.cudagraph_manager = EagleCudaGraphManager( - vllm_config, device, cudagraph_mode, self.draft_tokens + self.vllm_config, + self.device, + cudagraph_mode, + self.draft_tokens, ) def load_model(self, target_model: nn.Module) -> None: @@ -133,7 +134,7 @@ class EagleSpeculator: ) -> None: self.model_state = model_state self.kv_cache_config = kv_cache_config - _, self.attn_groups = init_attn_backend( + _, self.attn_groups, _ = init_attn_backend( kv_cache_config, self.vllm_config, self.device, @@ -242,29 +243,6 @@ class EagleSpeculator: idx_mapping, query_start_loc, pos, num_tokens_padded ) - def _dispatch_and_sync_dp( - self, - cudagraph_manager: EagleCudaGraphManager, - num_reqs: int, - num_tokens: int, - uniform_token_count: int | None, - ) -> tuple[BatchExecutionDescriptor, torch.Tensor | None]: - batch_desc = cudagraph_manager.dispatch( - num_reqs, num_tokens, uniform_token_count - ) - num_tokens_across_dp = None - if self.dp_size > 1: - batch_desc, num_tokens_across_dp = sync_cudagraph_and_dp_padding( - cudagraph_manager, - batch_desc, - num_tokens, - num_reqs, - uniform_token_count, - self.dp_size, - self.dp_rank, - ) - return batch_desc, num_tokens_across_dp - def _build_draft_attn_metadata( self, num_reqs: int, @@ -303,8 +281,10 @@ class EagleSpeculator: return attn_metadata def capture_model(self) -> None: + assert self.cudagraph_manager is not None if self.num_speculative_steps == 1: return + logger.info("Capturing model for Eagle speculator...") self.cudagraph_manager.capture( self.generate_draft, @@ -342,6 +322,7 @@ class EagleSpeculator: dummy_run: bool = False, skip_attn_for_dummy_run: bool = False, mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, + is_profile: bool = False, ) -> torch.Tensor: # NOTE(woosuk): To avoid CPU-GPU synchronization without CPU knowing the # number of rejected tokens, we maintain the size of eagle's input_ids and @@ -430,11 +411,14 @@ class EagleSpeculator: # Each request produces exactly 1 token per draft decode step, # enabling FULL cudagraph. - decode_batch_desc, num_tokens_across_dp = self._dispatch_and_sync_dp( + decode_batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp( self.cudagraph_manager, num_reqs, num_reqs, uniform_token_count=1, + dp_size=self.dp_size, + dp_rank=self.dp_rank, + need_eager=is_profile, ) attn_metadata_updated = None @@ -461,6 +445,7 @@ class EagleSpeculator: ) if decode_batch_desc.cg_mode == CUDAGraphMode.FULL: + assert self.cudagraph_manager is not None self.cudagraph_manager.run_fullgraph(decode_batch_desc) else: self.generate_draft( diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index d1ec9126ceb..4f3e192772f 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -2783,7 +2783,20 @@ class GPUModelRunner( ) self.lora_manager.set_active_adapters(lora_requests, tower_mapping) - if hasattr(self.model, "get_num_mm_connector_tokens"): + # Only set connector mapping if the model actually has a connector. + # Some multimodal models inherit a stub `get_num_mm_connector_tokens` + # from `SupportsMultiModal`, which returns None and should not be + # treated as a signal that connector LoRA is supported. + mm_mapping = ( + self.model.get_mm_mapping() # type: ignore[attr-defined] + if hasattr(self.model, "get_mm_mapping") + else None + ) + if ( + mm_mapping is not None + and mm_mapping.connector + and hasattr(self.model, "get_num_mm_connector_tokens") + ): post_op_counts = [ self.model.get_num_mm_connector_tokens(num_tokens) # type: ignore[attr-defined] for num_tokens in encoder_token_counts @@ -6294,7 +6307,7 @@ class GPUModelRunner( cudagraph_mode. """ min_cg_support = AttentionCGSupport.ALWAYS - min_cg_backend_name = None + min_cg_attn_backend = None for attn_backend_set, kv_cache_group in zip( attention_backends, kv_cache_groups @@ -6307,152 +6320,18 @@ class GPUModelRunner( ) if cg_support.value < min_cg_support.value: min_cg_support = cg_support - min_cg_backend_name = attn_backend.__name__ - # Flexible resolve the cudagraph mode - cudagraph_mode = self.compilation_config.cudagraph_mode - assert cudagraph_mode is not None - # check cudagraph for mixed batch is supported - if ( - cudagraph_mode.mixed_mode() == CUDAGraphMode.FULL - and min_cg_support != AttentionCGSupport.ALWAYS - ): - msg = ( - f"CUDAGraphMode.{cudagraph_mode.name} is not supported " - f"with {min_cg_backend_name} backend (support: " - f"{min_cg_support})" - ) - if min_cg_support == AttentionCGSupport.NEVER: - # if not supported any full cudagraphs, just raise it. - msg += ( - "; please try cudagraph_mode=PIECEWISE, and " - "make sure compilation mode is VLLM_COMPILE" - ) - raise ValueError(msg) - - # attempt to resolve the full cudagraph related mode - if self.compilation_config.splitting_ops_contain_attention(): - msg += "; setting cudagraph_mode=FULL_AND_PIECEWISE" - cudagraph_mode = self.compilation_config.cudagraph_mode = ( - CUDAGraphMode.FULL_AND_PIECEWISE - ) - else: - msg += "; setting cudagraph_mode=FULL_DECODE_ONLY" - cudagraph_mode = self.compilation_config.cudagraph_mode = ( - CUDAGraphMode.FULL_DECODE_ONLY - ) - logger.warning(msg) - - # check that if we are doing decode full-cudagraphs it is supported - if ( - cudagraph_mode.decode_mode() == CUDAGraphMode.FULL - and min_cg_support == AttentionCGSupport.NEVER - ): - msg = ( - f"CUDAGraphMode.{cudagraph_mode.name} is not supported " - f"with {min_cg_backend_name} backend (support: " - f"{min_cg_support})" - ) - if self.compilation_config.mode == CompilationMode.VLLM_COMPILE and ( - self.compilation_config.splitting_ops_contain_attention() - or self.compilation_config.use_inductor_graph_partition - ): - msg += ( - "; setting cudagraph_mode=PIECEWISE because " - "attention is compiled piecewise" - ) - cudagraph_mode = self.compilation_config.cudagraph_mode = ( - CUDAGraphMode.PIECEWISE - ) - else: - msg += ( - "; setting cudagraph_mode=NONE because " - "attention is not compiled piecewise" - ) - cudagraph_mode = self.compilation_config.cudagraph_mode = ( - CUDAGraphMode.NONE - ) - logger.warning(msg) - - # check that if we are doing spec-decode + decode full-cudagraphs it is - # supported - if ( - cudagraph_mode.decode_mode() == CUDAGraphMode.FULL - and self.uniform_decode_query_len > 1 - and min_cg_support.value < AttentionCGSupport.UNIFORM_BATCH.value - ): - msg = ( - f"CUDAGraphMode.{cudagraph_mode.name} is not supported" - f" with spec-decode for attention backend " - f"{min_cg_backend_name} (support: {min_cg_support})" - ) - if self.compilation_config.splitting_ops_contain_attention(): - msg += "; setting cudagraph_mode=PIECEWISE" - cudagraph_mode = self.compilation_config.cudagraph_mode = ( - CUDAGraphMode.PIECEWISE - ) - else: - msg += "; setting cudagraph_mode=NONE" - cudagraph_mode = self.compilation_config.cudagraph_mode = ( - CUDAGraphMode.NONE - ) - logger.warning(msg) - - # double check that we can support full cudagraph if they are requested - # even after automatic downgrades - if ( - cudagraph_mode.has_full_cudagraphs() - and min_cg_support == AttentionCGSupport.NEVER - ): - raise ValueError( - f"CUDAGraphMode.{cudagraph_mode.name} is not " - f"supported with {min_cg_backend_name} backend (" - f"support:{min_cg_support}) " - "; please try cudagraph_mode=PIECEWISE, " - "and make sure compilation mode is VLLM_COMPILE" - ) - - # if we have dedicated decode cudagraphs, and spec-decode is enabled, - # we need to adjust the cudagraph sizes to be a multiple of the uniform - # decode query length to avoid: https://github.com/vllm-project/vllm/issues/28207 - # temp-fix: https://github.com/vllm-project/vllm/issues/28207#issuecomment-3504004536 - # Will be removed in the near future when we have separate cudagraph capture - # sizes for decode and mixed prefill-decode. - if ( - cudagraph_mode.decode_mode() == CUDAGraphMode.FULL - and cudagraph_mode.separate_routine() - and self.uniform_decode_query_len > 1 - ): - self.compilation_config.adjust_cudagraph_sizes_for_spec_decode( - self.uniform_decode_query_len, self.parallel_config.tensor_parallel_size - ) - - # For Mamba models with FULL decode cudagraphs, each decode - # sequence needs one Mamba cache block. The decode cudagraph - # dispatcher already caps batch sizes at max_num_seqs, so we just - # need to verify that enough blocks exist. Raising here instead - # of silently capping cudagraph_capture_sizes avoids unintended - # restrictions on PIECEWISE (prefill) cudagraphs. - # See: https://github.com/vllm-project/vllm/issues/34094 - if cudagraph_mode.has_full_cudagraphs() and not is_profiling: - has_mamba = any( - isinstance(g.kv_cache_spec, MambaSpec) for g in kv_cache_groups - ) - if has_mamba and self.kv_cache_config is not None: - num_blocks = self.kv_cache_config.num_blocks - if self.max_num_reqs > num_blocks: - raise ValueError( - f"max_num_seqs ({self.max_num_reqs}) exceeds " - f"available Mamba cache blocks ({num_blocks}). " - f"Each decode sequence requires one Mamba cache " - f"block, so CUDA graph capture cannot proceed. " - f"Please lower max_num_seqs to at most " - f"{num_blocks} or increase " - f"gpu_memory_utilization." - ) - + min_cg_attn_backend = attn_backend.__name__ + cudagraph_mode = self.compilation_config.resolve_cudagraph_mode_and_sizes( + min_cg_support, + min_cg_attn_backend, + self.uniform_decode_query_len, + self.parallel_config.tensor_parallel_size, + self.kv_cache_config, + self.max_num_reqs, + is_profiling=is_profiling, + ) # Trigger cudagraph dispatching keys initialization after # resolved cudagraph mode. - self.compilation_config.cudagraph_mode = cudagraph_mode self.cudagraph_dispatcher.initialize_cudagraph_keys( cudagraph_mode, self.uniform_decode_query_len )