From 4ac0d0eac25eecc98ff519887f2aa3d659dc75f3 Mon Sep 17 00:00:00 2001 From: AIwork4me Date: Sat, 3 Oct 2026 23:44:17 +0800 Subject: [PATCH 01/22] [ROCm][RDNA3] Fix W4A16 split-K accuracy and determinism (#54706) Signed-off-by: AIwork4me Co-authored-by: AIwork4me Co-authored-by: JartX --- csrc/rocm/moe_q_gemm_rdna3.cu | 82 ++- csrc/rocm/q_gemm_rdna3.cu | 582 +++++++++++++---- csrc/rocm/q_gemm_rdna3_wmma.cu | 435 +++++++++++-- csrc/rocm/qdq_4_rdna3.cuh | 106 ++-- .../test_rdna3_w4a16_determinism.py | 585 ++++++++++++++++++ 5 files changed, 1524 insertions(+), 266 deletions(-) create mode 100644 tests/kernels/quantization/test_rdna3_w4a16_determinism.py diff --git a/csrc/rocm/moe_q_gemm_rdna3.cu b/csrc/rocm/moe_q_gemm_rdna3.cu index 61ffbb9987..9573fe731b 100644 --- a/csrc/rocm/moe_q_gemm_rdna3.cu +++ b/csrc/rocm/moe_q_gemm_rdna3.cu @@ -57,16 +57,6 @@ __forceinline__ __device__ bf16_t tzero() { return __float2bfloat16(0.0f); } -__forceinline__ __device__ float dot22_8_f(half2 (&dq)[4], const half* a_ptr) { - float result = 0.0f; - const half2* a2_ptr = (const half2*)a_ptr; - #pragma unroll - for (int i = 0; i < 4; i++) { - result = __builtin_amdgcn_fdot2(dq[i], *a2_ptr++, result, /*clamp=*/false); - } - return result; -} - __forceinline__ __device__ void atomic_add_pk4_f16(half* addr, half2 v01, half2 v23) { unsigned long long* addr_u = reinterpret_cast(addr); @@ -177,7 +167,8 @@ __global__ void moe_gemm_q4_kernel_rdna3( "BLOCK_KN_SIZE must equal THREADS_X"); // For bf16 M=1, we can skip LDS and read A from global (same as dense). - // fp16 always needs LDS due to the dot22_8_f indexing pattern. + // fp16 always needs LDS: the exact-factored inner loop indexes + // block_a[m][a_off] directly. constexpr bool USE_LDS_A = (BLOCK_SIZE_M > 1) || std::is_same::value; const int offset_m_base = token_block * BLOCK_SIZE_M; @@ -211,8 +202,10 @@ __global__ void moe_gemm_q4_kernel_rdna3( int qk = offset_k / 8; const uint32_t* b_ptr = expert_weights + qk * size_n + n; - // Per-column dequant constants (4 columns per thread) - half2 z1z16_h[4][2], y1y16_h[4][2]; + // Per-column dequant constants (4 columns per thread). fp16 uses the + // exact factored form (fp32 y/z pairs, no baked bias — see + // qdq_4_rdna3.cuh); bf16 keeps fp32 scalars as before. + float yf_h[4][2], zf_h[4][2]; float z_b_f[4], y_b_f[4]; // GPTQv1: zero_offset = 1 @@ -228,8 +221,8 @@ __global__ void moe_gemm_q4_kernel_rdna3( if constexpr (std::is_same::value) { #pragma unroll for (int i = 0; i < 4; ++i) { - gptq_rdna3::prep_zero_scale_fp16((uint32_t)(zeros[i] + zero_offset), - scales[i], z1z16_h[i], y1y16_h[i]); + gptq_rdna3::prep_zero_scale_fp16_f32((uint32_t)(zeros[i] + zero_offset), + scales[i], yf_h[i], zf_h[i]); } } else { #pragma unroll @@ -266,29 +259,42 @@ __global__ void moe_gemm_q4_kernel_rdna3( } b_ptr += 4 * size_n; + // fp16 exact factored form: activation sums per row, split by nibble + // slot, accumulated across the 32-K round; the z correction is applied + // once after the j loop (see qdq_4_rdna3.cuh). + float sum_lo[BLOCK_SIZE_M] = {}, sum_hi[BLOCK_SIZE_M] = {}; + #pragma unroll for (int j = 0; j < 4; ++j) { const int a_off = (k - offset_k) + 8 * j; if constexpr (std::is_same::value) { - // fp16 path: dequant via bit-trick, dot via v_dot2_f32_f16 - half2 dq[4][4]; - gptq_rdna3::dequant_4bit_8_fp16((uint32_t)b_w[j].x, dq[0], z1z16_h[0], - y1y16_h[0]); - gptq_rdna3::dequant_4bit_8_fp16((uint32_t)b_w[j].y, dq[1], z1z16_h[1], - y1y16_h[1]); - gptq_rdna3::dequant_4bit_8_fp16((uint32_t)b_w[j].z, dq[2], z1z16_h[2], - y1y16_h[2]); - gptq_rdna3::dequant_4bit_8_fp16((uint32_t)b_w[j].w, dq[3], z1z16_h[3], - y1y16_h[3]); + // fp16 path: magic values without bias, dots via v_dot2_f32_f16, + // the (y, z) correction applied in fp32 (qdq_4_rdna3.cuh). + half2 qm[4][4]; + gptq_rdna3::magic_4bit_8_fp16((uint32_t)b_w[j].x, qm[0]); + gptq_rdna3::magic_4bit_8_fp16((uint32_t)b_w[j].y, qm[1]); + gptq_rdna3::magic_4bit_8_fp16((uint32_t)b_w[j].z, qm[2]); + gptq_rdna3::magic_4bit_8_fp16((uint32_t)b_w[j].w, qm[3]); + + const half2 ones = gptq_rdna3::ones_half2_fp16(); #pragma unroll for (int m = 0; m < BLOCK_SIZE_M; ++m) { - const half* a_ptr = reinterpret_cast(&block_a[m][a_off]); - block_c[m][0] += dot22_8_f(dq[0], a_ptr); - block_c[m][1] += dot22_8_f(dq[1], a_ptr); - block_c[m][2] += dot22_8_f(dq[2], a_ptr); - block_c[m][3] += dot22_8_f(dq[3], a_ptr); + const half2* a2 = reinterpret_cast(&block_a[m][a_off]); + sum_lo[m] = __builtin_amdgcn_fdot2(a2[0], ones, sum_lo[m], false); + sum_hi[m] = __builtin_amdgcn_fdot2(a2[1], ones, sum_hi[m], false); + sum_lo[m] = __builtin_amdgcn_fdot2(a2[2], ones, sum_lo[m], false); + sum_hi[m] = __builtin_amdgcn_fdot2(a2[3], ones, sum_hi[m], false); + #pragma unroll + for (int c = 0; c < 4; ++c) { + float pl = 0.0f, ph = 0.0f; + pl = __builtin_amdgcn_fdot2(qm[c][0], a2[0], pl, /*clamp=*/false); + ph = __builtin_amdgcn_fdot2(qm[c][1], a2[1], ph, /*clamp=*/false); + pl = __builtin_amdgcn_fdot2(qm[c][2], a2[2], pl, /*clamp=*/false); + ph = __builtin_amdgcn_fdot2(qm[c][3], a2[3], ph, /*clamp=*/false); + block_c[m][c] += yf_h[c][0] * pl + yf_h[c][1] * ph; + } } } else if constexpr (BLOCK_SIZE_M == 1) { // bf16 M=1: v_dot2_f32_bf16 with InstCombine-defeating opacity @@ -415,6 +421,18 @@ __global__ void moe_gemm_q4_kernel_rdna3( } } } + + if constexpr (std::is_same::value) { + // z is a group constant and sum_a is complete for this round: apply + // the fp16 z correction once per round (see qdq_4_rdna3.cuh). + #pragma unroll + for (int m = 0; m < BLOCK_SIZE_M; ++m) { + #pragma unroll + for (int c = 0; c < 4; ++c) { + block_c[m][c] += zf_h[c][0] * sum_lo[m] + zf_h[c][1] * sum_hi[m]; + } + } + } k += 32; } @@ -587,6 +605,10 @@ void moe_gptq_gemm_rdna3(torch::Tensor a, torch::Tensor c, int size_k = (int)a.size(1); int size_n = (int)b_q_weight.size(2); int groups = (int)b_scales.size(1); + TORCH_CHECK(size_k % (groups * 32) == 0, + "group size (K/groups = ", size_k / groups, + ") must be a multiple of 32: the kernel checks group " + "transitions at 32-K granularity"); // Per-expert strides int expert_weight_stride = (int)(b_q_weight.size(1) * b_q_weight.size(2)); diff --git a/csrc/rocm/q_gemm_rdna3.cu b/csrc/rocm/q_gemm_rdna3.cu index bf23682699..e6f1c2fb17 100644 --- a/csrc/rocm/q_gemm_rdna3.cu +++ b/csrc/rocm/q_gemm_rdna3.cu @@ -5,31 +5,42 @@ // activation dtype (half or __hip_bfloat16). Adapted from exllamav2's 4-bit // kernel (csrc/quantization/gptq/q_gemm.cu) with the following changes: // -// 1. Direct write to the T-typed output via packed CAS-loop on a 64-bit -// word (atomic_add_pk4_{f16,bf16}). gfx11 has no native -// v_global_atomic_pk_add_{f16,bf16}, so the kernel emulates one with -// global_atomic_cmpswap_b64. This avoids the M*N*4-byte FP32 scratch -// buffer + memset + cast-pass that an fp32-accumulator design would -// need; the caller passes a zero-initialised T-typed output tensor -// and every block atomically adds its partial sum into it. +// 1. Deterministic split-K epilogue. K is split across gridDim.z blocks; +// each block stores its FP32 block partial to a scratch tensor, and the +// last z block of each (x, y) tile reduces the z-slices in fixed +// ascending order with a single final cast to T — the reduction runs +// in the GEMM kernel itself, gated by a per-call ticket counter (no +// cross-launch state; see the epilogue). When gridDim.z == 1 each +// output element has exactly one writer and the kernel stores the +// rounded accumulator directly — no scratch, no reduce, no atomics. +// (Legacy design, pre-#54706: a packed CAS-loop atomic add on a 64-bit +// word emulating v_global_atomic_pk_add_{f16,bf16}, which gfx11 lacks. +// It narrowed every split partial to bf16/fp16 BEFORE accumulation, so +// the execution-dependent CAS completion order changed the rounded +// result for identical inputs. The CAS epilogue is kept below the +// partials branch for A/B comparison only; it requires a pre-zeroed +// output tensor.) // // 2. The bf16 path uses a dedicated bit-trick that avoids the fp16-only // "upper nibble * 16" trick, which would overflow the 7-bit bf16 -// mantissa. See qdq_4_rdna3.cuh for details. +// mantissa. See qdq_4_rdna3.cuh for details. The fp16 path uses the +// exact factored dequant (also in qdq_4_rdna3.cuh): magic values in +// fp16, the (y, z) bias correction in fp32. // // 3. Wave32 geometry sized for high CU saturation: THREADS_X=256 // (8 waves per block) and BLOCK_KN_SIZE=256, with each thread // computing 4 N output columns. gridDim.z = K / BLOCK_KN_SIZE -// splits K and the output is atomically accumulated. fp16 uses +// splits K; the epilogue stores FP32 partials for the fixed-order +// reduce (see item 1). fp16 uses // v_dot2_f32_f16 (__builtin_amdgcn_fdot2) for the inner dot; // bf16 widens to fp32 (no v_pk_fma_bf16 on gfx11) and accumulates // with v_fma_f32. M_COUNT ∈ {1,2,4,8} is selected at launch // based on size_m. // -// 4. The bf16 dispatch with M >= 16 forwards to the WMMA kernel in +// 4. The dispatch with M >= 12 forwards to the WMMA kernel in // q_gemm_rdna3_wmma.cu (separate translation unit) where -// v_wmma_f32_16x16x16_bf16_w32 wins. The fp16 path always stays -// scalar (the bit-trick dequant beats WMMA below M=64). +// v_wmma_f32_16x16x16_bf16_w32 wins. Below that the scalar path wins +// and is the more accurate of the two. #include #include @@ -92,40 +103,16 @@ __forceinline__ __device__ bf16_t tzero() { return __float2bfloat16(0.0f); } -__forceinline__ __device__ float dot22_8_f(half2 (&dq)[4], const half* a_ptr) { - // RDNA3 has v_dot2_f32_f16 (`__builtin_amdgcn_fdot2`) which computes - // fp32 += a.x*b.x + a.y*b.y in a single instruction with the accumulator - // staying in fp32 throughout. hipcc 7.2 does NOT peephole the obvious - // `__hfma2 + cast + add` pattern into v_dot2 (verified by ISA - // disassembly: 0 v_dot2_f32_f16 vs 256 v_cvt_f32_f16 + 218 v_add_f32 in - // the M_COUNT=8 kernel before this change), so we issue the builtin - // explicitly. Saves the trailing 2× v_cvt_f32_f16 + v_add_f32 (3 ops) - // per dot22_8_f call vs the half2-accumulator form. With 128 calls per - // K=32 step that's ~384 ops/K-step less issue pressure on the VALU. - // - // Numerical bonus: accumulator stays fp32 throughout the dot. The old - // form accumulated 8 muladds in fp16 (10-bit mantissa) before casting, - // which could lose ~3 bits of precision on borderline magnitudes. - float result = 0.0f; - const half2* a2_ptr = (const half2*)a_ptr; - #pragma unroll - for (int i = 0; i < 4; i++) { - result = __builtin_amdgcn_fdot2(dq[i], *a2_ptr++, result, /*clamp=*/false); - } - return result; -} - // --------------------------------------------------------------------------- // Packed atomic-add via CAS-loop on a 64-bit word (4 fp16/bf16 lanes per CAS). // RDNA3 (gfx11) does NOT have native v_global_atomic_pk_add_f16 / _bf16 (those // landed on gfx940 / gfx1250 respectively), so this lowers to -// global_atomic_cmpswap_b64 plus retry. We use this in the kernel epilogue to -// write 4 output columns per row in a single atomic operation — half the -// atomic instruction count and half the contention vs two 32-bit CAS calls. -// -// Writing directly to fp16/bf16 (instead of through an FP32 scratch buffer + -// cast pass) saves M*N*4 bytes of allocation, the memset, and the epilogue -// cast pass that an fp32-accumulator design would need. +// global_atomic_cmpswap_b64 plus retry. This is the LEGACY (pre-#54706) +// epilogue: the caller must pre-zero the output, and every split block +// atomically adds its low-precision partial into it. Because the addition +// happens in bf16/fp16 AFTER narrowing, the result depends on CAS completion +// order — the nondeterminism PR #54706 fixes. Kept for A/B comparison with +// the deterministic path; the shipped dispatch never selects it. // // 64-bit alignment: the kernel writes at `out + n` where n = offset_n + t*4 // (always multiple of 4), and partition_weight_shape[1] is required to be a @@ -198,13 +185,26 @@ __forceinline__ __device__ void load4_scales(const T* scales_row, int n, // --------------------------------------------------------------------------- template -__global__ void gemm_q4_kernel_rdna3(const T* __restrict__ a, - const uint32_t* __restrict__ b_q_weight, - const uint32_t* __restrict__ b_qzeros, - const T* __restrict__ b_scales, - T* __restrict__ c, const int size_m, - const int size_n, const int size_k, - const int groups, const int zero_offset) { +__global__ void gemm_q4_kernel_rdna3( + const T* __restrict__ a, const uint32_t* __restrict__ b_q_weight, + const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales, + T* __restrict__ c, const int size_m, const int size_n, const int size_k, + const int groups, const int zero_offset, + // Deterministic split-K epilogue: when + // non-null, each split block stores its + // FP32 partial to + // partials[(z*size_m + m)*size_n + n] + // instead of CAS-atomically accumulating + // a low-precision partial into c. The + // last z block of each (x, y) tile then + // performs the fixed ascending-z + // reduction in-kernel (see the epilogue + // below). `tickets` (same length as the + // number of tiles) is a per-call + // zero-filled arrival counter array — + // one ticket per tile — so concurrent + // invocations never share state. + float* __restrict__ partials, int* __restrict__ tickets) { const int t = threadIdx.x; const int offset_n = blockIdx.x * BLOCK_KN_SIZE * 4; const int offset_m = blockIdx.y * M_COUNT; @@ -270,27 +270,24 @@ __global__ void gemm_q4_kernel_rdna3(const T* __restrict__ a, int qk = offset_k / 8; const uint32_t* b_ptr = b_q_weight + qk * size_n + n; - // Per-column dequant constants. We hold one set of (z, y) pairs per column. - // fp16 uses the exllama (z1z16, y1y16) double-pair to enable the upper- - // nibble-*16 trick. bf16 uses fp32 scalars (z, y) for the inline v_dot2 - // bias correction. - half2 z1z16_h[4][2], y1y16_h[4][2]; + // Per-column dequant constants, fp32 for both dtypes. fp16 needs two (y, z) + // pairs because the upper-nibble trick carries a factor of 16; bf16 needs one + // because its dequant produces fp32 directly — see prep_zero_scale_bf16_f32 / + // the FMA bypass for the missing v_pk_fma_bf16 on gfx11. Neither narrows the + // bias to the activation dtype: see the exact factored form in + // qdq_4_rdna3.cuh for why that matters. float z_b_f[4], y_b_f[4]; + float yf_h[4][2], zf_h[4][2]; + // bf16-only group refresh (fp16 preps from prefetched raw words below). auto refresh_group = [&](int g) { - const uint32_t* qz_row = b_qzeros + g * (size_n / 8); - const T* sc_row = b_scales + g * size_n; - int zeros[4]; - T scales[4]; - load4_zeros(qz_row, n, zeros); - load4_scales(sc_row, n, scales); - if constexpr (std::is_same::value) { - #pragma unroll - for (int i = 0; i < 4; ++i) { - prep_zero_scale_fp16((uint32_t)(zeros[i] + zero_offset), scales[i], - z1z16_h[i], y1y16_h[i]); - } - } else { + if constexpr (!std::is_same::value) { + const uint32_t* qz_row = b_qzeros + g * (size_n / 8); + const T* sc_row = b_scales + g * size_n; + int zeros[4]; + T scales[4]; + load4_zeros(qz_row, n, zeros); + load4_scales(sc_row, n, scales); #pragma unroll for (int i = 0; i < 4; ++i) { prep_zero_scale_bf16_f32((uint32_t)(zeros[i] + zero_offset), scales[i], @@ -299,7 +296,32 @@ __global__ void gemm_q4_kernel_rdna3(const T* __restrict__ a, } }; - refresh_group(group); + // fp16: the raw zero/scale words of the NEXT group are loaded during the + // round before the boundary, so a group change no longer stalls on a full + // memory round trip before the weights of the round are even issued. + uint32_t rz = 0; + T rs[4]; + auto load_raw_group = [&](int g) { + rz = b_qzeros[g * (size_n / 8) + n / 8] >> ((n & 7) * 4); + #pragma unroll + for (int i = 0; i < 4; ++i) rs[i] = b_scales[g * size_n + n + i]; + }; + auto prep_raw_fp16 = [&]() { + if constexpr (std::is_same::value) { + #pragma unroll + for (int i = 0; i < 4; ++i) { + prep_zero_scale_fp16_f32(((rz >> (4 * i)) & 0xF) + zero_offset, rs[i], + yf_h[i], zf_h[i]); + } + } + }; + + if constexpr (std::is_same::value) { + load_raw_group(group); + prep_raw_fp16(); + } else { + refresh_group(group); + } float block_c[M_COUNT][4]; #pragma unroll @@ -309,11 +331,12 @@ __global__ void gemm_q4_kernel_rdna3(const T* __restrict__ a, } // Note on group-transition granularity: we check `k == nextgroup` at the - // start of each outer iteration (which advances K by 32). This is correct - // when group_size >= 32 OR group_size divides 32 evenly (groupsize is one - // of {1,2,4,8,16,32,64,128,...}). For group_size in {16, 8, 4, ...} the - // inner loop would cross a group boundary between j-iterations; we require - // group_size >= 32 here, mirroring exllama's assumption. + // start of each outer iteration (which advances K by 32), so a group + // boundary must never fall strictly inside a 32-K round: groupsize must + // be a multiple of 32 (all real GPTQ group sizes 32/64/128/256, or + // groups == 1). groupsize in {16, 8, ...} or non-multiples like 48 would + // skip the boundary check and silently use stale constants; the host-side + // TORCH_CHECK below enforces it. Same assumption as exllama. // // Software pipelining: we issue all 4 vectorized weight loads up front // before any dequant/FMA depends on them. This gives the AMDGPU backend @@ -325,7 +348,11 @@ __global__ void gemm_q4_kernel_rdna3(const T* __restrict__ a, if (k == nextgroup) { group++; nextgroup += groupsize; - refresh_group(group); + if constexpr (std::is_same::value) { + prep_raw_fp16(); // from the words prefetched in the previous round + } else { + refresh_group(group); + } } // Prefetch all four j-iterations' weight words. The compiler emits 4 @@ -337,25 +364,55 @@ __global__ void gemm_q4_kernel_rdna3(const T* __restrict__ a, b_w[j] = *(const int4*)(b_ptr + j * size_n); } b_ptr += 4 * size_n; + if constexpr (std::is_same::value) { + if (k + 32 == nextgroup && k + 32 < end_k) load_raw_group(group + 1); + } + + // fp16 exact factored form: raw dot products and activation sums, both + // split by nibble slot, accumulated across the whole 32-K round. Neither + // depends on the output column, and the (y, z) constants are group + // constants, so both corrections are applied once after the j loop + // instead of four times inside it — see qdq_4_rdna3.cuh. + float sum_lo[M_COUNT] = {}, sum_hi[M_COUNT] = {}; + float acc_lo[M_COUNT][4] = {}, acc_hi[M_COUNT][4] = {}; #pragma unroll for (int j = 0; j < 4; ++j) { const int a_off = (k - offset_k) + 8 * j; if constexpr (std::is_same::value) { - half2 dq[4][4]; - dequant_4bit_8_fp16((uint32_t)b_w[j].x, dq[0], z1z16_h[0], y1y16_h[0]); - dequant_4bit_8_fp16((uint32_t)b_w[j].y, dq[1], z1z16_h[1], y1y16_h[1]); - dequant_4bit_8_fp16((uint32_t)b_w[j].z, dq[2], z1z16_h[2], y1y16_h[2]); - dequant_4bit_8_fp16((uint32_t)b_w[j].w, dq[3], z1z16_h[3], y1y16_h[3]); + // Magic values with no bias folded in; the (y, z) correction is + // applied in fp32 below — see the exact factored form in + // qdq_4_rdna3.cuh. + half2 qm[4][4]; + magic_4bit_8_fp16((uint32_t)b_w[j].x, qm[0]); + magic_4bit_8_fp16((uint32_t)b_w[j].y, qm[1]); + magic_4bit_8_fp16((uint32_t)b_w[j].z, qm[2]); + magic_4bit_8_fp16((uint32_t)b_w[j].w, qm[3]); + + const half2 ones = ones_half2_fp16(); #pragma unroll for (int m = 0; m < M_COUNT; ++m) { - const half* a_ptr = reinterpret_cast(&block_a[m][a_off]); - block_c[m][0] += dot22_8_f(dq[0], a_ptr); - block_c[m][1] += dot22_8_f(dq[1], a_ptr); - block_c[m][2] += dot22_8_f(dq[2], a_ptr); - block_c[m][3] += dot22_8_f(dq[3], a_ptr); + const half2* a2 = reinterpret_cast(&block_a[m][a_off]); + // sum_a split by nibble slot (the high pairs carry a factor of + // 16). It does not depend on the output column, so one v_dot2 per + // pair is shared by the four columns this thread owns. + sum_lo[m] = __builtin_amdgcn_fdot2(a2[0], ones, sum_lo[m], false); + sum_hi[m] = __builtin_amdgcn_fdot2(a2[1], ones, sum_hi[m], false); + sum_lo[m] = __builtin_amdgcn_fdot2(a2[2], ones, sum_lo[m], false); + sum_hi[m] = __builtin_amdgcn_fdot2(a2[3], ones, sum_hi[m], false); + #pragma unroll + for (int c = 0; c < 4; ++c) { + acc_lo[m][c] = + __builtin_amdgcn_fdot2(qm[c][0], a2[0], acc_lo[m][c], false); + acc_hi[m][c] = + __builtin_amdgcn_fdot2(qm[c][1], a2[1], acc_hi[m][c], false); + acc_lo[m][c] = + __builtin_amdgcn_fdot2(qm[c][2], a2[2], acc_lo[m][c], false); + acc_hi[m][c] = + __builtin_amdgcn_fdot2(qm[c][3], a2[3], acc_hi[m][c], false); + } } } else if constexpr (M_COUNT == 1) { // bf16 decode (M=1), v_dot2_f32_bf16 path. Mirrors the data-flow of @@ -520,18 +577,84 @@ __global__ void gemm_q4_kernel_rdna3(const T* __restrict__ a, } } } + + if constexpr (std::is_same::value) { + // y and z are group constants and this round's dots and activation + // sums are complete: one pair of FMAs per (row, column) per round + // instead of per j-iteration. Per-round z applications sum exactly to + // the per-group correction because z does not change inside a group. + #pragma unroll + for (int m = 0; m < M_COUNT; ++m) { + #pragma unroll + for (int c = 0; c < 4; ++c) { + block_c[m][c] += yf_h[c][0] * acc_lo[m][c] + + yf_h[c][1] * acc_hi[m][c] + zf_h[c][0] * sum_lo[m] + + zf_h[c][1] * sum_hi[m]; + } + } + } k += 32; // 4 weight words * 8 nibbles = 32 K elements } - // Pack the 4 FP32 partial sums into 2 packed pairs and atomically add all - // four lanes in a single 64-bit CAS write directly to the T-typed output - // (caller pre-zeros it). On gfx11 the packed atomic is a CAS-loop, but with - // a single b64 op we halve the atomic instruction count vs two b32 CAS - // calls, AND save the FP32 buffer + memset + cast pass entirely. + // ---- Epilogue: three store modes, selected by launch shape ---- + // + // partials != nullptr: deterministic split-K. Plain FP32 stores to + // partials[(blockIdx.z * size_m + m) * size_n + n] + // (one writer per in-range slot; threads with n >= size_n returned before + // the epilogue and rows past size_m are skipped — exactly the slots the + // in-kernel fixed-order reduce below never reads). n is a multiple of 4 + // and size_n % 8 == 0, so the 4-lane store never crosses the right edge. + // After the stores, the block publishes its arrival on the tile's ticket + // (a per-call zero-filled counter); the LAST z block of the tile then + // sums the z slices in fixed ascending order with a single final cast — + // the reduction order is a pure function of the launch shape, so the + // result stays bit-reproducible for identical inputs. Tickets are + // per-invocation state (see launch_gemm_q4_deterministic); nothing is + // shared across kernel launches, so concurrent invocations on other + // streams cannot interfere. + // + // partials == nullptr && gridDim.z == 1: single split block per output + // element — direct store of the rounded accumulator. No scratch, no + // reduce, no atomics; c may be left uninitialized (torch::empty). + // + // partials == nullptr && gridDim.z > 1: LEGACY pre-#54706 CAS epilogue + // (A/B control only; never selected by launch_gemm_q4_deterministic). + // Adds the narrowed 4-lane partial into c with one 64-bit CAS per 4 + // columns; REQUIRES a pre-zeroed output because the CAS adds into + // whatever is already there. #pragma unroll for (int m = 0; m < M_COUNT; ++m) { if (offset_m + m >= size_m) continue; // skip padding rows past size_m + if (partials != nullptr) { + float* p = + partials + ((long)blockIdx.z * size_m + (offset_m + m)) * size_n + n; + p[0] = block_c[m][0]; + p[1] = block_c[m][1]; + p[2] = block_c[m][2]; + p[3] = block_c[m][3]; + continue; + } T* out = c + (offset_m + m) * size_n + n; + if (gridDim.z == 1) { + // Single writer per element: round once and store the packed 4 lanes + // (8-byte aligned, see the note above) in one go. + if constexpr (std::is_same::value) { + half2 packed[2] = {__halves2half2(__float2half_rn(block_c[m][0]), + __float2half_rn(block_c[m][1])), + __halves2half2(__float2half_rn(block_c[m][2]), + __float2half_rn(block_c[m][3]))}; + __builtin_memcpy(out, packed, sizeof(packed)); + } else { + bf162_t packed[2]; + packed[0].x = __float2bfloat16(block_c[m][0]); + packed[0].y = __float2bfloat16(block_c[m][1]); + packed[1].x = __float2bfloat16(block_c[m][2]); + packed[1].y = __float2bfloat16(block_c[m][3]); + __builtin_memcpy(out, packed, sizeof(packed)); + } + continue; + } + // Legacy CAS-atomic epilogue: see the store-mode comment above. if constexpr (std::is_same::value) { half2 r01 = __halves2half2(__float2half_rn(block_c[m][0]), __float2half_rn(block_c[m][1])); @@ -548,6 +671,67 @@ __global__ void gemm_q4_kernel_rdna3(const T* __restrict__ a, atomic_add_pk4_bf16(out, r01, r23); } } + + // Deterministic in-kernel split-K reduction (fp16 only — on gfx1100 the + // bf16 scalar kernel measured neutral at M=1 and up to 10% slower at + // M=4-8 with the reduction fused, so bf16 keeps the separate reducer; + // see launch_gemm_q4_deterministic). Every (x, y) tile has gridDim.z + // producer blocks; the last one to finish sums the tile's z slices in + // fixed ascending order and writes the output with a single final cast. + // Memory ordering: each thread's partials stores precede its + // __threadfence() (release, device scope); the fence precedes the block + // barrier, so thread 0's ticket increment happens after every store of + // the block is device-visible. The last block's second __threadfence() + // pairs with the writers' fences (acquire side), making the other blocks' + // partials visible before they are read. The ticket is re-armed after the + // reduction so the same per-call array serves the next row tile + // (launches on a stream are ordered, so the re-arm is complete before any + // later launch reads the ticket). + if constexpr (std::is_same::value) { + if (partials != nullptr) { + __shared__ int s_last; + __threadfence(); + __syncthreads(); + if (t == 0) { + s_last = atomicAdd(tickets + blockIdx.y * gridDim.x + blockIdx.x, 1) == + (int)gridDim.z - 1; + } + __syncthreads(); + if (!s_last) return; + __threadfence(); + + typedef float f4v __attribute__((ext_vector_type(4))); + const long zs = ((long)size_m * size_n) / 4; // z-slice stride, f4 units + const int zc = (int)gridDim.z; + #pragma unroll + for (int m = 0; m < M_COUNT; ++m) { + if (offset_m + m >= size_m) continue; // skip padding rows past size_m + const f4v* p = + (const f4v*)(partials + (long)(offset_m + m) * size_n + n); + f4v acc = {}; + int z = 0; + // Ascending-z accumulation with batched loads (one L2 round trip per + // eight slices); the addition order is fixed, so the result is + // bit-reproducible. + for (; z + 8 <= zc; z += 8) { + f4v v[8]; + #pragma unroll + for (int u = 0; u < 8; ++u) + v[u] = __builtin_nontemporal_load(p + (long)(z + u) * zs); + #pragma unroll + for (int u = 0; u < 8; ++u) acc += v[u]; + } + for (; z < zc; ++z) acc += __builtin_nontemporal_load(p + (long)z * zs); + + T* out = c + (long)(offset_m + m) * size_n + n; + half2 packed[2] = { + __halves2half2(__float2half_rn(acc[0]), __float2half_rn(acc[1])), + __halves2half2(__float2half_rn(acc[2]), __float2half_rn(acc[3]))}; + __builtin_memcpy(out, packed, sizeof(packed)); + } + if (t == 0) *(tickets + blockIdx.y * gridDim.x + blockIdx.x) = 0; + } + } } #else // non-RDNA3 device pass: empty __global__ for symbol parity. @@ -555,7 +739,8 @@ __global__ void gemm_q4_kernel_rdna3(const T* __restrict__ a, template __global__ void gemm_q4_kernel_rdna3(const T*, const uint32_t*, const uint32_t*, const T*, T*, const int, const int, - const int, const int, const int) {} + const int, const int, const int, float*, + int*) {} #endif // __HIP__RDNA3__ || !__HIP_DEVICE_COMPILE__ @@ -567,20 +752,28 @@ template void launch_gemm_q4_for_mcount(const T* a, const uint32_t* b_q_weight, const uint32_t* b_qzeros, const T* b_scales, T* c, int size_m, int size_n, int size_k, - int groups, int zero_offset, - cudaStream_t stream) { + int groups, int zero_offset, float* partials, + int* tickets, cudaStream_t stream) { dim3 block(THREADS_X); dim3 grid((size_n + BLOCK_KN_SIZE * 4 - 1) / (BLOCK_KN_SIZE * 4), (size_m + M_COUNT - 1) / M_COUNT, (size_k + BLOCK_KN_SIZE - 1) / BLOCK_KN_SIZE); - gemm_q4_kernel_rdna3 - <<>>(a, b_q_weight, b_qzeros, b_scales, c, size_m, - size_n, size_k, groups, zero_offset); + gemm_q4_kernel_rdna3<<>>( + a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups, + zero_offset, partials, tickets); +} + +// The M_COUNT template the launcher below picks for a given row count. +static inline int mcount_for_rows(int rows) { + if (rows == 1) return 1; + if (rows <= 3) return 2; + if (rows <= 7) return 4; + return 8; } // Dispatch to the largest M_COUNT template that doesn't waste more than -// half a tile. Caps at 8: above that, the WMMA-prefill kernel (M >= 16) is +// half a tile. Caps at 8: above that, the WMMA-prefill kernel (M >= 12) is // the right tool, not bigger M_COUNT in the scalar dot-product path. // // Tile-waste table: @@ -594,28 +787,130 @@ template void launch_gemm_q4(const T* a, const uint32_t* b_q_weight, const uint32_t* b_qzeros, const T* b_scales, T* c, int size_m, int size_n, int size_k, int groups, - bool use_v2_format, cudaStream_t stream) { + bool use_v2_format, float* partials, int* tickets, + cudaStream_t stream) { const int zero_offset = use_v2_format ? 0 : 1; + switch (mcount_for_rows(size_m)) { + case 1: + launch_gemm_q4_for_mcount(a, b_q_weight, b_qzeros, b_scales, c, + size_m, size_n, size_k, groups, + zero_offset, partials, tickets, stream); + break; + case 2: + launch_gemm_q4_for_mcount(a, b_q_weight, b_qzeros, b_scales, c, + size_m, size_n, size_k, groups, + zero_offset, partials, tickets, stream); + break; + case 4: + launch_gemm_q4_for_mcount(a, b_q_weight, b_qzeros, b_scales, c, + size_m, size_n, size_k, groups, + zero_offset, partials, tickets, stream); + break; + default: + // M_COUNT=8 covers the whole scalar domain (M <= 11 in practice: the + // public dispatch routes M >= 12 to WMMA); this branch is only reached + // when the caller falls through, and still produces correct output. + launch_gemm_q4_for_mcount(a, b_q_weight, b_qzeros, b_scales, c, + size_m, size_n, size_k, groups, + zero_offset, partials, tickets, stream); + break; + } +} - if (size_m == 1) { - launch_gemm_q4_for_mcount(a, b_q_weight, b_qzeros, b_scales, c, - size_m, size_n, size_k, groups, zero_offset, - stream); - } else if (size_m <= 3) { - launch_gemm_q4_for_mcount(a, b_q_weight, b_qzeros, b_scales, c, - size_m, size_n, size_k, groups, zero_offset, - stream); - } else if (size_m <= 7) { - launch_gemm_q4_for_mcount(a, b_q_weight, b_qzeros, b_scales, c, - size_m, size_n, size_k, groups, zero_offset, - stream); +// Deterministic split-K reduction for the bf16 launches: one thread per +// output element sums the grid.z FP32 partial slices in fixed ascending-z +// order and rounds to the output dtype exactly once. The order is a pure +// function of the launch shape, so the result is bit-reproducible. +template +__global__ void reduce_partials_rdna3(const float* __restrict__ partials, + T* __restrict__ c, const int z_count, + const int size_m, const int size_n) { + const long idx = (long)blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= (long)size_m * size_n) return; + const int m = (int)(idx / size_n); + const int n = (int)(idx % size_n); + float acc = 0.0f; + for (int z = 0; z < z_count; ++z) + acc += partials[((long)z * size_m + m) * size_n + n]; + if constexpr (std::is_same::value) { + c[idx] = __float2half_rn(acc); } else { - // M_COUNT=8 covers M up to 15 here; M >= 16 should ideally take the - // WMMA path, but if it falls through we still produce correct output — - // just leaving 3-5× of throughput on the table for prefill workloads. - launch_gemm_q4_for_mcount(a, b_q_weight, b_qzeros, b_scales, c, - size_m, size_n, size_k, groups, zero_offset, - stream); + c[idx] = __float2bfloat16(acc); + } +} + +// Deterministic scalar GEMM: split-K blocks write FP32 partials to scratch +// (no atomics, no intermediate low-precision rounding), and the last z +// block of each (x, y) tile performs the fixed ascending-z reduction +// in-kernel with a single final cast (see the kernel epilogue). The scalar +// kernel writes every in-range (z, m, n) partial exactly once, so a plain +// at::empty scratch suffices; row-tile bound: +// scratch_bytes = z_count * TILE_M * size_n * 4 +// is independent of the caller's M (the scalar domain is M < 64, so a single +// tile covers it). The PyTorch caching allocator (including its CUDA-graph +// capture pool) owns scratch reuse and lifetime. +// +// The arrival tickets are per-call at::zeros state: every launch of this +// call counts on tickets that start at zero (the kernel re-arms them for +// the next row tile), and concurrent invocations get distinct allocations +// from the stream-aware caching allocator, so no synchronization state is +// ever shared between invocations. When z_count == 1 the kernel's +// direct-store epilogue (gridDim.z == 1) is already deterministic, so the +// scratch and ticket allocations are skipped entirely. +template +void launch_gemm_q4_deterministic(const T* a, const uint32_t* b_q_weight, + const uint32_t* b_qzeros, const T* b_scales, + T* c, int size_m, int size_n, int size_k, + int groups, bool use_v2_format, + cudaStream_t stream) { + constexpr int TILE_M = 64; // single tile covers the scalar domain + const int z_count = (size_k + BLOCK_KN_SIZE - 1) / BLOCK_KN_SIZE; + if (z_count == 1) { + launch_gemm_q4(a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, + groups, use_v2_format, + /*partials=*/nullptr, /*tickets=*/nullptr, stream); + return; + } + at::TensorOptions dev_opts = at::TensorOptions().device( + at::Device(at::kCUDA, c10::cuda::current_device())); + at::Tensor partials = at::empty({z_count, std::min(TILE_M, size_m), size_n}, + dev_opts.dtype(at::kFloat)); + float* partials_ptr = partials.data_ptr(); + + int* tickets_ptr = nullptr; + if constexpr (std::is_same::value) { + // fp16: the reduction runs in the kernel (see the epilogue). One + // ticket per (x, y) tile, sized for the widest row-tile launch. + const int gx = (size_n + BLOCK_KN_SIZE * 4 - 1) / (BLOCK_KN_SIZE * 4); + int max_gy = 1; + for (int row0 = 0; row0 < size_m; row0 += TILE_M) { + const int rows = std::min(TILE_M, size_m - row0); + const int mc = mcount_for_rows(rows); + max_gy = std::max(max_gy, (rows + mc - 1) / mc); + } + at::Tensor tickets = + at::zeros({(long)gx * max_gy}, dev_opts.dtype(at::kInt)); + for (int row0 = 0; row0 < size_m; row0 += TILE_M) { + const int rows = std::min(TILE_M, size_m - row0); + launch_gemm_q4(a + (long)row0 * size_k, b_q_weight, b_qzeros, b_scales, + c + (long)row0 * size_n, rows, size_n, size_k, groups, + use_v2_format, partials_ptr, tickets.data_ptr(), + stream); + } + } else { + // bf16: separate fixed-order reduce (measured faster than the fused + // form on this path). + for (int row0 = 0; row0 < size_m; row0 += TILE_M) { + const int rows = std::min(TILE_M, size_m - row0); + launch_gemm_q4(a + (long)row0 * size_k, b_q_weight, b_qzeros, b_scales, + c + (long)row0 * size_n, rows, size_n, size_k, groups, + use_v2_format, partials_ptr, tickets_ptr, stream); + const long total = (long)rows * size_n; + const int threads = 256; + const int blocks = (int)((total + threads - 1) / threads); + reduce_partials_rdna3<<>>( + partials_ptr, c + (long)row0 * size_n, z_count, rows, size_n); + } } } @@ -640,13 +935,40 @@ torch::Tensor gptq_gemm_rdna3_wmma(torch::Tensor a, torch::Tensor b_q_weight, torch::Tensor b_qzeros, torch::Tensor b_scales, bool use_v2_format); +// WMMA dispatch threshold, from the measured scalar/WMMA crossover on +// gfx1100 (ratio = scalar_us / wmma_us, >1 means WMMA faster; median of +// 300 CUDA-event timings per arm, identical inputs, both ops called +// directly): +// +// K / N (dtype) M=1 M=4 M=8 M=11 M=12 M=16 M=24 M=32 +// 1536/5120 fp16 0.69 0.97 1.36 1.60 1.59 1.60 1.22 1.62 +// 4096/4096 fp16 0.46 0.69 1.01 1.30 1.30 1.32 1.05 1.40 +// 4352/5120 fp16 0.42 0.63 0.93 1.37 1.37 1.37 1.19 1.74 +// 5120/2048 fp16 0.44 0.62 0.87 1.00 1.00 1.01 1.15 1.34 +// 5120/8704 fp16 0.30 0.53 0.84 1.39 1.42 1.43 1.24 1.58 +// 1536/5120 bf16 1.07 1.23 1.50 1.71 1.69 1.00 1.00 1.00 +// 4096/4096 bf16 0.70 0.82 1.08 1.33 1.34 1.00 1.00 1.00 +// 4352/5120 bf16 0.62 0.74 0.98 1.40 1.40 1.00 1.00 1.00 +// 5120/2048 bf16 0.64 0.73 0.89 1.00 1.00 1.00 1.01 1.00 +// 5120/8704 bf16 0.45 0.65 0.92 1.46 1.49 1.00 1.00 1.00 +// +// At M >= 12 WMMA wins or ties on every shape for both dtypes (exact 1.00 +// cells and the one sub-1.0 sample, bf16 4096/4096 at M=24 = 0.996, are +// within measurement noise); below +// that the scalar path wins on most shapes (and is the more accurate of +// the two since the exact fp16 factored dequant, so it keeps the band). +// fp16 used to stay scalar until M=64 because the old baked-dequant +// scalar kernel was faster there — that speed came with ~100x more +// dequant error, and with the exact dequant the crossover moved down to +// the same 12 as bf16. A shape-aware (N/K) dispatch would do better +// still; that is tuning work beyond this fix. +constexpr int64_t WMMA_MIN_M = 12; + torch::Tensor gptq_gemm_rdna3(torch::Tensor a, torch::Tensor b_q_weight, torch::Tensor b_qzeros, torch::Tensor b_scales, bool use_v2_format) { if (a.dim() == 2 && b_q_weight.dim() == 2 && a.size(1) % 16 == 0 && - b_q_weight.size(1) % 16 == 0 && - ((a.scalar_type() == torch::kBFloat16 && a.size(0) >= 16) || - (a.scalar_type() == torch::kHalf && a.size(0) >= 64))) { + b_q_weight.size(1) % 16 == 0 && a.size(0) >= WMMA_MIN_M) { return gptq_gemm_rdna3_wmma(a, b_q_weight, b_qzeros, b_scales, use_v2_format); } @@ -673,22 +995,38 @@ torch::Tensor gptq_gemm_rdna3(torch::Tensor a, torch::Tensor b_q_weight, TORCH_CHECK(b_q_weight.size(0) * 8 == size_k, "b_q_weight first dim must be K/8"); + TORCH_CHECK(size_k % (groups * 32) == 0, + "group size (K/groups = ", size_k / groups, + ") must be a multiple of 32: the kernel checks group " + "transitions at 32-K granularity"); TORCH_CHECK(b_scales.size(0) == groups, "b_scales must have same group count as qzeros"); TORCH_CHECK(b_scales.size(1) == size_n, "b_scales last dim must be N"); - TORCH_CHECK(size_n % 8 == 0, "N must be a multiple of 8 (64-bit atomic CAS)"); + TORCH_CHECK(size_n % 8 == 0, + "N must be a multiple of 8 (packed qzeros layout: 8 4-bit zero " + "points per uint32 along N; also keeps the 4-column epilogue " + "store in bounds)"); auto opts = torch::TensorOptions().dtype(a.dtype()).device(a.device()); - at::Tensor c = torch::zeros({size_m, size_n}, opts); + // The deterministic epilogue writes every output element exactly once + // (reduce pass below — or the direct store when z_count == 1), so c needs + // no zero-initialization. This is a one-way door: the legacy CAS epilogue + // (see the kernel's store-mode comment) ADDS into c and would require + // torch::zeros here if it were ever re-selected. + at::Tensor c = torch::empty({size_m, size_n}, opts); + // Deterministic split-K: FP32 partials + fixed-order reduction (see + // launch_gemm_q4_deterministic). The CAS-atomic low-precision epilogue + // this replaces was order-dependent and produced different results for + // identical inputs on gfx11 once more than a few split blocks contended. if (a.scalar_type() == torch::kHalf) { - vllm::gptq_rdna3::launch_gemm_q4( + vllm::gptq_rdna3::launch_gemm_q4_deterministic( (const half*)a.data_ptr(), (const uint32_t*)b_q_weight.data_ptr(), (const uint32_t*)b_qzeros.data_ptr(), (const half*)b_scales.data_ptr(), (half*)c.data_ptr(), size_m, size_n, size_k, groups, use_v2_format, stream); } else { - vllm::gptq_rdna3::launch_gemm_q4( + vllm::gptq_rdna3::launch_gemm_q4_deterministic( (const vllm::gptq_rdna3::bf16_t*)a.data_ptr(), (const uint32_t*)b_q_weight.data_ptr(), (const uint32_t*)b_qzeros.data_ptr(), diff --git a/csrc/rocm/q_gemm_rdna3_wmma.cu b/csrc/rocm/q_gemm_rdna3_wmma.cu index 3876cd2225..116feeb04e 100644 --- a/csrc/rocm/q_gemm_rdna3_wmma.cu +++ b/csrc/rocm/q_gemm_rdna3_wmma.cu @@ -24,13 +24,15 @@ // of the 16x16 output, with 8 elements alternating M rows by lane_hi // (lanes 0..15 = even rows, lanes 16..31 = odd rows). See the layout // diagram on `gemm_q4_wmma_kernel_16x16_1w` below for the full mapping. -// * No native v_global_atomic_pk_add_{f16,bf16} on gfx11; the K-split -// epilogue (gridDim.z > 1) emulates packed atomic add via a CAS-32 -// retry loop on a uint32 word covering 2 fp16/bf16 lanes. Within a -// block we shuffle adjacent lanes via shfl_xor first so each pair of -// output cols goes through a single atomic — no intra-block -// contention. K_SPLIT == 1 keeps the original direct-write path with -// each block owning its 16M × 16N output tile. +// * No native v_global_atomic_pk_add_{f16,bf16} on gfx11. The split-K +// epilogue therefore stores FP32 per-split partials to scratch and a +// separate pass reduces them in fixed ascending-z order with a single +// final low-precision rounding — deterministic by construction +// (PR #54706). K_SPLIT == 1 keeps the original direct-write path with +// each block owning its 16M × 16N output tile. The legacy pre-#54706 +// design (a CAS-32 retry loop accumulating narrowed partials directly +// into c, order-dependent) is kept in the store path below the +// partials branch for A/B comparison only. #include @@ -182,16 +184,19 @@ __forceinline__ __device__ void dequant_4bit_8_bf16_to_bf16(uint32_t qa, } // --------------------------------------------------------------------------- -// Packed atomic-add helpers used by the K-split epilogue. +// Packed atomic-add helpers — LEGACY (pre-#54706) K-split epilogue. // -// When the kernel is launched with gridDim.z > 1, multiple K-segments -// accumulate into the same 16x16 output tile and need atomic write-back. -// gfx11 has no native v_global_atomic_pk_add_{f16,bf16}, so we issue a +// When the kernel is launched with partials == nullptr and gridDim.z > 1, +// multiple K-segments accumulate into the same 16x16 output tile and the +// legacy path atomically adds narrowed partials into a pre-zeroed c via a // CAS-loop on a 32-bit word covering 2 packed fp16/bf16 lanes. Within a // block the kernel pairs adjacent lanes via shfl_xor first, so each pair // of cols (n=lane_lo even, lane_lo+1) goes through a SINGLE atomic — no // intra-block contention on the same uint32 target. Inter-block // contention from gridDim.z (4-way at K_SPLIT=4) is the residual cost. +// The addition happens in bf16/fp16 AFTER narrowing, so the result depends +// on CAS completion order — the nondeterminism PR #54706 replaces. Kept +// for A/B comparison only; the shipped launchers never select it. // --------------------------------------------------------------------------- __forceinline__ __device__ void atomic_add_pk_f16(half2* addr, half2 val) { @@ -230,12 +235,15 @@ __forceinline__ __device__ void atomic_add_pk_bf16(bf162_t* addr, bf162_t val) { // K-split factor heuristic. Returns the gridDim.z to use for a given K. // Aim: each block does at least ~16 K-tiles (= K=256) so the per-block // constant overhead (LDS init, kernel prologue) is amortised. Upper -// bound K_SPLIT=4 to cap inter-block atomic contention to 4-way. +// bound K_SPLIT=4 — historically to cap inter-block atomic contention to +// 4-way under the legacy CAS epilogue; now it caps the FP32 scratch and +// reduce-pass traffic the deterministic epilogue adds. // // For typical Qwen-class shapes K ∈ {4096, 5120, 11008}, all return 4. -// Smaller K (e.g., embedding lookups) fall back to 1 (no split, no -// atomic). K must be divisible by (K_SPLIT × 16) for the split to be -// valid; the heuristic checks divisibility before raising the factor. +// Smaller K (e.g., embedding lookups) fall back to 1: no split, no +// scratch, direct store. K must be divisible by (K_SPLIT × 16) for the +// split to be valid; the heuristic checks divisibility before raising +// the factor. __host__ __device__ static inline int compute_wmma_k_split(int size_k) { if (size_k >= 1024 && size_k % 64 == 0) return 4; if (size_k >= 512 && size_k % 32 == 0) return 2; @@ -245,17 +253,18 @@ __host__ __device__ static inline int compute_wmma_k_split(int size_k) { // M-and-N-aware K-split heuristic for the v3/v4/v5 launchers. // // The original `compute_wmma_k_split` was K-only and always returns 4 for -// Qwen-class K, which over-subscribes wave slots and pays the atomic CAS -// epilogue once-per-K-segment per output cell. With v3/v4/v5's larger +// Qwen-class K, which over-subscribes wave slots and pays the split-K +// epilogue (historically the atomic CAS; now FP32 scratch + reduce) once +// per K-segment per output cell. With v3/v4/v5's larger // tiles (64M × 16/32/64N) and 4 resident waves per block, the no-split // grid is often already well-saturated on gfx1100's 96 CUs / 3072 wave -// slots — adding gridDim.z just adds atomic overhead. +// slots — adding gridDim.z just adds epilogue overhead. // // Heuristic: compute the no-split block count gridDim.x × gridDim.y, then // pick the smallest K_SPLIT that brings total waves to at least // ~2× over-subscription (~6000 waves for our 3072 slots, i.e. 1500 blocks -// at 4 waves/block). Above that threshold, K_SPLIT=1 — direct write, no -// atomic. +// at 4 waves/block). Above that threshold, K_SPLIT=1: no scratch, no +// reduce pass, direct store. // // Args: // size_m, size_n, size_k — GEMM dims @@ -373,7 +382,11 @@ __global__ void gemm_q4_wmma_kernel_16x16_1w( const T* __restrict__ a, const uint32_t* __restrict__ b_q, const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales, T* __restrict__ c, const int size_m, const int size_n, const int size_k, - const int groups, const int zero_offset) { + const int groups, const int zero_offset, + // Deterministic split-K: non-null => each block stores its FP32 partial + // to partials[(z*size_m + m)*size_n + n]; a separate pass then reduces + // the z-slices in fixed order with a single low-precision rounding. + float* __restrict__ partials) { using E = typename WmmaNative::elem; using V16 = typename WmmaNative::v16; @@ -525,7 +538,21 @@ __global__ void gemm_q4_wmma_kernel_16x16_1w( // lane_hi == 0 → rows 0, 2, 4, ..., 14 (even rows) // lane_hi == 1 → rows 1, 3, 5, ..., 15 (odd rows) // c_acc[i] corresponds to actual row m = 2*i + lane_hi at column lane_lo. - if (gridDim.z > 1) { + if (partials != nullptr) { + // Deterministic split-K: plain FP32 store per (z, m, n); the reduce + // pass fixes the accumulation order. + const int out_n = n_tile + lane_lo; + if (out_n < size_n) { + #pragma unroll + for (int i = 0; i < 8; i++) { + const int out_m = m_tile + 2 * i + lane_hi; + if (out_m < size_m) { + partials[((long)blockIdx.z * size_m + out_m) * size_n + out_n] = + c_acc[i]; + } + } + } + } else if (gridDim.z > 1) { // K-split path: 4 (or whatever the split factor is) K-segments per // output cell contend → atomic accumulation. Caller has zero-init'd c. // @@ -588,9 +615,78 @@ template __global__ void gemm_q4_wmma_kernel_16x16_1w(const T*, const uint32_t*, const uint32_t*, const T*, T*, const int, const int, const int, - const int, const int) {} + const int, const int, float*) {} #endif +// Deterministic split-K reduction for the WMMA path: one thread per output +// element sums the grid.z FP32 partial slices in fixed ascending-z order and +// rounds to the output dtype exactly once. Order is a pure function of the +// launch shape, so the result is bit-reproducible for identical inputs. +// (Kernel bodies live in q_gemm_rdna3_wmma.cu alongside their users; HIP +// kernel templates are not shared across translation units here.) +template +__global__ void reduce_partials_wmma(const float* __restrict__ partials, + T* __restrict__ c, const int z_count, + const int size_m, const int size_n) { + const long idx = (long)blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= (long)size_m * size_n) return; + const int m = (int)(idx / size_n); + const int n = (int)(idx % size_n); + float acc = 0.0f; + for (int z = 0; z < z_count; ++z) + acc += partials[((long)z * size_m + m) * size_n + n]; + if constexpr (std::is_same::value) { + c[idx] = __float2half_rn(acc); + } else { + c[idx] = __float2bfloat16(acc); + } +} + +// FP32 split-partial scratch for one deterministic WMMA launch. +// +// Coverage invariant — why at::empty is safe: every scratch slot that the +// fixed-order reducer reads, i.e. every (z, m, n) with z < k_split, +// m < size_m, n < size_n, has EXACTLY ONE writer: +// * z: grid.z == k_split and each z-block stores slice blockIdx.z; +// * m: grid.y tiles M by the kernel's M_TILE; within a tile, wave w owns +// rows [16w, 16w+16) and (i, lane_hi) with out_m = m_tile + 16w + 2i + +// lane_hi enumerates the 16 rows of the wave uniquely (i in [0,8), +// lane_hi in {0,1}); +// * n: grid.x tiles N by N_TILE; each 16-wide accumulator slice stores +// columns n_base + lane_lo with lane_lo in [0,16) unique per column. +// The out_m >= size_m / out_n >= size_n guards in the store paths skip +// only slots the reducer never reads (it iterates m < size_m, +// n < size_n), and the store index ((z * size_m + m) * size_n + n) stays +// inside [0, k_split * size_m * size_n) — no out-of-bounds writes either. +// This holds for every WMMA kernel variant in this TU (16x16_1w, 32x16_2w, +// 64x16_4w, 64x32_4w, 64x64_4w, 128x64_k16, 128x64_k32) because they all +// share the wave/lane output mapping above; the tiled launchers pass +// size_m = rows of the current row tile so the same argument applies per +// tile. A stale at::empty scratch therefore never reaches the reducer. +// +// One-way door: the legacy CAS epilogue (partials == nullptr && +// gridDim.z > 1, kept for A/B comparison) accumulates into the OUTPUT +// instead and would require a zero-initialized c if re-selected. +static inline at::Tensor alloc_wmma_partials(int k_split, int size_m, + int size_n) { + return at::empty( + {k_split, size_m, size_n}, + at::TensorOptions() + .dtype(at::kFloat) + .device(at::Device(at::kCUDA, c10::cuda::current_device()))); +} + +template +static inline void launch_wmma_reduce(const at::Tensor& partials, T* c, + int k_split, int size_m, int size_n, + cudaStream_t stream) { + const long total = (long)size_m * size_n; + const int threads = 256; + const int blocks = (int)((total + threads - 1) / threads); + reduce_partials_wmma<<>>( + partials.data_ptr(), c, k_split, size_m, size_n); +} + template void launch_gemm_q4_wmma_16x16_1w(const T* a, const uint32_t* b_q_weight, const uint32_t* b_qzeros, const T* b_scales, @@ -599,13 +695,24 @@ void launch_gemm_q4_wmma_16x16_1w(const T* a, const uint32_t* b_q_weight, cudaStream_t stream) { // 1 wave per block (32 lanes), 16x16 C tile per block. gridDim.z splits // K so that more blocks (and therefore more waves) are in flight; with - // K_SPLIT > 1 the kernel switches to atomic write-back at the epilogue. + // K_SPLIT > 1 the kernel stores FP32 partials and a reduce pass fixes + // the accumulation order. const int k_split = compute_wmma_k_split(size_k); dim3 block(32); dim3 grid((size_n + 15) / 16, (size_m + 15) / 16, k_split); - gemm_q4_wmma_kernel_16x16_1w - <<>>(a, b_q_weight, b_qzeros, b_scales, c, size_m, - size_n, size_k, groups, zero_offset); + // k_split == 1: single writer per cell, the kernel's direct-store path is + // already deterministic. k_split > 1: FP32 partials + fixed-order reduce. + at::Tensor partials; + float* partials_ptr = nullptr; + if (k_split > 1) { + partials = alloc_wmma_partials(k_split, size_m, size_n); + partials_ptr = partials.data_ptr(); + } + gemm_q4_wmma_kernel_16x16_1w<<>>( + a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups, + zero_offset, partials_ptr); + if (k_split > 1) + launch_wmma_reduce(partials, c, k_split, size_m, size_n, stream); } #if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__) @@ -647,7 +754,11 @@ __global__ void gemm_q4_wmma_kernel_32x16_2w( const T* __restrict__ a, const uint32_t* __restrict__ b_q, const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales, T* __restrict__ c, const int size_m, const int size_n, const int size_k, - const int groups, const int zero_offset) { + const int groups, const int zero_offset, + // Deterministic split-K: non-null => each block stores its FP32 partial + // to partials[(z*size_m + m)*size_n + n]; a separate pass then reduces + // the z-slices in fixed order with a single low-precision rounding. + float* __restrict__ partials) { using E = typename WmmaNative::elem; using V16 = typename WmmaNative::v16; @@ -775,7 +886,21 @@ __global__ void gemm_q4_wmma_kernel_32x16_2w( // Each wave owns rows [m_tile + wave_id*16 .. + 16) of the output tile. const int m_tile_wave = m_tile + wave_id * 16; - if (gridDim.z > 1) { + if (partials != nullptr) { + // Deterministic split-K: plain FP32 store per (z, m, n); the reduce + // pass fixes the accumulation order. + const int out_n = n_tile + lane_lo; + if (out_n < size_n) { + #pragma unroll + for (int i = 0; i < 8; i++) { + const int out_m = m_tile_wave + 2 * i + lane_hi; + if (out_m < size_m) { + partials[((long)blockIdx.z * size_m + out_m) * size_n + out_n] = + c_acc[i]; + } + } + } + } else if (gridDim.z > 1) { // K-split atomic path. Pair-shuffle within wave to halve atomic count. // shfl_xor here is wave-local (wave32 semantics) so each wave does its // own pairing — the two waves don't interact during the store. @@ -826,7 +951,7 @@ template __global__ void gemm_q4_wmma_kernel_32x16_2w(const T*, const uint32_t*, const uint32_t*, const T*, T*, const int, const int, const int, - const int, const int) {} + const int, const int, float*) {} #endif template @@ -856,9 +981,17 @@ void launch_gemm_q4_wmma_32x16_2w(const T* a, const uint32_t* b_q_weight, const int k_split = compute_wmma_k_split(size_k); dim3 block(64); dim3 grid((size_n + 15) / 16, (size_m + 31) / 32, k_split); - gemm_q4_wmma_kernel_32x16_2w - <<>>(a, b_q_weight, b_qzeros, b_scales, c, size_m, - size_n, size_k, groups, zero_offset); + at::Tensor partials; + float* partials_ptr = nullptr; + if (k_split > 1) { + partials = alloc_wmma_partials(k_split, size_m, size_n); + partials_ptr = partials.data_ptr(); + } + gemm_q4_wmma_kernel_32x16_2w<<>>( + a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups, + zero_offset, partials_ptr); + if (k_split > 1) + launch_wmma_reduce(partials, c, k_split, size_m, size_n, stream); } #if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__) @@ -898,7 +1031,11 @@ __global__ void gemm_q4_wmma_kernel_64x16_4w( const T* __restrict__ a, const uint32_t* __restrict__ b_q, const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales, T* __restrict__ c, const int size_m, const int size_n, const int size_k, - const int groups, const int zero_offset) { + const int groups, const int zero_offset, + // Deterministic split-K: non-null => each block stores its FP32 partial + // to partials[(z*size_m + m)*size_n + n]; a separate pass then reduces + // the z-slices in fixed order with a single low-precision rounding. + float* __restrict__ partials) { using E = typename WmmaNative::elem; using V16 = typename WmmaNative::v16; @@ -1024,7 +1161,21 @@ __global__ void gemm_q4_wmma_kernel_64x16_4w( // ---- Store C ---- Each wave owns 16 M-rows of the output tile. const int m_tile_wave = m_tile + wave_id * 16; - if (gridDim.z > 1) { + if (partials != nullptr) { + // Deterministic split-K: plain FP32 store per (z, m, n); the reduce + // pass fixes the accumulation order. + const int out_n = n_tile + lane_lo; + if (out_n < size_n) { + #pragma unroll + for (int i = 0; i < 8; i++) { + const int out_m = m_tile_wave + 2 * i + lane_hi; + if (out_m < size_m) { + partials[((long)blockIdx.z * size_m + out_m) * size_n + out_n] = + c_acc[i]; + } + } + } + } else if (gridDim.z > 1) { // K-split atomic path. Pair-shuffle within wave to halve atomic count. const bool is_even_lane = (lane_lo & 1) == 0; const int out_n_pair = n_tile + lane_lo; @@ -1073,7 +1224,7 @@ template __global__ void gemm_q4_wmma_kernel_64x16_4w(const T*, const uint32_t*, const uint32_t*, const T*, T*, const int, const int, const int, - const int, const int) {} + const int, const int, float*) {} #endif template @@ -1094,9 +1245,17 @@ void launch_gemm_q4_wmma_64x16_4w(const T* a, const uint32_t* b_q_weight, const int k_split = compute_wmma_k_split_mn(size_m, size_n, size_k, 64, 16); dim3 block(128); dim3 grid((size_n + 15) / 16, (size_m + 63) / 64, k_split); - gemm_q4_wmma_kernel_64x16_4w - <<>>(a, b_q_weight, b_qzeros, b_scales, c, size_m, - size_n, size_k, groups, zero_offset); + at::Tensor partials; + float* partials_ptr = nullptr; + if (k_split > 1) { + partials = alloc_wmma_partials(k_split, size_m, size_n); + partials_ptr = partials.data_ptr(); + } + gemm_q4_wmma_kernel_64x16_4w<<>>( + a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups, + zero_offset, partials_ptr); + if (k_split > 1) + launch_wmma_reduce(partials, c, k_split, size_m, size_n, stream); } #if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__) @@ -1130,7 +1289,11 @@ __global__ void gemm_q4_wmma_kernel_64x32_4w( const T* __restrict__ a, const uint32_t* __restrict__ b_q, const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales, T* __restrict__ c, const int size_m, const int size_n, const int size_k, - const int groups, const int zero_offset) { + const int groups, const int zero_offset, + // Deterministic split-K: non-null => each block stores its FP32 partial + // to partials[(z*size_m + m)*size_n + n]; a separate pass then reduces + // the z-slices in fixed order with a single low-precision rounding. + float* __restrict__ partials) { using E = typename WmmaNative::elem; using V16 = typename WmmaNative::v16; @@ -1266,6 +1429,22 @@ __global__ void gemm_q4_wmma_kernel_64x32_4w( // Helper: store one v8fp32 accumulator's 8 outputs (covers 16 N-cols at // n_base via lane_lo + interleaved M rows m_tile_wave + 2i + lane_hi). auto store_acc = [&](const v8fp32& acc, int n_base) { + if (partials != nullptr) { + // Deterministic split-K: plain FP32 store per (z, m, n). Each element + // has one writer per z-slice and the reduce pass fixes the + // accumulation order, so the result is bit-reproducible. + const int out_n = n_base + lane_lo; + if (out_n >= size_n) return; + #pragma unroll + for (int i = 0; i < 8; i++) { + const int out_m = m_tile_wave + 2 * i + lane_hi; + if (out_m < size_m) { + partials[((long)blockIdx.z * size_m + out_m) * size_n + out_n] = + acc[i]; + } + } + return; + } if (gridDim.z > 1) { const bool is_even_lane = (lane_lo & 1) == 0; const int out_n_pair = n_base + lane_lo; @@ -1316,7 +1495,7 @@ template __global__ void gemm_q4_wmma_kernel_64x32_4w(const T*, const uint32_t*, const uint32_t*, const T*, T*, const int, const int, const int, - const int, const int) {} + const int, const int, float*) {} #endif template @@ -1339,9 +1518,17 @@ void launch_gemm_q4_wmma_64x32_4w(const T* a, const uint32_t* b_q_weight, const int k_split = compute_wmma_k_split_mn(size_m, size_n, size_k, 64, 32); dim3 block(128); dim3 grid((size_n + 31) / 32, (size_m + 63) / 64, k_split); - gemm_q4_wmma_kernel_64x32_4w - <<>>(a, b_q_weight, b_qzeros, b_scales, c, size_m, - size_n, size_k, groups, zero_offset); + at::Tensor partials; + float* partials_ptr = nullptr; + if (k_split > 1) { + partials = alloc_wmma_partials(k_split, size_m, size_n); + partials_ptr = partials.data_ptr(); + } + gemm_q4_wmma_kernel_64x32_4w<<>>( + a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups, + zero_offset, partials_ptr); + if (k_split > 1) + launch_wmma_reduce(partials, c, k_split, size_m, size_n, stream); } #if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__) @@ -1375,7 +1562,11 @@ __global__ void gemm_q4_wmma_kernel_64x64_4w( const T* __restrict__ a, const uint32_t* __restrict__ b_q, const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales, T* __restrict__ c, const int size_m, const int size_n, const int size_k, - const int groups, const int zero_offset) { + const int groups, const int zero_offset, + // Deterministic split-K: non-null => each block stores its FP32 partial + // to partials[(z*size_m + m)*size_n + n]; a separate pass then reduces + // the z-slices in fixed order with a single low-precision rounding. + float* __restrict__ partials) { using E = typename WmmaNative::elem; using V16 = typename WmmaNative::v16; @@ -1502,6 +1693,22 @@ __global__ void gemm_q4_wmma_kernel_64x64_4w( // ---- Store C ---- Each wave owns 16M × 64N. Helper writes one acc slice. const int m_tile_wave = m_tile + wave_id * 16; auto store_acc = [&](const v8fp32& acc, int n_base) { + if (partials != nullptr) { + // Deterministic split-K: plain FP32 store per (z, m, n). Each element + // has one writer per z-slice and the reduce pass fixes the + // accumulation order, so the result is bit-reproducible. + const int out_n = n_base + lane_lo; + if (out_n >= size_n) return; + #pragma unroll + for (int i = 0; i < 8; i++) { + const int out_m = m_tile_wave + 2 * i + lane_hi; + if (out_m < size_m) { + partials[((long)blockIdx.z * size_m + out_m) * size_n + out_n] = + acc[i]; + } + } + return; + } if (gridDim.z > 1) { const bool is_even_lane = (lane_lo & 1) == 0; const int out_n_pair = n_base + lane_lo; @@ -1566,7 +1773,11 @@ __global__ void gemm_q4_wmma_kernel_128x64_k16( const T* __restrict__ a, const uint32_t* __restrict__ b_q, const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales, T* __restrict__ c, const int size_m, const int size_n, const int size_k, - const int groups, const int zero_offset) { + const int groups, const int zero_offset, + // Deterministic split-K: non-null => each block stores its FP32 partial + // to partials[(z*size_m + m)*size_n + n]; a separate pass then reduces + // the z-slices in fixed order with a single low-precision rounding. + float* __restrict__ partials) { using E = typename WmmaNative::elem; using V16 = typename WmmaNative::v16; @@ -1695,6 +1906,22 @@ __global__ void gemm_q4_wmma_kernel_128x64_k16( // ---- Store C ---- Each wave owns 16M × 64N. const int m_tile_wave = m_tile + wave_id * 16; auto store_acc = [&](const v8fp32& acc, int n_base) { + if (partials != nullptr) { + // Deterministic split-K: plain FP32 store per (z, m, n). Each element + // has one writer per z-slice and the reduce pass fixes the + // accumulation order, so the result is bit-reproducible. + const int out_n = n_base + lane_lo; + if (out_n >= size_n) return; + #pragma unroll + for (int i = 0; i < 8; i++) { + const int out_m = m_tile_wave + 2 * i + lane_hi; + if (out_m < size_m) { + partials[((long)blockIdx.z * size_m + out_m) * size_n + out_n] = + acc[i]; + } + } + return; + } if (gridDim.z > 1) { const bool is_even_lane = (lane_lo & 1) == 0; const int out_n_pair = n_base + lane_lo; @@ -1764,7 +1991,11 @@ __global__ void gemm_q4_wmma_kernel_128x64_k32( const T* __restrict__ a, const uint32_t* __restrict__ b_q, const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales, T* __restrict__ c, const int size_m, const int size_n, const int size_k, - const int groups, const int zero_offset) { + const int groups, const int zero_offset, + // Deterministic split-K: non-null => each block stores its FP32 partial + // to partials[(z*size_m + m)*size_n + n]; a separate pass then reduces + // the z-slices in fixed order with a single low-precision rounding. + float* __restrict__ partials) { using E = typename WmmaNative::elem; using V16 = typename WmmaNative::v16; @@ -1907,6 +2138,22 @@ __global__ void gemm_q4_wmma_kernel_128x64_k32( // ---- Store C ---- Same as V7. const int m_tile_wave = m_tile + wave_id * 16; auto store_acc = [&](const v8fp32& acc, int n_base) { + if (partials != nullptr) { + // Deterministic split-K: plain FP32 store per (z, m, n). Each element + // has one writer per z-slice and the reduce pass fixes the + // accumulation order, so the result is bit-reproducible. + const int out_n = n_base + lane_lo; + if (out_n >= size_n) return; + #pragma unroll + for (int i = 0; i < 8; i++) { + const int out_m = m_tile_wave + 2 * i + lane_hi; + if (out_m < size_m) { + partials[((long)blockIdx.z * size_m + out_m) * size_n + out_n] = + acc[i]; + } + } + return; + } if (gridDim.z > 1) { const bool is_even_lane = (lane_lo & 1) == 0; const int out_n_pair = n_base + lane_lo; @@ -1960,17 +2207,17 @@ template __global__ void gemm_q4_wmma_kernel_64x64_4w(const T*, const uint32_t*, const uint32_t*, const T*, T*, const int, const int, const int, - const int, const int) {} + const int, const int, float*) {} template __global__ void gemm_q4_wmma_kernel_128x64_k16(const T*, const uint32_t*, const uint32_t*, const T*, T*, const int, const int, const int, - const int, const int) {} + const int, const int, float*) {} template __global__ void gemm_q4_wmma_kernel_128x64_k32(const T*, const uint32_t*, const uint32_t*, const T*, T*, const int, const int, const int, - const int, const int) {} + const int, const int, float*) {} #endif template @@ -1994,15 +2241,46 @@ void launch_gemm_q4_wmma_64x64_4w(const T* a, const uint32_t* b_q_weight, compute_wmma_k_split_mn(size_m, size_n, size_k, 128, 64); const int groupsize = size_k / groups; dim3 block(256); - dim3 grid((size_n + 63) / 64, (size_m + 127) / 128, k_split); - if (size_k % 32 == 0 && groupsize >= 32 && (size_k / k_split) % 32 == 0) { - gemm_q4_wmma_kernel_128x64_k32<<>>( - a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups, - zero_offset); - } else { - gemm_q4_wmma_kernel_128x64_k16<<>>( - a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups, - zero_offset); + if (k_split == 1) { + // Single writer per cell: the kernel's direct-store path is already + // deterministic — no scratch, no reduce pass (same invariant as the + // 64x64 path below). No row tiling either: tiling exists only to + // bound the FP32 scratch, which is not allocated here; grid.y tiles + // M by 128 internally for any size_m. + dim3 grid((size_n + 63) / 64, (size_m + 127) / 128, 1); + if (size_k % 32 == 0 && groupsize >= 32) { + gemm_q4_wmma_kernel_128x64_k32<<>>( + a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, + groups, zero_offset, /*partials=*/nullptr); + } else { + gemm_q4_wmma_kernel_128x64_k16<<>>( + a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, + groups, zero_offset, /*partials=*/nullptr); + } + return; + } + // Deterministic split-K, row-tiled so the FP32 scratch bound + // scratch_bytes = k_split * TILE_M * size_n * 4 + // stays independent of the caller's M (k_split <= 4 here). + constexpr int TILE_M = 512; + at::Tensor partials = + alloc_wmma_partials(k_split, std::min(TILE_M, size_m), size_n); + float* partials_ptr = partials.data_ptr(); + for (int row0 = 0; row0 < size_m; row0 += TILE_M) { + const int rows = std::min(TILE_M, size_m - row0); + const T* a_t = a + (long)row0 * size_k; + T* c_t = c + (long)row0 * size_n; + dim3 grid((size_n + 63) / 64, (rows + 127) / 128, k_split); + if (size_k % 32 == 0 && groupsize >= 32 && (size_k / k_split) % 32 == 0) { + gemm_q4_wmma_kernel_128x64_k32<<>>( + a_t, b_q_weight, b_qzeros, b_scales, c_t, rows, size_n, size_k, + groups, zero_offset, partials_ptr); + } else { + gemm_q4_wmma_kernel_128x64_k16<<>>( + a_t, b_q_weight, b_qzeros, b_scales, c_t, rows, size_n, size_k, + groups, zero_offset, partials_ptr); + } + launch_wmma_reduce(partials, c_t, k_split, rows, size_n, stream); } return; } @@ -2010,10 +2288,29 @@ void launch_gemm_q4_wmma_64x64_4w(const T* a, const uint32_t* b_q_weight, // 4 waves per block (128 threads), 64M × 64N tile per block. const int k_split = compute_wmma_k_split_mn(size_m, size_n, size_k, 64, 64); dim3 block(128); - dim3 grid((size_n + 63) / 64, (size_m + 63) / 64, k_split); - gemm_q4_wmma_kernel_64x64_4w - <<>>(a, b_q_weight, b_qzeros, b_scales, c, size_m, - size_n, size_k, groups, zero_offset); + if (k_split == 1) { + // Single writer per cell: the kernel's direct-store path is already + // deterministic — no scratch, no reduce pass. + dim3 grid((size_n + 63) / 64, (size_m + 63) / 64, 1); + gemm_q4_wmma_kernel_64x64_4w<<>>( + a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups, + zero_offset, /*partials=*/nullptr); + return; + } + constexpr int TILE_M = 512; + at::Tensor partials = + alloc_wmma_partials(k_split, std::min(TILE_M, size_m), size_n); + float* partials_ptr = partials.data_ptr(); + for (int row0 = 0; row0 < size_m; row0 += TILE_M) { + const int rows = std::min(TILE_M, size_m - row0); + const T* a_t = a + (long)row0 * size_k; + T* c_t = c + (long)row0 * size_n; + dim3 grid((size_n + 63) / 64, (rows + 63) / 64, k_split); + gemm_q4_wmma_kernel_64x64_4w<<>>( + a_t, b_q_weight, b_qzeros, b_scales, c_t, rows, size_n, size_k, groups, + zero_offset, partials_ptr); + launch_wmma_reduce(partials, c_t, k_split, rows, size_n, stream); + } } } // namespace gptq_rdna3_wmma @@ -2070,11 +2367,13 @@ torch::Tensor gptq_gemm_rdna3_wmma(torch::Tensor a, torch::Tensor b_q_weight, TORCH_CHECK(size_k % 16 == 0, "WMMA path requires K % 16 == 0"); auto opts = torch::TensorOptions().dtype(a.dtype()).device(a.device()); - // Always zero-init the output: some V3-V8 boundary threads may exit - // without writing their output cell (e.g. out_m >= size_m), leaving - // uninitialized garbage when torch::empty is used. The cost is - // negligible (< 1.5% of prefill time on gfx1100). - at::Tensor c = torch::zeros({size_m, size_n}, opts); + // Every path writes each output element exactly once: direct store when + // k_split == 1, otherwise the fixed-order FP32 reduce over per-split + // partials (whose scratch coverage invariant is documented at + // alloc_wmma_partials). c therefore needs no zero-initialization. + // One-way door: the legacy CAS epilogue ADDS into c and would require + // torch::zeros here if it were ever re-selected. + at::Tensor c = torch::empty({size_m, size_n}, opts); const int zero_offset = use_v2_format ? 0 : 1; diff --git a/csrc/rocm/qdq_4_rdna3.cuh b/csrc/rocm/qdq_4_rdna3.cuh index 43a7b33ba1..237782bb35 100644 --- a/csrc/rocm/qdq_4_rdna3.cuh +++ b/csrc/rocm/qdq_4_rdna3.cuh @@ -52,61 +52,75 @@ __forceinline__ __device__ void shuffle_4bit_8(uint32_t* q) { } // --------------------------------------------------------------------------- -// fp16 path +// fp16 path: exact factored form. +// --------------------------------------------------------------------------- +// The pre-#54706-v2 kernels baked the zero/scale bias into every weight, +// which forced z = scale * (-1024 - zero) to be *stored* as fp16. At +// scale ~0.02 that costs ~0.008 abs per (group, column) constant and the +// error then accumulates along the K/groups axis (measured max abs +// error ~1.0 vs an FP32 dequantized reference at K=4096). Computing the +// product in fp32 does not help: the error is in the fp16 *representation* +// of z, not in the multiply. +// +// Instead use the identity, exact for any grouping: +// +// sum_k a_k * w_k = y * sum_k a_k * (1024 + q_k) + z * sum_k a_k +// +// with one (y, z) pair per nibble slot (the high pairs carry a factor of +// 16): y_lo = s, z_lo = s*(-1024 - zero), y_hi = s/16, z_hi = s*(-64 - +// zero). The y/z constants live in fp32 and are never narrowed; the magic +// values (1024 + q) / (1024 + 16q) stay in fp16 where they are exact. The +// z * sum_a correction is applied once per 32-K round in the kernel (z is +// constant within a group, so per-round applications sum to the per-group +// application). sum_k a_k does not depend on the output column, so the +// extra v_dot2 against half2(1,1) is shared by all four columns a thread +// owns. This is the same structure the bf16 branches of q_gemm_rdna3.cu +// already use; fp16 needs two (y, z) pairs and two activation sums because +// its magic values carry the upper-nibble *16 trick. // --------------------------------------------------------------------------- -// Precompute scale-baked constants for a single zero/scale pair. -// z1z16[0] = scale * (-1024 - zero) (used for "low" pairs) -// z1z16[1] = scale * (-64 - zero) (used for "high" pairs) -// y1y16[0] = scale * 1 (low pairs are q + 1024) -// y1y16[1] = scale * (1/16) (high pairs are q*16 + 1024) -__forceinline__ __device__ void prep_zero_scale_fp16(uint32_t zero, half scale, - half2 (&z1z16)[2], - half2 (&y1y16)[2]) { - // half(-1024 - zero) via the exllamav2 bit-trick: - // half bits 0xE400 == -1024.0 ; ORing the zero into mantissa subtracts it. - union { - uint16_t u; - half h; - } z1u; - z1u.u = (uint16_t)(0xE400 | zero); - half z1 = z1u.h; - half z16 = __hsub(__int2half_rn(-64), __int2half_rn((int)zero)); - - half2 scale2 = __half2half2(scale); - z1z16[0] = __hmul2(scale2, __half2half2(z1)); - z1z16[1] = __hmul2(scale2, __half2half2(z16)); - - half y1 = __float2half_rn(1.0f); - half y16 = __float2half_rn(1.0f / 16.0f); - y1y16[0] = __hmul2(scale2, __half2half2(y1)); - y1y16[1] = __hmul2(scale2, __half2half2(y16)); +// fp32 (y, z) pairs for the exact factored form above. +__forceinline__ __device__ void prep_zero_scale_fp16_f32(uint32_t zero, + half scale, + float (&y)[2], + float (&z)[2]) { + const float s = __half2float(scale); + const float zf = (float)(int)zero; + y[0] = s; // low pairs: (1024 + q) * s + z[0] = s * (-1024.0f - zf); + y[1] = s * (1.0f / 16.0f); // high pairs: (1024 + 16q) * s/16 + z[1] = s * (-64.0f - zf); } -// Dequantize one int32 (8 shuffled 4-bit weights) into 4 half2 pairs: -// dq[0] = (q[0], q[1]) * scale - zero*scale -// dq[1] = (q[2], q[3]) * scale - zero*scale -// dq[2] = (q[4], q[5]) * scale - zero*scale -// dq[3] = (q[6], q[7]) * scale - zero*scale -__forceinline__ __device__ void dequant_4bit_8_fp16(uint32_t qa, half2 (&dq)[4], - half2 (&z1z16)[2], - half2 (&y1y16)[2]) { +// Magic values only, with no y/z folded in: low pairs land in [0] and [2], +// high pairs in [1] and [3]. The lane mapping assumes the host-side +// gptq_shuffle interleave (see shuffle_4bit_8 above): even elements sit in +// the low nibbles, odd elements in the high nibbles. +__forceinline__ __device__ void magic_4bit_8_fp16(uint32_t qa, half2 (&qm)[4]) { const uint32_t c0 = 0x64006400; - union { uint32_t u; half2 h2; - } q0, q1, q2, q3; - q0.u = (qa & 0x000F000F) | c0; // half2(q[0]+1024, q[1]+1024) - q1.u = (qa & 0x00F000F0) | c0; // half2(q[2]*16+1024, q[3]*16+1024) - uint32_t qa_hi = qa >> 8; - q2.u = (qa_hi & 0x000F000F) | c0; // half2(q[4]+1024, q[5]+1024) - q3.u = (qa_hi & 0x00F000F0) | c0; // half2(q[6]*16+1024, q[7]*16+1024) + } t; + t.u = (qa & 0x000F000F) | c0; // half2(1024 + q[0], 1024 + q[1]) + qm[0] = t.h2; + t.u = (qa & 0x00F000F0) | c0; // half2(1024 + q[2]*16, 1024 + q[3]*16) + qm[1] = t.h2; + const uint32_t qa_hi = qa >> 8; + t.u = (qa_hi & 0x000F000F) | c0; // half2(1024 + q[4], 1024 + q[5]) + qm[2] = t.h2; + t.u = (qa_hi & 0x00F000F0) | c0; // half2(1024 + q[6]*16, 1024 + q[7]*16) + qm[3] = t.h2; +} - dq[0] = __hfma2(q0.h2, y1y16[0], z1z16[0]); - dq[1] = __hfma2(q1.h2, y1y16[1], z1z16[1]); - dq[2] = __hfma2(q2.h2, y1y16[0], z1z16[0]); - dq[3] = __hfma2(q3.h2, y1y16[1], z1z16[1]); +// half2(1.0, 1.0) for the activation-sum v_dot2 in the exact factored form. +__forceinline__ __device__ half2 ones_half2_fp16() { + union { + uint32_t u; + half2 h2; + } t; + t.u = 0x3C003C00u; + return t.h2; } // --------------------------------------------------------------------------- diff --git a/tests/kernels/quantization/test_rdna3_w4a16_determinism.py b/tests/kernels/quantization/test_rdna3_w4a16_determinism.py new file mode 100644 index 0000000000..89b3130fba --- /dev/null +++ b/tests/kernels/quantization/test_rdna3_w4a16_determinism.py @@ -0,0 +1,585 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Determinism + FP32-reference correctness tests for the RDNA3 W4A16 GEMM. + +The split-K epilogues of ``gptq_gemm_rdna3`` (scalar path) and the WMMA +path previously accumulated per-split low-precision partials directly into +the output via CAS atomics. With multiple writers per output element the +accumulation order — and therefore the rounded result — varied per call: +repeated fixed-input calls produced a different output almost every time on +gfx1100 once more than a few split blocks contended. + +Both paths now store FP32 split partials and reduce them in a fixed order +with a single final low-precision rounding, which is deterministic by +construction. When the split count is 1 every path stores directly +(single writer per output element). + +These tests assert, through the public op path: + +A. Repeatability — ``torch.equal`` across repeated identical calls. +B. Numerical correctness — against an FP32 dequantized reference. + +(B) exists because repeatability alone cannot see the failure mode the +deterministic epilogue introduces: the FP32 partials scratch is +``at::empty`` (its coverage invariant is documented at +``alloc_wmma_partials`` in ``csrc/rocm/q_gemm_rdna3_wmma.cu``), and across +repeated calls the caching allocator hands back the same block — a slot +that ever went unwritten would be bitwise-repeatable *and* wrong. + +Reference semantics (uint4b8, GPTQv1, synthesized zero points): + W[k, n] = (q[k, n] - 8) * float(round_to_dtype(scale[g, n])), g = k//G + ref[m, n] = x[m, :] @ W[:, n] (FP32 matmul) + +Tolerances are derived from the kernel's rounding structure, not tuned to +pass (measured max errors on W7900/gfx1100 in parentheses): + +* scalar bf16 — dequant keeps FP32 precision end-to-end (magic-value + v_dot2 path); the only rounding is the final cast. Measured max ~0.06 + on the tested shapes (|ref| up to ~30); tolerance 0.10 with headroom + for larger |ref| at the final-cast ulp. +* scalar fp16 — the dequant uses the exact factored form (magic values in + fp16, (y, z) correction in fp32; see qdq_4_rdna3.cuh), so the only fp16 + roundings are the final output cast and the B-tile-free dot inputs: + max_abs <= 0.5 * ulp_fp16(max|ref|) * 1.5 + 1e-3 (measured ~0.015 at + |ref|~37; the pre-fix baked-constant form measured ~1.0-1.4 on the same + shapes — 94-200x worse). +* WMMA bf16 — B is narrowed to bf16 per cell (one rounding, + <= 0.5*ulp_bf16(|B|)), accumulated in FP32; measured max ≈ 0.09 at + K=4096 — tolerance 0.25. +* WMMA fp16 — B narrowed to fp16 via the precise sub-then-mul variant + (one rounding, ~1e-4); measured max ≈ 0.015 — tolerance 0.05. + +All four are orders of magnitude below what an unwritten scratch slot +(caching-allocator garbage / NaN) or an indexing bug (errors of the same +scale as |ref| ≈ 7) would produce. + +Test hygiene: the op path needs a single-process distributed group and a +current VLLM config; tests use the repo-standard ``dist_init`` fixture from +``tests/conftest.py`` (real context manager, temp-file rendezvous so no +hardcoded MASTER_PORT can collide in parallel CI, and the suite's standard +post-test cleanup). +""" + +import pytest +import torch + +from vllm.platforms import current_platform + +if not current_platform.is_rocm(): + pytest.skip("RDNA3 W4A16 kernel is ROCm-only", allow_module_level=True) + +from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import ( # noqa: E402 + MPLinearLayerConfig, +) +from vllm.model_executor.kernels.linear.mixed_precision.rdna3_w4a16 import ( # noqa: E402 + RDNA3W4A16LinearKernel, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( # noqa: E402 + pack_quantized_values_into_int32, +) +from vllm.model_executor.parameter import ( # noqa: E402 + GroupQuantScaleParameter, + PackedvLLMParameter, +) +from vllm.platforms.rocm import on_gfx1100 # noqa: E402 +from vllm.scalar_type import scalar_types # noqa: E402 + +device = "cuda" +WEIGHT_TYPE = scalar_types.uint4b8 + +gfx1100_only = pytest.mark.skipif( + not ( + on_gfx1100() + and hasattr(torch.ops, "_rocm_C") + and hasattr(torch.ops._rocm_C, "gptq_gemm_rdna3") + ), + reason="requires gfx1100 with the _rocm_C.gptq_gemm_rdna3 op built in", +) + +REPEATS = 20 +GROUP = 128 +DTYPES = [torch.bfloat16, torch.float16] + +# max_abs_error tolerances vs the FP32 reference, per (path, dtype). +# Derivation in the module docstring; ~2-3x headroom over measured maxima. +TOL = { + ("scalar", torch.bfloat16): 0.10, + ("scalar", torch.float16): 0.025, + ("wmma", torch.bfloat16): 0.25, + ("wmma", torch.float16): 0.05, +} + + +def _build_layer(k, n, seed, dtype, group=GROUP): + torch.manual_seed(seed) + q_int4_kn = torch.randint(0, 16, (k, n), dtype=torch.int32) + scales_gn = (torch.randn(k // group, n) * 0.01 + 0.02).to(dtype) + qweight = pack_quantized_values_into_int32(q_int4_kn, WEIGHT_TYPE, packed_dim=0) + no_loader = lambda *a, **kw: None # noqa: E731 + + class DummyLayer(torch.nn.Module): + pass + + layer = DummyLayer() + layer.register_parameter( + "qweight", + PackedvLLMParameter( + data=qweight, + weight_loader=no_loader, + input_dim=0, + output_dim=1, + packed_dim=0, + packed_factor=8, + ), + ) + layer.register_parameter( + "scales", + GroupQuantScaleParameter( + data=scales_gn, weight_loader=no_loader, input_dim=0, output_dim=1 + ), + ) + layer.to(device) + return layer, q_int4_kn, scales_gn + + +def _prepare(layer, dtype, k, n, group=GROUP): + cfg = MPLinearLayerConfig( + full_weight_shape=(k, n), + partition_weight_shape=(k, n), + weight_type=WEIGHT_TYPE, + act_type=dtype, + group_size=group, + zero_points=False, + ) + kernel = RDNA3W4A16LinearKernel( + cfg, w_q_param_name="qweight", w_s_param_name="scales", w_zp_param_name=None + ) + kernel.process_weights_after_loading(layer) + return kernel._get_weight_params(layer) + + +def _reference(m, k, dtype, seed, q_int4_kn, scales_gn, group=GROUP): + torch.manual_seed(seed + 1) + x = torch.randn(m, k, device=device, dtype=dtype) + # FP32 reference; see module docstring for the dequant semantics. + w_f32 = (q_int4_kn.to(device).float() - 8.0) * scales_gn.to( + device + ).repeat_interleave(group, dim=0).float() + return x, x.float() @ w_f32 + + +def _run_op(x, w_q, w_zp, w_s): + return torch.ops._rocm_C.gptq_gemm_rdna3(x, w_q, w_zp, w_s, False) + + +def _outputs_and_ref(M, K, N, seed, dtype, repeats=1, group=GROUP): + layer, q_int4_kn, scales_gn = _build_layer(K, N, seed, dtype, group) + w_q, w_s, w_zp = _prepare(layer, dtype, K, N, group) + x, ref = _reference(M, K, dtype, seed, q_int4_kn, scales_gn, group) + outs = [_run_op(x, w_q, w_zp, w_s) for _ in range(repeats)] + return outs, ref + + +def _path_for(dtype, m): + """Public dispatch: both dtypes reach WMMA at M>=12 (see the + WMMA_MIN_M crossover table in csrc/rocm/q_gemm_rdna3.cu).""" + if m >= 12: + return "wmma" + return "scalar" + + +def _assert_repeatable(outs): + for o in outs[1:]: + assert torch.equal(o, outs[0]), ( + "RDNA3 W4A16 split-K produced differing outputs for identical inputs" + ) + + +def _assert_close_to_ref(out, ref, path, dtype): + err = (out.float() - ref).abs() + max_abs = err.max().item() + mean_abs = err.mean().item() + tol = TOL[(path, dtype)] + assert torch.isfinite(out.float()).all(), "non-finite output" + assert max_abs <= tol, ( + f"RDNA3 W4A16 {path} output deviates from the FP32 dequantized " + f"reference: max_abs={max_abs:.4f} > tol={tol} (mean_abs=" + f"{mean_abs:.4f}). This catches values that are bitwise-repeatable " + "but wrong (e.g. an unwritten FP32 scratch slot)." + ) + + +# --------------------------------------------------------------------------- +# A. Repeatability (bit-exact across identical calls) +# --------------------------------------------------------------------------- + + +@gfx1100_only +@pytest.mark.parametrize("dtype", DTYPES) +def test_scalar_splitk_bit_repeatable(dist_init, dtype): + """Scalar path at production-like contention: K=4096 splits into + ceil(4096/256) = 16 concurrent writers per output element (the old CAS + bug's onset was between 2 and 4 writers, so this sits well inside the + failure regime rather than at the empirical threshold).""" + outs, _ = _outputs_and_ref( + M=1, K=4096, N=4096, seed=1234, dtype=dtype, repeats=REPEATS + ) + _assert_repeatable(outs) + + +@gfx1100_only +@pytest.mark.parametrize("dtype", DTYPES) +def test_scalar_small_m_bit_repeatable(dist_init, dtype): + """Scalar M>1 tiles share the same deterministic epilogue.""" + outs, _ = _outputs_and_ref( + M=8, K=4096, N=4096, seed=1235, dtype=dtype, repeats=REPEATS + ) + _assert_repeatable(outs) + + +@gfx1100_only +@pytest.mark.parametrize("dtype", DTYPES) +def test_scalar_single_split_bit_repeatable(dist_init, dtype): + """z_count == 1 (K=256 < BLOCK_KN_SIZE): direct-store epilogue, no + scratch, no reduce pass.""" + outs, _ = _outputs_and_ref( + M=1, K=256, N=4096, seed=1238, dtype=dtype, repeats=REPEATS + ) + _assert_repeatable(outs) + + +@gfx1100_only +@pytest.mark.parametrize("dtype", DTYPES) +def test_wmma_splitk_bit_repeatable(dist_init, dtype): + """WMMA path (both dtypes M >= 12), split-K active. + + K=6656 gives K_SPLIT=4 under the upstream heuristic; the deterministic + epilogue must be bit-repeatable regardless of the split count. + """ + m = 12 # smallest M on the WMMA side of the dispatch boundary + outs, _ = _outputs_and_ref( + M=m, K=6656, N=4096, seed=1236, dtype=dtype, repeats=REPEATS + ) + _assert_repeatable(outs) + + +@gfx1100_only +@pytest.mark.parametrize("dtype", DTYPES) +def test_wmma_large_m_bit_repeatable(dist_init, dtype): + """Large-M WMMA tiles (128x64 kernels) with the deterministic epilogue.""" + outs, _ = _outputs_and_ref( + M=128, K=6656, N=4096, seed=1237, dtype=dtype, repeats=REPEATS + ) + _assert_repeatable(outs) + + +# --------------------------------------------------------------------------- +# B. Numerical correctness against the FP32 dequantized reference +# --------------------------------------------------------------------------- + + +@gfx1100_only +@pytest.mark.parametrize("dtype", DTYPES) +def test_scalar_matches_fp32_reference(dist_init, dtype): + """Scalar split-K (K=4096, 16 writers) vs the FP32 reference.""" + m = 1 + outs, ref = _outputs_and_ref(M=m, K=4096, N=512, seed=1240, dtype=dtype) + _assert_close_to_ref(outs[0], ref, _path_for(dtype, m), dtype) + + +@gfx1100_only +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("m", [2, 3, 4, 5, 8]) +def test_scalar_multi_m_matches_fp32_reference(dist_init, m, dtype): + """The exact factored dequant across every M_COUNT instantiation + (M=2,3 -> M_COUNT=2; M=4,5 -> M_COUNT=4 with a padded row; M=8 -> 8): + the (y, z) correction and the nibble-slot activation sums must hold + for every tile width, not just the M=1 fast path.""" + outs, ref = _outputs_and_ref(M=m, K=4096, N=512, seed=1247, dtype=dtype) + _assert_repeatable(outs) + _assert_close_to_ref(outs[0], ref, _path_for(dtype, m), dtype) + + +@gfx1100_only +@pytest.mark.parametrize("dtype", DTYPES) +def test_scalar_single_split_matches_fp32_reference(dist_init, dtype): + """Scalar z_count == 1 direct-store path vs the FP32 reference.""" + m = 1 + outs, ref = _outputs_and_ref(M=m, K=256, N=512, seed=1241, dtype=dtype) + _assert_close_to_ref(outs[0], ref, _path_for(dtype, m), dtype) + + +@gfx1100_only +@pytest.mark.parametrize("dtype", DTYPES) +def test_wmma_matches_fp32_reference(dist_init, dtype): + """WMMA split-K path (16x16_1w at the M=12 boundary) vs FP32 reference.""" + m = 12 # smallest M on the WMMA side of the dispatch boundary + outs, ref = _outputs_and_ref(M=m, K=4096, N=512, seed=1242, dtype=dtype) + _assert_close_to_ref(outs[0], ref, _path_for(dtype, m), dtype) + + +@gfx1100_only +@pytest.mark.parametrize("dtype", DTYPES) +def test_wmma_single_split_matches_fp32_reference(dist_init, dtype): + """WMMA k_split == 1 direct-store path (K=256 -> compute_wmma_k_split + returns 1 for every dispatch level) vs the FP32 reference.""" + m = 12 # smallest M on the WMMA side of the dispatch boundary + outs, ref = _outputs_and_ref(M=m, K=256, N=512, seed=1243, dtype=dtype) + _assert_close_to_ref(outs[0], ref, _path_for(dtype, m), dtype) + + +@gfx1100_only +@pytest.mark.parametrize("dtype", DTYPES) +def test_wmma_large_m_matches_fp32_reference(dist_init, dtype): + """Large-M WMMA (128x64_k32, row-tiled scratch + reduce) vs FP32.""" + m = 128 + outs, ref = _outputs_and_ref(M=m, K=4096, N=512, seed=1244, dtype=dtype) + _assert_close_to_ref(outs[0], ref, _path_for(dtype, m), dtype) + + +# --------------------------------------------------------------------------- +# C. V7/V8 (128x64) no-split path: k_split == 1 -> direct store, no scratch, +# no reduce pass (partials == nullptr). +# --------------------------------------------------------------------------- + +# The measured production-like shape from the W7900 benchmark campaign. +# Routing arithmetic (see _assert_v7v8_no_split_routing): with (128, 64) +# tiles, blocks_xy = ceil(25600/64) * ceil(512/128) = 400 * 4 = 1600 >= +# 1500, so compute_wmma_k_split_mn returns 1. Smaller no-split shapes are +# not practical: N would have to grow past ~24k at M=512 (or M past ~1k at +# N=12k) to reach the 1500-block threshold, so this is effectively the +# smallest production-realistic V7/V8 no-split shape. +V7V8_NOSPLIT = (512, 25600, 6656) # (M, N, K) + + +def _compute_wmma_k_split_mn(m, n, k, m_tile, n_tile): + """Faithful replica of compute_wmma_k_split_mn (C++ heuristic).""" + blocks_xy = ((n + n_tile - 1) // n_tile) * ((m + m_tile - 1) // m_tile) + if blocks_xy >= 1500: + return 1 + if blocks_xy * 2 >= 1500 and k >= 512 and k % 32 == 0: + return 2 + if blocks_xy * 4 >= 1500 and k >= 1024 and k % 64 == 0: + return 4 + if k >= 1024 and k % 64 == 0: + return 4 + if k >= 512 and k % 32 == 0: + return 2 + return 1 + + +def _assert_v7v8_no_split_routing(dtype, m, n, k): + """Assert the shape deterministically routes to the V7/V8 128x64 + kernel with k_split == 1: WMMA dispatch (both dtypes M >= 12), the + M >= 128 branch, and the no-split threshold.""" + assert m >= 12, "not the WMMA path" + assert m >= 128, "below the V7/V8 (128x64) branch" + assert _compute_wmma_k_split_mn(m, n, k, 128, 64) == 1, ( + "shape does not route to the k_split == 1 no-split path" + ) + + +@gfx1100_only +@pytest.mark.parametrize("dtype", DTYPES) +def test_wmma_v7v8_no_split_bit_repeatable(dist_init, dtype): + """V7/V8 128x64_k32 with k_split == 1: single direct writer per output + cell — no FP32 scratch, no reduce pass.""" + m, n, k = V7V8_NOSPLIT + _assert_v7v8_no_split_routing(dtype, m, n, k) + outs, _ = _outputs_and_ref(M=m, K=k, N=n, seed=1245, dtype=dtype, repeats=REPEATS) + _assert_repeatable(outs) + + +@gfx1100_only +@pytest.mark.parametrize("dtype", DTYPES) +def test_wmma_v7v8_no_split_matches_fp32_reference(dist_init, dtype): + """V7/V8 128x64_k32 no-split direct-store path vs the FP32 reference.""" + m, n, k = V7V8_NOSPLIT + _assert_v7v8_no_split_routing(dtype, m, n, k) + outs, ref = _outputs_and_ref(M=m, K=k, N=n, seed=1246, dtype=dtype) + _assert_close_to_ref(outs[0], ref, _path_for(dtype, m), dtype) + + +# --------------------------------------------------------------------------- +# D. Dispatch boundary: M=11 must stay scalar, M=12 must route to WMMA +# (WMMA_MIN_M in csrc/rocm/q_gemm_rdna3.cu). Observed via the CUDA +# profiler: the scalar and WMMA kernels have disjoint names. +# --------------------------------------------------------------------------- + + +def _launched_kernel_names(fn) -> set: + from torch.profiler import ProfilerActivity, profile + + fn() + torch.accelerator.synchronize() + with profile(activities=[ProfilerActivity.CUDA]) as prof: + fn() + torch.accelerator.synchronize() + return { + e.name + for e in prof.events() + if "gemm_q4" in e.name or "reduce_partials" in e.name + } + + +@pytest.mark.parametrize("dtype", DTYPES) +@gfx1100_only +def test_dispatch_boundary(dist_init, dtype): + """The scalar/WMMA crossover sits between M=11 and M=12 for both + dtypes; the boundary must be exact (see the WMMA_MIN_M table).""" + + def run(m): + outs, _ = _outputs_and_ref(M=m, K=4096, N=512, seed=1248, dtype=dtype) + return outs[0] + + m11 = run(11) + names11 = _launched_kernel_names(lambda: run(11)) + assert any("gemm_q4_kernel_rdna3" in n for n in names11), names11 + assert not any("wmma" in n for n in names11), names11 + + m12 = run(12) + names12 = _launched_kernel_names(lambda: run(12)) + assert any("wmma" in n for n in names12), names12 + assert not any("gemm_q4_kernel_rdna3" in n for n in names12), names12 + + # Both sides of the boundary stay correct against the reference. + _, ref11 = _outputs_and_ref(M=11, K=4096, N=512, seed=1248, dtype=dtype) + _, ref12 = _outputs_and_ref(M=12, K=4096, N=512, seed=1248, dtype=dtype) + _assert_close_to_ref(m11, ref11, "scalar", dtype) + _assert_close_to_ref(m12, ref12, "wmma", dtype) + + +# --------------------------------------------------------------------------- +# E. In-kernel split-K reduction: stress the per-call ticket epilogue. +# The fused reduction must stay deterministic and correct under +# back-to-back unsynchronized calls, concurrent invocations on two +# streams (per-call tickets must never share state), and CUDA-graph +# capture/replay. +# --------------------------------------------------------------------------- + + +@gfx1100_only +@pytest.mark.parametrize("dtype", DTYPES) +def test_scalar_splitk_100x_bit_repeatable(dist_init, dtype): + """100 back-to-back calls (no intermediate sync) on the split-K path.""" + outs, _ = _outputs_and_ref(M=1, K=4096, N=512, seed=1249, dtype=dtype, repeats=100) + _assert_repeatable(outs) + + +@gfx1100_only +@pytest.mark.parametrize("dtype", DTYPES) +def test_scalar_multistream_concurrent(dist_init, dtype): + """Two concurrent invocations on separate streams, distinct weights and + activations: outputs must be bit-identical to the serial invocation of + each (the arrival tickets are per-call, so cross-stream interference + would show up here as a corrupted or unreduced output).""" + layer_a, q_a, s_a = _build_layer(4096, 512, 4321, dtype) + layer_b, q_b, s_b = _build_layer(4096, 512, 4322, dtype) + wq_a, ws_a, wz_a = _prepare(layer_a, dtype, 4096, 512) + wq_b, ws_b, wz_b = _prepare(layer_b, dtype, 4096, 512) + torch.manual_seed(4323) + x_a = torch.randn(1, 4096, device=device, dtype=dtype) + x_b = torch.randn(1, 4096, device=device, dtype=dtype) + + serial_a = _run_op(x_a, wq_a, wz_a, ws_a) + serial_b = _run_op(x_b, wq_b, wz_b, ws_b) + torch.accelerator.synchronize() + + s1, s2 = torch.cuda.Stream(), torch.cuda.Stream() + for _ in range(25): + with torch.cuda.stream(s1): + out_a = _run_op(x_a, wq_a, wz_a, ws_a) + with torch.cuda.stream(s2): + out_b = _run_op(x_b, wq_b, wz_b, ws_b) + torch.accelerator.synchronize() + assert torch.equal(out_a, serial_a), "concurrent stream A corrupted" + assert torch.equal(out_b, serial_b), "concurrent stream B corrupted" + + +@gfx1100_only +@pytest.mark.parametrize("dtype", DTYPES) +def test_scalar_graph_capture_replay(dist_init, dtype): + """The op must survive CUDA-graph capture: the replayed output must be + bit-identical to the eager output and stable across repeated replays.""" + layer, q_int4, scales_gn = _build_layer(4096, 512, 4324, dtype) + wq, ws, wz = _prepare(layer, dtype, 4096, 512) + torch.manual_seed(4325) + x = torch.randn(1, 4096, device=device, dtype=dtype) + + eager = _run_op(x, wq, wz, ws) + torch.accelerator.synchronize() + + graph = torch.cuda.CUDAGraph() + _run_op(x, wq, wz, ws) # warm-up side streams/allocator on capture path + torch.accelerator.synchronize() + with torch.cuda.graph(graph): + captured = _run_op(x, wq, wz, ws) + graph.replay() + torch.accelerator.synchronize() + assert torch.equal(captured, eager), "graph replay differs from eager" + graph.replay() + graph.replay() + torch.accelerator.synchronize() + assert torch.equal(captured, eager), "repeated replay not repeatable" + + +@gfx1100_only +@pytest.mark.parametrize("dtype", DTYPES) +def test_scalar_multi_row_tile_bit_repeatable(dist_init, dtype): + """Row tiling with ticket reuse: N % 16 == 8 bypasses the WMMA + dispatch (N % 16 != 0), so M=100 runs scalar across two 64-row tiles + (the second with gy > 1 y-blocks). Every tile's tickets must be armed, + consumed, and re-armed correctly for all outputs to be written exactly + once — and the result must be bit-repeatable.""" + outs, ref = _outputs_and_ref( + M=100, K=4096, N=2056, seed=1250, dtype=dtype, repeats=20 + ) + _assert_repeatable(outs) + _assert_close_to_ref(outs[0], ref, "scalar", dtype) + + +# --------------------------------------------------------------------------- +# F. Split-K geometry corners: clipped non-first z blocks, small split +# counts, and non-128 group sizes. +# --------------------------------------------------------------------------- + + +@gfx1100_only +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("k,group", [(4224, 128), (2112, 32)]) +def test_scalar_clipped_tail_splitk(dist_init, dtype, k, group): + """K % 256 != 0 with z_count > 1: the LAST z block is clipped by end_k, + its group prefetch guard (`k + 32 < end_k`) is what keeps the raw + zero/scale loads in bounds, and the reduce tail loop runs (z=17 -> one + remainder slice; z=9 -> one batch + one). K=2112/G=32 also puts two + group switches inside the clipped tail block.""" + outs, ref = _outputs_and_ref(M=1, K=k, N=512, seed=1251, dtype=dtype, group=group) + _assert_repeatable(outs) + _assert_close_to_ref(outs[0], ref, "scalar", dtype) + + +@gfx1100_only +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("group", [32, 64]) +def test_scalar_small_group_matches_fp32_reference(dist_init, dtype, group): + """Group sizes below 128 make the group constants change every 32-K + round (G=32) or three times per z block (G=64); the factored-dequant + correction must track every change. G=32 is also the production MoE + group size.""" + outs, ref = _outputs_and_ref( + M=1, K=2048, N=512, seed=1252, dtype=dtype, group=group + ) + _assert_repeatable(outs) + _assert_close_to_ref(outs[0], ref, "scalar", dtype) + + +@gfx1100_only +@pytest.mark.parametrize("dtype", DTYPES) +def test_scalar_two_way_split_bit_repeatable(dist_init, dtype): + """z_count == 2 (K=512): the minimal split-K configuration — the old + CAS bug's onset regime — and the pure-remainder reduce tail.""" + outs, _ = _outputs_and_ref( + M=1, K=512, N=512, seed=1253, dtype=dtype, repeats=REPEATS + ) + _assert_repeatable(outs) From 84bcbc62644356270aaaa5e2d0237d03adc9bb3a Mon Sep 17 00:00:00 2001 From: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Date: Sat, 3 Oct 2026 13:26:07 -0400 Subject: [PATCH 02/22] [DCP] Enable TokenSpeed MLA with block-interleaved DCP (#59462) Signed-off-by: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- requirements/cuda.txt | 2 +- tests/v1/attention/test_mla_backends.py | 148 +++++++++++++++++- .../attention/backends/mla/tokenspeed_mla.py | 9 ++ 3 files changed, 156 insertions(+), 3 deletions(-) diff --git a/requirements/cuda.txt b/requirements/cuda.txt index fbc80b28e1..e0db38baad 100644 --- a/requirements/cuda.txt +++ b/requirements/cuda.txt @@ -30,7 +30,7 @@ nvidia-cutlass-dsl[cu13]==4.8.0 quack-kernels==0.6.5 # Required by MSA # Tokenspeed_MLA for faster mla with spec decode -tokenspeed-mla==0.1.8; platform_system == "Linux" +tokenspeed-mla==0.2.15; platform_system == "Linux" # Humming kernels for quantization gemm humming-kernels[cu13]==0.1.16 diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index 9ae0fc216a..1c192f3f1c 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -1098,8 +1098,14 @@ def test_flashinfer_mla_dspark_dcp_supports_target_and_draft(monkeypatch): pytest.param(False, 3, 1, 0, id="noncausal-multi-token"), ], ) +@pytest.mark.parametrize("cp_interleave_size", [1, 16, 896]) def test_tokenspeed_mla_decode_contract( - monkeypatch, causal, tokens_per_decode, dcp_world_size, dcp_rank + monkeypatch, + causal, + tokens_per_decode, + dcp_world_size, + dcp_rank, + cp_interleave_size, ): decode_call = None num_decodes = 2 @@ -1135,7 +1141,7 @@ def test_tokenspeed_mla_decode_contract( impl = object.__new__(tokenspeed_mla_module.TokenspeedMLAImpl) impl.dcp_world_size = dcp_world_size impl.dcp_rank = dcp_rank - impl.cp_kv_cache_interleave_size = 1 + impl._parallel_config = SimpleNamespace(cp_kv_cache_interleave_size=1) impl.need_to_return_lse_for_decode = True impl.kv_lora_rank = kv_lora_rank impl.qk_rope_head_dim = qk_rope_head_dim @@ -1164,6 +1170,8 @@ def test_tokenspeed_mla_decode_contract( dtype=torch.float8_e4m3fn, ) + # NIXL may resolve the interleave size after the implementation is created. + impl._parallel_config.cp_kv_cache_interleave_size = cp_interleave_size out, lse = impl.forward_mqa( q, kv_cache, @@ -1194,6 +1202,142 @@ def test_tokenspeed_mla_decode_contract( assert decode_call["return_lse"] is True assert decode_call["cp_world"] == dcp_world_size assert decode_call["cp_rank"] == dcp_rank + assert decode_call.get("cp_interleave_size", 1) == cp_interleave_size + + +@pytest.mark.skipif( + AttentionBackendEnum.TOKENSPEED_MLA not in BACKENDS_TO_TEST, + reason="TokenSpeed MLA requires an SM100-family GPU and tokenspeed-mla", +) +@pytest.mark.parametrize( + ("cp_interleave_size", "seq_len"), + [ + pytest.param(1, 5, id="token-interleave-empty-local-ranks"), + pytest.param(1, 8 + 17, id="token-interleave-full-cycle"), + pytest.param(896, 5, id="block-interleave-empty-local-ranks"), + pytest.param(896, 8 * 896 + 17, id="block-interleave-full-cycle"), + ], +) +def test_tokenspeed_mla_dcp_matches_unsharded_decode(cp_interleave_size, seq_len): + from tokenspeed_mla import tokenspeed_mla_decode + + from vllm.v1.attention.ops.dcp import ( + _lse_weighted_combine, + mask_dcp_empty_shards_, + ) + + torch.manual_seed(7) + device = torch.device("cuda") + dcp_world_size = 8 + kernel_block_size = 64 + num_heads = 128 + kv_lora_rank = 512 + qk_rope_head_dim = 64 + head_size = kv_lora_rank + qk_rope_head_dim + + def make_paged_cache(tokens: torch.Tensor): + # Empty DCP ranks still have the null block in vLLM's physical cache. + num_pages = max(1, cdiv(tokens.shape[0], kernel_block_size)) + cache = torch.zeros( + num_pages, + kernel_block_size, + head_size, + dtype=tokens.dtype, + device=device, + ) + cache.view(-1, head_size)[: tokens.shape[0]].copy_(tokens) + block_table = torch.arange( + num_pages, dtype=torch.int32, device=device + ).unsqueeze(0) + return cache, block_table + + # Four query positions exercise the MTP/DSpark causal-mask path. + query = (torch.randn(1, 4, num_heads, head_size, device=device) * 0.1).to( + torch.float8_e4m3fn + ) + global_tokens = (torch.randn(seq_len, head_size, device=device) * 0.1).to( + torch.float8_e4m3fn + ) + workspace = tokenspeed_mla_module._get_workspace(device, num_heads, kv_lora_rank) + global_cache, global_block_table = make_paged_cache(global_tokens) + global_seq_len = torch.tensor([seq_len], dtype=torch.int32, device=device) + reference, _ = tokenspeed_mla_decode( + query=query, + kv_cache=global_cache, + workspace_buffer=workspace, + kv_lora_rank=kv_lora_rank, + qk_rope_head_dim=qk_rope_head_dim, + block_tables=global_block_table, + seq_lens=global_seq_len, + max_seq_len=seq_len, + softmax_scale=head_size**-0.5, + return_lse=True, + ) + + positions = torch.arange(seq_len, device=device) + owners = positions.div(cp_interleave_size, rounding_mode="floor").remainder( + dcp_world_size + ) + max_local_seq_len = ( + cdiv(seq_len, dcp_world_size * cp_interleave_size) * cp_interleave_size + ) + partial_outputs = [] + partial_lses = [] + for dcp_rank in range(dcp_world_size): + local_tokens = global_tokens[owners == dcp_rank] + local_cache, local_block_table = make_paged_cache(local_tokens) + local_seq_len = torch.tensor( + [local_tokens.shape[0]], dtype=torch.int32, device=device + ) + impl = object.__new__(tokenspeed_mla_module.TokenspeedMLAImpl) + impl.dcp_world_size = dcp_world_size + impl.dcp_rank = dcp_rank + impl._parallel_config = SimpleNamespace( + cp_kv_cache_interleave_size=cp_interleave_size + ) + impl.need_to_return_lse_for_decode = True + impl.kv_lora_rank = kv_lora_rank + impl.qk_rope_head_dim = qk_rope_head_dim + impl.num_heads = num_heads + impl.scale = head_size**-0.5 + impl.softmax_scale = None + impl.output_scale = None + impl._workspace_buffer = workspace + metadata = SimpleNamespace( + num_decodes=1, + num_decode_tokens=query.shape[1], + max_seq_len=max_local_seq_len, + causal=True, + decode=SimpleNamespace( + block_table=local_block_table, + seq_lens=local_seq_len, + dcp_tot_seq_lens=global_seq_len, + ), + ) + output, lse = impl.forward_mqa( + query.view(-1, num_heads, head_size), + local_cache, + metadata, + SimpleNamespace(_q_scale_float=1.0, _k_scale_float=1.0), + ) + # Empty-shard normalization belongs to the downstream DCP combine path. + mask_dcp_empty_shards_( + lse, + local_seq_len, + torch.tensor([0, query.shape[1]], dtype=torch.int32, device=device), + ) + partial_outputs.append( + output.view(query.shape[1], num_heads, kv_lora_rank).float() + ) + partial_lses.append(lse.view(query.shape[1], num_heads).float()) + + # TokenSpeed returns base-2 LSE. Merge the per-rank partial attention states + # the same way MLADCPManager does before comparing to the unsharded result. + merged = _lse_weighted_combine( + torch.stack(partial_outputs), torch.stack(partial_lses), is_lse_base_on_e=False + ).unsqueeze(0) + # FP8 attention rounds differently when KV is partitioned across DCP ranks. + torch.testing.assert_close(merged, reference.float(), atol=1e-2, rtol=1e-2) @pytest.mark.parametrize("is_fp8_kvcache", [False, True], ids=["bf16", "fp8"]) diff --git a/vllm/v1/attention/backends/mla/tokenspeed_mla.py b/vllm/v1/attention/backends/mla/tokenspeed_mla.py index b3d982eac3..96859d7f98 100644 --- a/vllm/v1/attention/backends/mla/tokenspeed_mla.py +++ b/vllm/v1/attention/backends/mla/tokenspeed_mla.py @@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, ClassVar import torch +from vllm.config import get_current_vllm_config from vllm.config.cache import CacheDType from vllm.logger import init_logger from vllm.model_executor.layers.attention.mla_attention import ( @@ -155,6 +156,7 @@ class TokenspeedMLABackend(MLACommonBackend): class TokenspeedMLAImpl(MLACommonImpl[MLACommonMetadata]): can_return_lse_for_decode: bool = True supports_dcp: bool = True + supports_mtp_with_cp_non_trivial_interleave_size: bool = True # tokenspeed_mla_decode returns LSE in log2 units; its own DCP test merges # partial outputs with exp2(lse). lse_base_on_e: bool = False @@ -216,6 +218,9 @@ class TokenspeedMLAImpl(MLACommonImpl[MLACommonMetadata]): self._workspace_buffer: torch.Tensor | None = None self.softmax_scale: float | None = None self.output_scale: float | None = None + # NIXL resolves interleaving after model construction; retain the config + # rather than caching its initial interleave size. + self._parallel_config = get_current_vllm_config().parallel_config # Pre-JIT BF16 and FP8 prefill kernels here too — decode impl always # runs when tokenspeed is selected, prefill backend may not (user can @@ -311,9 +316,13 @@ class TokenspeedMLAImpl(MLACommonImpl[MLACommonMetadata]): causal_seqs=causal_seqs if self.dcp_world_size > 1 else None, cp_world=self.dcp_world_size, cp_rank=self.dcp_rank, + cp_interleave_size=self._parallel_config.cp_kv_cache_interleave_size, ) if return_lse: o, lse = kernel_out + # DCP ranks with no local KV may return undefined outputs (e.g., + # NaNs). The downstream DCP combine path must mask empty + # shards and ignore their outputs when merging partial attention. lse = lse.view(-1, lse.shape[-1]) else: o, lse = kernel_out, None From 9e0b6fcf77b57401a08224df89f893459ceee375 Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Sun, 4 Oct 2026 01:42:35 +0100 Subject: [PATCH 03/22] [CI] Bump Transformers version to 5.18.0 (#59621) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- requirements/common.txt | 2 +- requirements/test/cpu.txt | 2 +- requirements/test/cuda.in | 2 +- requirements/test/cuda.txt | 2 +- requirements/test/nightly-torch.txt | 2 +- requirements/test/rocm.in | 2 +- requirements/test/rocm.txt | 2 +- requirements/test/xpu.in | 2 +- requirements/test/xpu.txt | 2 +- tests/models/registry.py | 6 +- vllm/config/model.py | 2 + vllm/config/speculative.py | 2 +- vllm/model_executor/models/lightonocr.py | 13 ++ vllm/model_executor/models/longcat_flash.py | 166 ++---------------- .../models/longcat_flash_mtp.py | 13 +- .../models/longcat_flash_ngram.py | 22 +-- vllm/models/hy_v4/nvidia/attention.py | 8 +- vllm/models/qwen4_exp/amd/model.py | 10 +- vllm/models/qwen4_exp/nvidia/model.py | 10 +- vllm/transformers_utils/config.py | 22 +++ .../model_arch_config_convertor.py | 8 +- vllm/v1/worker/gpu/attn_utils.py | 5 +- 22 files changed, 104 insertions(+), 201 deletions(-) diff --git a/requirements/common.txt b/requirements/common.txt index 6e0b0db783..433494056d 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -7,7 +7,7 @@ requests >= 2.33.0 tqdm blake3 py-cpuinfo -transformers >= 5.16.1, < 5.18.0 +transformers >= 5.16.1, < 5.19.0 huggingface_hub >= 1.31.0 # Required for the `huggingface_hub.utils.httpx` re-export tokenizers >= 0.23.2 # Required for fast incremental detokenization. safetensors >= 0.6.2 # MXFP4/MXFP6 dtype support (F8_E8M0, F4) added in 0.6.0: https://github.com/huggingface/safetensors/pull/611 diff --git a/requirements/test/cpu.txt b/requirements/test/cpu.txt index 926a5ceb1c..91fac06966 100644 --- a/requirements/test/cpu.txt +++ b/requirements/test/cpu.txt @@ -1228,7 +1228,7 @@ traitlets==5.16.1 # via # jupyter-client # jupyter-core -transformers==5.17.0 +transformers==5.18.0 # via # -r requirements/test/../common.txt # -r requirements/test/cuda.in diff --git a/requirements/test/cuda.in b/requirements/test/cuda.in index cb5ceb248e..4732f3a505 100644 --- a/requirements/test/cuda.in +++ b/requirements/test/cuda.in @@ -40,7 +40,7 @@ datamodel_code_generator>=0.64.0 # required for minicpm3 test nltk>=3.10.3 # pin transitive rouge-score/lm-eval dep for known GHSAs lm-eval[api]>=0.4.12 # required for model evaluation test mteb[bm25s]>=2, <3 # required for mteb test -transformers==5.17.0 +transformers==5.18.0 schemathesis>=4.0.0 # Required for openai schema test. # quantization buildkite-test-collector==0.1.9 diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt index 8f39bb900e..ad185844be 100644 --- a/requirements/test/cuda.txt +++ b/requirements/test/cuda.txt @@ -1335,7 +1335,7 @@ traitlets==5.16.1 # via # jupyter-client # jupyter-core -transformers==5.17.0 +transformers==5.18.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt diff --git a/requirements/test/nightly-torch.txt b/requirements/test/nightly-torch.txt index 256683967e..a229116859 100644 --- a/requirements/test/nightly-torch.txt +++ b/requirements/test/nightly-torch.txt @@ -29,7 +29,7 @@ opencv-python-headless >= 4.13.0 # required for video test datamodel_code_generator>=0.64.0 # required for minicpm3 test lm-eval[api]>=0.4.12 # required for model evaluation test mteb[bm25s]>=2, <3 # required for mteb test -transformers==5.17.0 +transformers==5.18.0 schemathesis>=4.0.0 # Required for openai schema test. # quantization buildkite-test-collector==0.1.9 diff --git a/requirements/test/rocm.in b/requirements/test/rocm.in index e66343d47b..753fad0810 100644 --- a/requirements/test/rocm.in +++ b/requirements/test/rocm.in @@ -36,7 +36,7 @@ datamodel_code_generator>=0.64.0 # required for minicpm3 test nltk>=3.10.3 # pin transitive rouge-score/lm-eval dep for known GHSAs lm-eval[api]>=0.4.12 # required for model evaluation test mteb[bm25s]>=2, <3 # required for mteb test -transformers==5.17.0 +transformers==5.18.0 schemathesis>=4.0.0 # Required for openai schema test # quantization bitsandbytes==0.50.0 diff --git a/requirements/test/rocm.txt b/requirements/test/rocm.txt index 0df6f9c82a..e6c137ebf5 100644 --- a/requirements/test/rocm.txt +++ b/requirements/test/rocm.txt @@ -1282,7 +1282,7 @@ traitlets==5.16.1 # via # jupyter-client # jupyter-core -transformers==5.17.0 +transformers==5.18.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt diff --git a/requirements/test/xpu.in b/requirements/test/xpu.in index c39e8350dd..4b3d68ddfc 100644 --- a/requirements/test/xpu.in +++ b/requirements/test/xpu.in @@ -20,7 +20,7 @@ nltk>=3.10.3 # pin transitive rouge-score/lm-eval dep for known GHSAs modelscope<1.38 runai-model-streamer[s3,gcs,azure]>=0.15.7 tensorizer==2.10.1 -transformers==5.17.0 +transformers==5.18.0 # --- Audio Processing --- av diff --git a/requirements/test/xpu.txt b/requirements/test/xpu.txt index 33b89a85fc..19cd271017 100644 --- a/requirements/test/xpu.txt +++ b/requirements/test/xpu.txt @@ -1134,7 +1134,7 @@ traitlets==5.16.1 # via # jupyter-client # jupyter-core -transformers==5.17.0 +transformers==5.18.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt diff --git a/tests/models/registry.py b/tests/models/registry.py index 80ca856615..b764cad595 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -380,12 +380,9 @@ _TEXT_GENERATION_EXAMPLE_MODELS = { "Llama4ForCausalLM": _HfExamplesInfo( "meta-llama/Llama-4-Scout-17B-16E-Instruct", ), - "LongcatFlashForCausalLM": _HfExamplesInfo( - "meituan-longcat/LongCat-Flash-Chat", trust_remote_code=True - ), + "LongcatFlashForCausalLM": _HfExamplesInfo("meituan-longcat/LongCat-Flash-Chat"), "LongcatFlashNgramForCausalLM": _HfExamplesInfo( "meituan-longcat/LongCat-Flash-Lite", - trust_remote_code=True, # Shrink the ~62GB n-gram tables (ngram_vocab_size_ratio * vocab_size) # so the dummy-weight init test fits in CI memory. hf_overrides={"ngram_vocab_size_ratio": 1}, @@ -1752,7 +1749,6 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = { ), "LongCatFlashMTPModel": _HfExamplesInfo( "meituan-longcat/LongCat-Flash-Chat", - trust_remote_code=True, speculative_model="meituan-longcat/LongCat-Flash-Chat", ), "MiMoMTPModel": _HfExamplesInfo( diff --git a/vllm/config/model.py b/vllm/config/model.py index 75ca774fa1..3e481d8c91 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -111,6 +111,8 @@ LayerBlockType = Literal["attention", "linear_attention", "mamba"] _ATTENTION_LAYER_TYPES = frozenset( { "full_attention", + "indexed_attention", + # TODO: Delete below once Transformers 5.18.0 is the minimum required version. "deepseek_sparse_attention", "qwen_sparse_attention", } diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index ff4f76153d..47ffb3b7f7 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -960,7 +960,7 @@ class SpeculativeConfig: "architectures": [architecture], } ) - if hf_config.model_type in ("longcat_flash", "longcat_flash_ngram"): + if hf_config.model_type == "longcat_flash": hf_config.model_type = "longcat_flash_mtp" n_predict = getattr(hf_config, "num_nextn_predict_layers", 1) hf_config.update( diff --git a/vllm/model_executor/models/lightonocr.py b/vllm/model_executor/models/lightonocr.py index bc0292ed55..d89fa53299 100644 --- a/vllm/model_executor/models/lightonocr.py +++ b/vllm/model_executor/models/lightonocr.py @@ -50,6 +50,19 @@ class LightOnOCRMultiModalProcessor(BaseMultiModalProcessor[LightOnOCRProcessing def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: return self.dummy_inputs.get_dummy_text(mm_counts) + def _call_hf_processor( + self, + hf_data: Mapping[str, object], + hf_kwargs: Mapping[str, object], + ) -> BatchFeature: + # The checkpoint sets `do_pad: null`, which disables padding since + # Transformers v5.18, so images of different sizes cannot be stacked. + return self.info.ctx.call_hf_processor( + self.info.get_hf_processor(**hf_kwargs), + hf_data, + {**hf_kwargs, "do_pad": True}, + ) + def _postprocess_hf_mm_data( self, mm_data: Mapping[str, object], diff --git a/vllm/model_executor/models/longcat_flash.py b/vllm/model_executor/models/longcat_flash.py index 0c9802a8b3..50962e9728 100644 --- a/vllm/model_executor/models/longcat_flash.py +++ b/vllm/model_executor/models/longcat_flash.py @@ -39,7 +39,7 @@ from itertools import islice import torch from torch import nn -from transformers import PreTrainedConfig +from transformers import LongcatFlashConfig from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, VllmConfig @@ -80,121 +80,6 @@ from .utils import ( logger = init_logger(__name__) -class FlashConfig(PreTrainedConfig): - """Flash model configuration.""" - - moe_intermediate_size: int - - model_type = "longcat_flash" - keys_to_ignore_at_inference = ["past_key_values"] - - def __init__( - self, - vocab_size=131072, - hidden_size=4096, - intermediate_size=8192, - num_layers=28, - num_hidden_layers=None, - num_attention_heads=96, - num_key_value_heads=128, - ep_size=1, - kv_lora_rank=512, - q_lora_rank=1536, - qk_rope_head_dim=64, - v_head_dim=128, - qk_nope_head_dim=128, - num_experts_per_tok=None, - norm_topk_prob=False, - max_position_embeddings=8192, - initializer_range=0.02, - rms_norm_eps=1e-05, - use_cache=True, - pad_token_id=None, - bos_token_id=100000, - eos_token_id=100001, - pretraining_tp=1, - tie_word_embeddings=False, - rope_parameters=None, - attention_bias=False, - attention_dropout=0.0, - mla_scale_q_lora=False, - mla_scale_kv_lora=False, - dtype="bfloat16", - params_dtype="bfloat16", - router_dtype="float32", - router_bias=False, - topk_method=None, - routed_scaling_factor=1.0, - zero_expert_num=0, - zero_expert_type=None, - nextn_use_scmoe=False, - **kwargs, - ): - super().__init__( - pad_token_id=pad_token_id, - bos_token_id=bos_token_id, - eos_token_id=eos_token_id, - tie_word_embeddings=tie_word_embeddings, - dtype=dtype, - params_dtype=params_dtype, - router_dtype=router_dtype, - topk_method=topk_method, - router_bias=router_bias, - nextn_use_scmoe=nextn_use_scmoe, - **kwargs, - ) - self.vocab_size = vocab_size - self.max_position_embeddings = max_position_embeddings - self.hidden_size = hidden_size - self.num_hidden_layers = ( - num_hidden_layers if num_hidden_layers is not None else num_layers - ) - self.num_attention_heads = num_attention_heads - self.ep_size = ep_size - self.kv_lora_rank = kv_lora_rank - self.q_lora_rank = q_lora_rank - self.qk_rope_head_dim = qk_rope_head_dim - self.v_head_dim = v_head_dim - self.qk_nope_head_dim = qk_nope_head_dim - self.num_experts_per_tok = num_experts_per_tok - self.norm_topk_prob = norm_topk_prob - # for backward compatibility - if num_key_value_heads is None: - num_key_value_heads = num_attention_heads - - self.num_key_value_heads = num_key_value_heads - self.initializer_range = initializer_range - self.rms_norm_eps = rms_norm_eps - self.pretraining_tp = pretraining_tp - self.use_cache = use_cache - # Try to set `rope_scaling` if available, otherwise use `rope_parameters` - rope_scaling = kwargs.pop("rope_scaling", None) - rope_parameters = rope_scaling or rope_parameters or {"rope_type": "default"} - rope_theta = kwargs.pop("rope_theta", 1000000.0) - if "rope_theta" not in rope_parameters: - rope_parameters["rope_theta"] = rope_theta - self.rope_parameters = rope_parameters - self.attention_bias = attention_bias - self.attention_dropout = attention_dropout - self.mla_scale_q_lora = mla_scale_q_lora - self.mla_scale_kv_lora = mla_scale_kv_lora - self.zero_expert_num = zero_expert_num - self.zero_expert_type = zero_expert_type - self.routed_scaling_factor = routed_scaling_factor - self.hidden_act = "silu" - self.intermediate_size = ( - self.ffn_hidden_size - if hasattr(self, "ffn_hidden_size") - else intermediate_size - ) - if hasattr(self, "moe_intermediate_size"): - self.moe_intermediate_size = self.moe_intermediate_size - elif hasattr(self, "expert_ffn_hidden_size"): - self.moe_intermediate_size = self.expert_ffn_hidden_size - else: - self.moe_intermediate_size = self.intermediate_size - - class FlashMLP(nn.Module): """Flash MLP layer.""" @@ -242,22 +127,17 @@ class FlashMLP(nn.Module): class LongcatRouter(nn.Module): def __init__( self, - config: FlashConfig, + config: LongcatFlashConfig, zero_expert_num: int, router_params_dtype: torch.dtype, prefix: str = "", ): super().__init__() - self.n_routed_experts = ( - config.n_routed_experts - if hasattr(config, "n_routed_experts") - else config.num_experts[0] - ) - self.n_routed_experts = self.n_routed_experts + zero_expert_num + self.n_routed_experts = config.n_routed_experts + zero_expert_num self.classifier = GateLinear( config.hidden_size, self.n_routed_experts, - bias=config.router_bias, + bias=False, out_dtype=router_params_dtype, params_dtype=router_params_dtype, prefix=f"{prefix}.classifier", @@ -274,7 +154,7 @@ class LongcatRouter(nn.Module): class LongcatMoe(nn.Module): def __init__( self, - config: FlashConfig, + config: LongcatFlashConfig, num_experts: int, top_k: int, hidden_size: int, @@ -286,10 +166,7 @@ class LongcatMoe(nn.Module): ): super().__init__() self.hidden_size = hidden_size - # Gate always runs at half / full precision for now. - self.router_params_dtype = params_dtype - if config.router_dtype == "float32": - self.router_params_dtype = torch.float32 + self.router_params_dtype = torch.float32 self.router = LongcatRouter( config=config, @@ -354,7 +231,7 @@ class FlashDecoderLayer(nn.Module): def __init__( self, vllm_config: VllmConfig, - config: FlashConfig, + config: LongcatFlashConfig, cache_config: CacheConfig | None = None, quant_config: QuantizationConfig | None = None, prefix: str = "", @@ -376,9 +253,7 @@ class FlashDecoderLayer(nn.Module): qk_nope_head_dim=config.qk_nope_head_dim, qk_rope_head_dim=config.qk_rope_head_dim, v_head_dim=config.v_head_dim, - q_lora_rank=( - config.q_lora_rank if hasattr(config, "q_lora_rank") else None - ), + q_lora_rank=config.q_lora_rank, kv_lora_rank=config.kv_lora_rank, max_position_embeddings=max_position_embeddings, cache_config=cache_config, @@ -415,12 +290,8 @@ class FlashDecoderLayer(nn.Module): self.mlp = LongcatMoe( config=config, - num_experts=config.n_routed_experts - if hasattr(config, "n_routed_experts") - else config.num_experts[self.layer_idx], - top_k=config.moe_topk - if hasattr(config, "moe_topk") - else config.num_experts_per_tok, + num_experts=config.n_routed_experts, + top_k=config.moe_topk, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, quant_config=quant_config, @@ -482,7 +353,7 @@ class FlashModel(nn.Module): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() - config = FlashConfig(**vllm_config.model_config.hf_config.__dict__) + config = vllm_config.model_config.hf_config cache_config = vllm_config.cache_config quant_config = vllm_config.quant_config self.config = config @@ -499,7 +370,7 @@ class FlashModel(nn.Module): else: self.embed_tokens = PPMissingLayer() self.start_layer, self.end_layer, self.layers = make_layers( - config.num_hidden_layers, + config.num_layers, lambda prefix: FlashDecoderLayer( vllm_config, config, @@ -561,9 +432,7 @@ class FlashModel(nn.Module): ckpt_gate_proj_name="gate_proj", ckpt_down_proj_name="down_proj", ckpt_up_proj_name="up_proj", - num_experts=self.config.n_routed_experts - if hasattr(self.config, "n_routed_experts") - else self.config.num_experts[0], + num_experts=self.config.n_routed_experts, ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: @@ -659,7 +528,7 @@ class FlashModel(nn.Module): ) weight_loader(param, loaded_weight) loaded_params.add(name) - for layer_id in range(self.config.num_hidden_layers): + for layer_id in range(self.config.num_layers): for i in range(2): if isinstance(self.layers[layer_id], PPMissingLayer): continue @@ -727,15 +596,10 @@ class LongcatFlashForCausalLM(nn.Module, SupportsLoRA, SupportsPP): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() - config = FlashConfig(**vllm_config.model_config.hf_config.__dict__) + config = vllm_config.model_config.hf_config quant_config = vllm_config.quant_config self.config = config - config.intermediate_size = ( - config.ffn_hidden_size - if hasattr(config, "ffn_hidden_size") - else config.intermediate_size - ) self.quant_config = quant_config diff --git a/vllm/model_executor/models/longcat_flash_mtp.py b/vllm/model_executor/models/longcat_flash_mtp.py index 770f931e81..506f7a7451 100644 --- a/vllm/model_executor/models/longcat_flash_mtp.py +++ b/vllm/model_executor/models/longcat_flash_mtp.py @@ -20,7 +20,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( VocabParallelEmbedding, ) from vllm.model_executor.model_loader.weight_utils import default_weight_loader -from vllm.model_executor.models.longcat_flash import FlashConfig from vllm.sequence import IntermediateTensors from .deepseek_v2 import DeepseekV2DecoderLayer @@ -80,9 +79,9 @@ class LongCatMultiTokenPredictor(nn.Module): prefix: str = "", ): super().__init__() - config = FlashConfig(**vllm_config.model_config.hf_config.__dict__) - vllm_config.model_config.hf_config.intermediate_size = config.intermediate_size - self.mtp_start_layer_idx = config.num_hidden_layers * 2 + config = vllm_config.model_config.hf_config + # TODO: Use `num_hidden_layers` once Transformers 5.18.0 is the minimum + self.mtp_start_layer_idx = config.num_layers * 2 self.num_mtp_layers = 1 self.layers = torch.nn.ModuleDict( { @@ -130,7 +129,7 @@ class LongCatFlashMTP(nn.Module): # builds a dense MLP. object.__setattr__ bypasses the ngram remote # config's strict validation (it rejects setting the int field to None). object.__setattr__(vllm_config.model_config.hf_config, "n_routed_experts", None) - self.config = FlashConfig(**vllm_config.model_config.hf_config.__dict__) + self.config = vllm_config.model_config.hf_config self.quant_config = ( None if "mtp" in getattr(self.config, "disable_quant_module", []) @@ -264,7 +263,7 @@ class LongCatFlashMTP(nn.Module): weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, loaded_weight) loaded_params.add(name) - spec_layer_id = self.config.num_hidden_layers * 2 + spec_layer_id = self.model.mtp_start_layer_idx self_attn = self.model.layers[str(spec_layer_id)].mtp_block.self_attn if ( self.quant_config is not None @@ -356,5 +355,5 @@ class LongCatFlashMTP(nn.Module): self, config: PreTrainedConfig, weight_name: str ) -> int | None: if "model.mtp" in weight_name: - return config.num_hidden_layers * 2 + return self.model.mtp_start_layer_idx return None diff --git a/vllm/model_executor/models/longcat_flash_ngram.py b/vllm/model_executor/models/longcat_flash_ngram.py index 435a11ee17..aab615845d 100644 --- a/vllm/model_executor/models/longcat_flash_ngram.py +++ b/vllm/model_executor/models/longcat_flash_ngram.py @@ -15,6 +15,7 @@ from typing import Any import torch from torch import nn +from transformers import LongcatFlashConfig from vllm import _custom_ops as ops from vllm.config import VllmConfig @@ -30,15 +31,15 @@ from vllm.v1.worker.gpu.model_states.default import DefaultModelState from vllm.v1.worker.gpu.states import RequestState from .interfaces import SupportsLoRA, SupportsPP -from .longcat_flash import FlashConfig, FlashModel +from .longcat_flash import FlashModel from .utils import AutoWeightsLoader, PPMissingLayer, WeightsMapper, maybe_prefix -def uses_ngram_embedding(config: FlashConfig) -> bool: +def uses_ngram_embedding(config: LongcatFlashConfig) -> bool: return getattr(config, "ngram_vocab_size_ratio", None) is not None -def _config_dtype(config: FlashConfig) -> torch.dtype: +def _config_dtype(config: LongcatFlashConfig) -> torch.dtype: dt = getattr(config, "torch_dtype", None) or getattr(config, "dtype", None) if isinstance(dt, torch.dtype): return dt @@ -54,7 +55,7 @@ class NgramEmbedding(nn.Module): single ``bmm``. Hashing math is ported from the HF reference. """ - def __init__(self, config: FlashConfig, base_embeddings: nn.Module) -> None: + def __init__(self, config: LongcatFlashConfig, base_embeddings: nn.Module) -> None: super().__init__() self.config = config self.word_embeddings = base_embeddings @@ -170,14 +171,6 @@ class FlashNgramModel(FlashModel): """FlashModel whose input embedding is an :class:`NgramEmbedding`.""" def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: - # Each FlashDecoderLayer is a *dual* layer (2 attentions), so the number - # of decoder layers is ``num_layers``. The ngram HF config sets - # ``num_hidden_layers`` to a multiple of that (attention-module count), - # which FlashModel would otherwise build as too many (dead) layers. - hf = vllm_config.model_config.hf_config - num_layers = getattr(hf, "num_layers", None) - if num_layers is not None and hf.num_hidden_layers != num_layers: - hf.num_hidden_layers = num_layers super().__init__(vllm_config=vllm_config, prefix=prefix) self.ngram_embeddings: NgramEmbedding | None if get_pp_group().is_first_rank and uses_ngram_embedding(self.config): @@ -222,10 +215,7 @@ class LongcatFlashNgramForCausalLM(nn.Module, SupportsLoRA, SupportsPP): "V2 model runner for its n-gram embedding state; it is selected " "automatically unless VLLM_USE_V2_MODEL_RUNNER=0 is set." ) - config = FlashConfig(**vllm_config.model_config.hf_config.__dict__) - config.intermediate_size = getattr( - config, "ffn_hidden_size", config.intermediate_size - ) + config = vllm_config.model_config.hf_config self.config = config self.quant_config = vllm_config.quant_config diff --git a/vllm/models/hy_v4/nvidia/attention.py b/vllm/models/hy_v4/nvidia/attention.py index 1e9be2e3bf..48d769a856 100644 --- a/vllm/models/hy_v4/nvidia/attention.py +++ b/vllm/models/hy_v4/nvidia/attention.py @@ -47,7 +47,13 @@ from vllm.v1.attention.selector import get_attn_backend logger = init_logger(__name__) -_SPARSE_LAYER_TYPES = ("sparse_attention", "sparse", "deepseek_sparse_attention") +_SPARSE_LAYER_TYPES = ( + "sparse_attention", + "sparse", + "indexed_attention", + # TODO: Delete below once Transformers 5.18.0 is the minimum required version. + "deepseek_sparse_attention", +) _WEIGHT_LAYER_INDEX_RE = re.compile(r"(?:^|\.)layers\.(\d+)(?:\.|$)") diff --git a/vllm/models/qwen4_exp/amd/model.py b/vllm/models/qwen4_exp/amd/model.py index 16509d609c..b3a2b75053 100644 --- a/vllm/models/qwen4_exp/amd/model.py +++ b/vllm/models/qwen4_exp/amd/model.py @@ -79,6 +79,10 @@ from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm from .ple_layer import Qwen4ExpPLELayer from .qsa import Qwen4ExpQSAAttention +# Transformers v5.18 renamed `qwen_sparse_attention` to `indexed_attention` +# TODO: Delete qwen_... once Transformers 5.18.0 is the minimum required version. +_QSA_LAYER_TYPES = ("qwen_sparse_attention", "indexed_attention") + def without_modelopt_fp4( quant_config: QuantizationConfig | None, @@ -207,7 +211,7 @@ class Qwen4ExpDecoderLayer(nn.Module): prefix=f"{prefix}.linear_attn", gqa_interleaved_layout=False, ) - elif layer_type == "qwen_sparse_attention": + elif layer_type in _QSA_LAYER_TYPES: use_qsa = getattr(config, "indexer_n_heads", None) is not None if not use_qsa: self.self_attn = Qwen3NextAttention( @@ -289,7 +293,7 @@ class Qwen4ExpDecoderLayer(nn.Module): if self.layer_type == "linear_attention": attn_out = self.linear_attn(hidden_states=block_input) - elif self.layer_type == "qwen_sparse_attention": + elif self.layer_type in _QSA_LAYER_TYPES: attn_out = self.self_attn( hidden_states=block_input, positions=positions, @@ -380,7 +384,7 @@ class Qwen4ExpModel(nn.Module): self._qsa_layer_ids = frozenset( layer_idx for layer_idx, layer_type in enumerate(config.layer_types) - if layer_type == "qwen_sparse_attention" + if layer_type in _QSA_LAYER_TYPES and getattr(config, "indexer_n_heads", None) is not None ) self.embed_tokens = VocabParallelEmbedding(self.vocab_size, config.hidden_size) diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py index 95e920bc21..73fe2cc468 100644 --- a/vllm/models/qwen4_exp/nvidia/model.py +++ b/vllm/models/qwen4_exp/nvidia/model.py @@ -79,6 +79,10 @@ from .ops.cute_dsl.hc_down_silu import request_hc_down_silu_warmup from .ple_layer import Qwen4ExpPLELayer from .qsa import Qwen4ExpQSAAttention +# Transformers v5.18 renamed `qwen_sparse_attention` to `indexed_attention` +# TODO: Delete qwen_... once Transformers 5.18.0 is the minimum required version. +_QSA_LAYER_TYPES = ("qwen_sparse_attention", "indexed_attention") + def without_modelopt_fp4( quant_config: QuantizationConfig | None, @@ -209,7 +213,7 @@ class Qwen4ExpDecoderLayer(nn.Module): prefix=f"{prefix}.linear_attn", gqa_interleaved_layout=False, ) - elif layer_type == "qwen_sparse_attention": + elif layer_type in _QSA_LAYER_TYPES: use_qsa = getattr(config, "indexer_n_heads", None) is not None if not use_qsa: self.self_attn = Qwen3NextAttention( @@ -293,7 +297,7 @@ class Qwen4ExpDecoderLayer(nn.Module): if self.layer_type == "linear_attention": attn_out = self.linear_attn(hidden_states=block_input) - elif self.layer_type == "qwen_sparse_attention": + elif self.layer_type in _QSA_LAYER_TYPES: attn_out = self.self_attn( hidden_states=block_input, positions=positions, @@ -373,7 +377,7 @@ class Qwen4ExpModel(nn.Module): self._qsa_layer_ids = frozenset( layer_idx for layer_idx, layer_type in enumerate(config.layer_types) - if layer_type == "qwen_sparse_attention" + if layer_type in _QSA_LAYER_TYPES and getattr(config, "indexer_n_heads", None) is not None ) self.embed_tokens = VocabParallelEmbedding(self.vocab_size, config.hidden_size) diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index 0766f9b9ba..e49f894a85 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -147,6 +147,12 @@ _CONFIG_REGISTRY: dict[str, type[PreTrainedConfig]] = LazyConfigDict( _SPECULATIVE_DECODING_CONFIGS: set[str] = {"eagle", "speculators", "medusa"} +# Checkpoints whose config.json has no model_type, only remote code in auto_map +_ARCHITECTURE_TO_MODEL_TYPE: dict[str, str] = { + "LongcatFlashForCausalLM": "longcat_flash", + "LongcatFlashNgramForCausalLM": "longcat_flash", +} + _PATCH_HF_VALIDATE_ROPE: set[str] = {"sarvam_mla"} # Model types whose checkpoints carry shared RoPE parameters alongside the @@ -352,6 +358,8 @@ class HFConfigParser(ConfigParserBase): if config_dict.get("speculators_config") is not None else model_type ) + if model_type is None and (architectures := config_dict.get("architectures")): + model_type = _ARCHITECTURE_TO_MODEL_TYPE.get(architectures[0]) # Allow hf_overrides to override model_type before checking _CONFIG_REGISTRY if (hf_overrides := kwargs.pop("hf_overrides", None)) is not None: if isinstance(hf_overrides, dict) and "model_type" in hf_overrides: @@ -374,6 +382,15 @@ class HFConfigParser(ConfigParserBase): if extra_layer_types := _PATCH_HF_ALLOWED_LAYER_TYPES.get(model_type): _patch_hf_transformers_allowed_layer_types(extra_layer_types) + if model_type == "vlm": + # HyperCLOVAX remote code registers this alias in a bare + # `try/except` that fails silently since transformers 5.18 + from transformers import CONFIG_MAPPING, Qwen2_5_VLVisionConfig + + CONFIG_MAPPING.register( + "qwen2_5_vl_visual", Qwen2_5_VLVisionConfig, exist_ok=True + ) + rope_parameters = config_dict.get("rope_parameters") or {} if model_type == "gemma4_text" and "full_attention" in rope_parameters: from transformers import Gemma4TextConfig @@ -387,6 +404,11 @@ class HFConfigParser(ConfigParserBase): } kwargs.setdefault("name_or_path", str(model)) config = Gemma4TextConfig.from_dict(config_dict, **kwargs) + elif model_type in _ARCHITECTURE_TO_MODEL_TYPE.values(): + from transformers import CONFIG_MAPPING + + kwargs.setdefault("name_or_path", str(model)) + config = CONFIG_MAPPING[model_type].from_dict(config_dict, **kwargs) elif model_type in _SPECULATIVE_DECODING_CONFIGS: config_class = _CONFIG_REGISTRY[model_type] config = config_class.from_pretrained( diff --git a/vllm/transformers_utils/model_arch_config_convertor.py b/vllm/transformers_utils/model_arch_config_convertor.py index 9386e9a328..9dee4af924 100644 --- a/vllm/transformers_utils/model_arch_config_convertor.py +++ b/vllm/transformers_utils/model_arch_config_convertor.py @@ -329,7 +329,6 @@ class ModelArchConfigConvertorBase: "kimi_k2", "kimi_linear", "longcat_flash", - "longcat_flash_ngram", "pangu_ultra_moe", "pangu_ultra_moe_mtp", "bailing_hybrid", @@ -714,6 +713,12 @@ class PanguUltraMoeMTPModelArchConfigConvertor(ModelArchConfigConvertorBase): return getattr(self.hf_text_config, "num_nextn_predict_layers", 0) +class LongCatFlashModelArchConfigConvertor(ModelArchConfigConvertorBase): + def get_num_hidden_layers(self) -> int: + # `num_hidden_layers` counts attention sublayers, two per decoder layer + return self.hf_text_config.num_layers + + class LongCatFlashMTPModelArchConfigConvertor(ModelArchConfigConvertorBase): def get_num_hidden_layers(self) -> int: return getattr(self.hf_text_config, "num_nextn_predict_layers", 1) @@ -825,6 +830,7 @@ MODEL_ARCH_CONFIG_CONVERTORS = { "gemma4_unified_text": Gemma4ModelArchConfigConvertor, "glm4_moe_mtp": GLM4MoeMTPModelArchConfigConvertor, "glm_ocr_mtp": GLM4MoeMTPModelArchConfigConvertor, + "longcat_flash": LongCatFlashModelArchConfigConvertor, "longcat_flash_mtp": LongCatFlashMTPModelArchConfigConvertor, "mamba": MambaModelArchConfigConvertor, "medusa": MedusaModelArchConfigConvertor, diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index 9e9ffd7be1..4fa308bcf4 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -398,10 +398,7 @@ def init_kv_cache( # Dual-attention models (e.g. LongCat-Flash) put two Attention modules per # decoder layer, so a layer name carries two integers (layer + module index). num_attn_module = ( - 2 - if vllm_config.model_config.hf_config.model_type - in ("longcat_flash", "longcat_flash_ngram") - else 1 + 2 if vllm_config.model_config.hf_config.model_type == "longcat_flash" else 1 ) bindable_caches = { name: cache for name, cache in kv_caches.items() if name in forward_context From 5e56e9af2aaede223a5f723cba88fa35999ac813 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Sun, 4 Oct 2026 09:13:10 +0800 Subject: [PATCH 04/22] [Feature] Release the CUDA graph pool on sleep (#59160) Signed-off-by: aoshen02 Co-authored-by: Claude Opus 5.5 --- docs/features/sleep_mode.md | 35 +++++++++ .../memory/cumem/test_cumem.py | 55 ++++++++++++++ .../memory/sleep_mode/test_sleep_mode.py | 75 ++++++++++++++++++- tests/distributed/test_engram_dp_shard.py | 4 +- tests/v1/cudagraph/test_cudagraph_manager.py | 1 + vllm/compilation/breakable_cudagraph.py | 13 ++-- vllm/compilation/cuda_graph.py | 11 ++- vllm/compilation/cudagraph_pool.py | 45 +++++++++++ vllm/config/model.py | 3 + vllm/config/vllm.py | 23 ++++++ vllm/device_allocator/cumem.py | 17 +++++ vllm/device_allocator/sleep_mode_backend.py | 7 +- .../device_communicators/cuda_communicator.py | 6 ++ .../device_communicators/custom_all_reduce.py | 6 +- vllm/engine/arg_utils.py | 6 ++ vllm/v1/worker/gpu/cudagraph_utils.py | 19 +++-- 16 files changed, 299 insertions(+), 27 deletions(-) create mode 100644 vllm/compilation/cudagraph_pool.py diff --git a/docs/features/sleep_mode.md b/docs/features/sleep_mode.md index 3e419707f2..a8aa15af7b 100644 --- a/docs/features/sleep_mode.md +++ b/docs/features/sleep_mode.md @@ -101,6 +101,41 @@ llm.release_kv_cache_memory() llm.wake_up(tags=["kv_cache"]) # Reallocate KV cache and resume scheduling. ``` +#### Offloading CUDA graph memory + +By default, CUDA graph memory stays on the GPU while asleep. With +`sleep_mode_offload_cudagraph=True` (off by default), CUDA graphs are captured +into a cuMem pool that sleep backs up to CPU memory at both levels and any wake +restores in place, so graphs are replayed, not recaptured. It needs the default +`cumem` backend, CUDA and CUDA graphs; otherwise it has no effect. + +```python +llm = LLM("Qwen/Qwen3-8B", enable_sleep_mode=True, sleep_mode_offload_cudagraph=True) +``` + +or `vllm serve --enable-sleep-mode --sleep-mode-offload-cudagraph`. + +The cost is pinned host memory for the pool's backup, also at level 2, and one +extra copy per captured custom allreduce (about 3% decode latency at batch +size 1). `NCCL_GRAPH_REGISTER` defaults to `0`, since NCCL graph registration +would pin the pool; an explicit value is kept with a warning. Graph executables +outside PyTorch pools stay resident. + +#### Releasing NCCL communicator memory + +By default, NCCL communicators keep their GPU buffers while asleep. With +`enable_nccl_comm_suspend=True` (off by default, experimental), sleep releases +them with `ncclCommSuspend` and wake restores them with `ncclCommResume`. The +communicators keep their topology, so they are not re-created. It needs NCCL +2.29.7 or newer; with an older library, a warning is logged and the memory stays +on the GPU. + +```python +llm = LLM("Qwen/Qwen3-8B", enable_sleep_mode=True, enable_nccl_comm_suspend=True) +``` + +or `vllm serve --enable-sleep-mode --enable-nccl-comm-suspend`. + ### Online Serving To enable sleep mode in a vLLM server you need to initialize it with the flag `VLLM_SERVER_DEV_MODE=1` and pass `--enable-sleep-mode` to the vLLM server. diff --git a/tests/basic_correctness/memory/cumem/test_cumem.py b/tests/basic_correctness/memory/cumem/test_cumem.py index fa06facf29..df8c604882 100644 --- a/tests/basic_correctness/memory/cumem/test_cumem.py +++ b/tests/basic_correctness/memory/cumem/test_cumem.py @@ -423,3 +423,58 @@ def test_cumem_with_cudagraph(): # output content is as expected assert torch.allclose(y, x + 1) + + +@pytest.mark.parametrize("level", [1, 2], ids=["sleep-1", "sleep-2"]) +@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn") +@pytest.mark.skipif(not current_platform.is_cuda(), reason="cuMem CUDA graph pool") +def test_cudagraph_pool_sleep(level): + """Routing, and the graph pool backed up at both sleep levels and restored + in place by any wake.""" + from contextlib import nullcontext + from types import SimpleNamespace + + import vllm.distributed.device_communicators.pynccl_allocator as nccl_alloc + from vllm.compilation.cudagraph_pool import ( + capture_outside_cumem_pool, + capture_pool, + ) + from vllm.device_allocator.sleep_mode_backend import CuMemBackend + + allocator = get_mem_allocator_instance() + on, off = (SimpleNamespace(use_cumem_cudagraph_pool=v) for v in (True, False)) + handle = current_platform.graph_pool_handle() + for cfg, ctx in ((off, nullcontext()), (on, capture_outside_cumem_pool())): + with ctx, capture_pool(handle, cfg) as used: + assert used == handle and allocator.current_tag != "cudagraph" + with allocator.use_memory_pool("weights"): + weight = torch.full((1 << 20,), 2.0, device=DEVICE_TYPE) + x = torch.ones_like(weight) + + def capture() -> list: + graph = torch.cuda.CUDAGraph() + stream = torch.cuda.Stream() + with ( + capture_pool(handle, on) as pool, + torch.cuda.graph(graph, pool=pool, stream=stream), + ): + assert nccl_alloc._graph_pool_id == pool != handle + const = torch.empty_like(x) # Never written by replay: needs a backup. + y = x * weight + const + return [graph, y, const.fill_(3.0)] + + def graph_pool() -> dict[int, int]: + data = allocator.pointer_to_data.items() + return { + p: d.handle[1] for p, d in data if d.tag == "cudagraph" and not d.is_asleep + } + + held, backend = capture(), CuMemBackend() + mapped = graph_pool() + backend.suspend(level=level) + backend.resume(tags=["kv_cache"]) + assert graph_pool() == mapped + backend.resume(tags=["weights"]) + weight.fill_(2.0) # Level 2 discards weights; emulate the reload. + held[0].replay() + assert torch.equal(held[1], torch.full_like(x, 5.0)) diff --git a/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py b/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py index 3ed4703c8b..73af6e3910 100644 --- a/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py +++ b/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py @@ -3,6 +3,7 @@ import asyncio import math +import os import pytest import torch @@ -14,7 +15,12 @@ from vllm.model_executor.model_loader import get_model_loader from vllm.platforms import current_platform from vllm.utils.mem_constants import GiB_bytes -from ....utils import create_new_process_for_each_test, multi_gpu_test, requires_fp8 +from ....utils import ( + create_new_process_for_each_test, + multi_gpu_marks, + multi_gpu_test, + requires_fp8, +) @create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn") @@ -195,6 +201,73 @@ def test_deep_sleep_lora_tp2(monkeypatch): assert output[0].outputs[0].text == output2[0].outputs[0].text +def _cudagraph_bytes(worker) -> int: + data = cumem.CuMemAllocator.get_instance().pointer_to_data.values() + return sum(d.handle[1] for d in data if d.tag == "cudagraph") + + +def _custom_ar_active(worker) -> bool: + from vllm.distributed.parallel_state import get_tp_group + + ca_comm = get_tp_group().device_communicator.ca_comm + return ca_comm is not None and not ca_comm.disabled + + +@pytest.mark.parametrize( + ("mode", "breakable", "tp", "offload"), + [ + ("FULL", False, 1, True), + ("PIECEWISE", False, 1, True), + ("PIECEWISE", True, 1, True), + pytest.param( + "FULL_AND_PIECEWISE", False, 2, True, marks=multi_gpu_marks(num_gpus=2) + ), + ("FULL_AND_PIECEWISE", False, 1, False), + ], + ids=["full", "piecewise", "breakable", "custom-ar-tp2", "off-by-default"], +) +@create_new_process_for_each_test() +@pytest.mark.skipif(not current_platform.is_cuda(), reason="cuMem CUDA graph pool") +def test_sleep_cudagraph_pool(monkeypatch, mode, breakable, tp, offload): + """Each capture site uses the pool, exact across sleeps; TP=2 needs + unregistered custom AR; off changes nothing.""" + for name, value in [ + ("VLLM_USE_V2_MODEL_RUNNER", "1"), + ("VLLM_USE_BREAKABLE_CUDAGRAPH", "1" if breakable else "0"), + ("VLLM_ALLOW_INSECURE_SERIALIZATION", "1"), + ("VLLM_ALLREDUCE_USE_FLASHINFER", "0"), + ("VLLM_ALLREDUCE_USE_SYMM_MEM", "0"), + ]: + monkeypatch.setenv(name, value) + monkeypatch.delenv("NCCL_GRAPH_REGISTER", raising=False) + llm = LLM( + "hmellor/tiny-random-LlamaForCausalLM", + enable_sleep_mode=True, + sleep_mode_offload_cudagraph=offload, + tensor_parallel_size=tp, + compilation_config={ + "cudagraph_mode": mode, + "pass_config": {"fuse_allreduce_rms": False}, + }, + ) + graph_bytes = llm.collective_rpc(_cudagraph_bytes) + assert os.environ.get("NCCL_GRAPH_REGISTER") == ("0" if offload else None) + assert all(graph_bytes) if offload else not any(graph_bytes) + if not offload: + return + if tp > 1 and not all(llm.collective_rpc(_custom_ar_active)): + pytest.skip("Custom allreduce is unavailable on these GPUs") + prompt, params = "How are you?", SamplingParams(temperature=0, max_tokens=10) + expected = llm.generate(prompt, params)[0].outputs[0].text + for level in (1, 2): + llm.sleep(level=level) + llm.wake_up(tags=["weights"]) + if level == 2: + llm.collective_rpc("reload_weights") + llm.wake_up(tags=["kv_cache"]) + assert llm.generate(prompt, params)[0].outputs[0].text == expected + + @create_new_process_for_each_test() def test_deep_sleep_async(gpu_memory_cleared): async def test(): diff --git a/tests/distributed/test_engram_dp_shard.py b/tests/distributed/test_engram_dp_shard.py index ac73188534..637d179a5d 100644 --- a/tests/distributed/test_engram_dp_shard.py +++ b/tests/distributed/test_engram_dp_shard.py @@ -177,7 +177,9 @@ def _worker(rank: int, tp_size: int, port: int) -> None: # Supply the resolved Engram settings without constructing a full model. vllm_config.engram_config = EngramConfig() vllm_config.model_config = SimpleNamespace( - architecture="DeepseekV41ForCausalLM", is_moe=True + architecture="DeepseekV41ForCausalLM", + is_moe=True, + sleep_mode_offload_cudagraph=False, ) try: init_distributed_environment() diff --git a/tests/v1/cudagraph/test_cudagraph_manager.py b/tests/v1/cudagraph/test_cudagraph_manager.py index 8f58462d6b..154116b80f 100644 --- a/tests/v1/cudagraph/test_cudagraph_manager.py +++ b/tests/v1/cudagraph/test_cudagraph_manager.py @@ -50,6 +50,7 @@ def _create_vllm_config() -> MagicMock: vllm_config.parallel_config = ParallelConfig() vllm_config.speculative_config = None vllm_config.num_speculative_tokens = 0 + vllm_config.use_cumem_cudagraph_pool = False return vllm_config diff --git a/vllm/compilation/breakable_cudagraph.py b/vllm/compilation/breakable_cudagraph.py index 781868804e..b70fe70d8c 100644 --- a/vllm/compilation/breakable_cudagraph.py +++ b/vllm/compilation/breakable_cudagraph.py @@ -34,9 +34,9 @@ from typing import Any, ClassVar, TypeVar import torch import vllm.envs as envs +from vllm.compilation.cudagraph_pool import capture_pool from vllm.compilation.monitor import validate_cudagraph_capturing_enabled from vllm.config import CUDAGraphMode, VllmConfig -from vllm.distributed.device_communicators.pynccl_allocator import set_graph_pool_id from vllm.forward_context import ( BatchDescriptor, get_forward_context, @@ -370,11 +370,6 @@ class BreakableCUDAGraphWrapper: entry.input_addresses = self._collect_tensor_addresses(args, kwargs) - if self.graph_pool is not None: - set_graph_pool_id(self.graph_pool) - else: - set_graph_pool_id(current_platform.graph_pool_handle()) - # Match torch.cuda.graph()'s pre-capture cleanup, which we bypass. # Skip it when gc is disabled: bulk capture runs under # freeze_gc_for_cudagraph_capture, which already did this cleanup, @@ -386,8 +381,10 @@ class BreakableCUDAGraphWrapper: # pre-capture prefetches are complete and don't leak into the graph. get_offloader().sync_prev_onload() - capture = BreakableCUDAGraphCapture(pool=self.graph_pool) - with capture: + with ( + capture_pool(self.graph_pool, self.vllm_config) as pool, + BreakableCUDAGraphCapture(pool=pool) as capture, + ): output = self.runnable(*args, **kwargs) # Join the offloader's copy stream while we still hold the last # segment open, so the join is captured into the graph (otherwise diff --git a/vllm/compilation/cuda_graph.py b/vllm/compilation/cuda_graph.py index dc8a8a9c84..2dea34e0ff 100644 --- a/vllm/compilation/cuda_graph.py +++ b/vllm/compilation/cuda_graph.py @@ -13,9 +13,9 @@ from unittest.mock import patch import torch from vllm.compilation.counter import compilation_counter +from vllm.compilation.cudagraph_pool import capture_pool from vllm.compilation.monitor import validate_cudagraph_capturing_enabled from vllm.config import CUDAGraphMode, VllmConfig -from vllm.distributed.device_communicators.pynccl_allocator import set_graph_pool_id from vllm.forward_context import ( BatchDescriptor, get_forward_context, @@ -302,10 +302,9 @@ class CUDAGraphWrapper: ) ) - if self.graph_pool is not None: - set_graph_pool_id(self.graph_pool) - else: - set_graph_pool_id(current_platform.graph_pool_handle()) + graph_pool = stack.enter_context( + capture_pool(self.graph_pool, self.vllm_config) + ) # Sync offloader's copy stream before capture. # Ensure any pre-capture prefetches from offloader are complete. @@ -314,7 +313,7 @@ class CUDAGraphWrapper: # mind-exploding: carefully manage the reference and memory. with torch.cuda.graph( cudagraph, - pool=self.graph_pool, + pool=graph_pool, stream=current_stream(), ): # `output` is managed by pytorch's cudagraph pool diff --git a/vllm/compilation/cudagraph_pool.py b/vllm/compilation/cudagraph_pool.py new file mode 100644 index 0000000000..c4cfeabd81 --- /dev/null +++ b/vllm/compilation/cudagraph_pool.py @@ -0,0 +1,45 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Graph pool routing for CUDA graph capture under sleep mode.""" + +from collections.abc import Iterator +from contextlib import contextmanager +from contextvars import ContextVar +from typing import TYPE_CHECKING, cast + +from vllm.config import VllmConfig +from vllm.device_allocator import get_mem_allocator_instance +from vllm.distributed.device_communicators.pynccl_allocator import set_graph_pool_id +from vllm.platforms import current_platform + +if TYPE_CHECKING: + from vllm.device_allocator.cumem import CuMemAllocator + +_outside_cumem: ContextVar[bool] = ContextVar("capture_outside_cumem", default=False) + + +@contextmanager +def capture_outside_cumem_pool() -> Iterator[None]: + """CUDA graphs captured here stay out of cuMem, so that memory profiling can + free its throwaway graphs normally.""" + token = _outside_cumem.set(True) + try: + yield + finally: + _outside_cumem.reset(token) + + +@contextmanager +def capture_pool( + pool: tuple[int, int] | None, vllm_config: VllmConfig +) -> Iterator[tuple[int, int] | None]: + """Yield the pool to capture into, the cuMem graph pool when sleep offloads + graph memory, and point NCCL's graph allocator at it.""" + if vllm_config.use_cumem_cudagraph_pool and not _outside_cumem.get(): + allocator = cast("CuMemAllocator", get_mem_allocator_instance()) + with allocator.cudagraph_pool() as cumem_pool: + set_graph_pool_id(cumem_pool) + yield cumem_pool + else: + set_graph_pool_id(pool or current_platform.graph_pool_handle()) + yield pool diff --git a/vllm/config/model.py b/vllm/config/model.py index 3e481d8c91..538bdcc834 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -378,6 +378,9 @@ class ModelConfig: (default) uses the built-in ``CuMemAllocator`` and is behavior-compatible with prior releases. Additional backends (CUDA checkpoint, CRIU, durable snapshot) may be registered in-tree or by plugins (RFC #34303).""" + sleep_mode_offload_cudagraph: bool = False + """Back up CUDA graph memory to CPU during sleep, restored in place on wake. + Takes effect with enable_sleep_mode, the cumem backend and CUDA graphs.""" enable_nccl_comm_suspend: bool = False """Enable releasing NCCL communicator memory during sleep mode (``ncclCommSuspend``/``ncclCommResume``). Experimental; when disabled diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 281666a633..32b49f4272 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -697,6 +697,21 @@ class VllmConfig: """ return 1 + self.num_speculative_tokens + @property + def use_cumem_cudagraph_pool(self) -> bool: + """Whether CUDA graphs go to the cuMem pool that sleep offloads.""" + from vllm.platforms import current_platform + + model_config = self.model_config + return ( + model_config is not None + and model_config.sleep_mode_offload_cudagraph + and model_config.enable_sleep_mode + and model_config.sleep_mode_backend == "cumem" + and self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE + and current_platform.is_cuda() + ) + @property def use_v2_model_runner(self) -> bool: if self.attention_config.hisparse_config is not None: @@ -2321,6 +2336,14 @@ class VllmConfig: custom_ops.append("+quant_fp8") self._verify_kv_transfer_compat() + if self.use_cumem_cudagraph_pool: + # NCCL graph registration pins the offloaded pool; workers inherit this. + value = os.environ.setdefault("NCCL_GRAPH_REGISTER", "0") + if value != "0": + logger.warning( + "NCCL_GRAPH_REGISTER=%s pins the CUDA graph pool during sleep.", + value, + ) # Log the custom passes that are enabled self.compilation_config.pass_config.log_enabled_passes() diff --git a/vllm/device_allocator/cumem.py b/vllm/device_allocator/cumem.py index 6bf0ea71f8..b85be08d28 100644 --- a/vllm/device_allocator/cumem.py +++ b/vllm/device_allocator/cumem.py @@ -430,6 +430,23 @@ class CuMemAllocator: if allocation["allocated_size"] == 0 and data and not data.is_asleep: unmap_and_release(self._python_free_callback(allocation["address"])) + @contextmanager + def cudagraph_pool(self) -> Iterator[tuple[int, int]]: + """Tag CUDA graph capture allocations and yield the graph pool id.""" + if "cudagraph" not in self.allocator_and_pools: + allocator = get_pluggable_allocator( + self.python_malloc_callback, self.python_free_callback + ) + mem_pool = torch.cuda.memory.MemPool(allocator._allocator) + self.allocator_and_pools["cudagraph"] = [(mem_pool, allocator)] + old_tag = self.current_tag + self.current_tag = "cudagraph" + try: + # capture_begin routes allocations to this pool; no use_mem_pool. + yield self.allocator_and_pools["cudagraph"][0][0].id + finally: + self.current_tag = old_tag + def get_current_usage(self) -> int: """Get the total number of bytes allocated in the memory pool.""" sum_bytes: int = 0 diff --git a/vllm/device_allocator/sleep_mode_backend.py b/vllm/device_allocator/sleep_mode_backend.py index 8ec9986ffc..2132b880de 100644 --- a/vllm/device_allocator/sleep_mode_backend.py +++ b/vllm/device_allocator/sleep_mode_backend.py @@ -129,8 +129,11 @@ class CuMemBackend(SleepModeBackend): self._state = "SUSPENDED" allocator = get_mem_allocator_instance() - # Runtime state is kept at every level; weights only at level 1. - offload_tags = ("weights", "runtime") if level == 1 else ("runtime",) + # Runtime state and graph pools are kept at every level; weights only + # at level 1. + offload_tags: tuple[str, ...] = ("runtime", "cudagraph") + if level == 1: + offload_tags = ("weights", *offload_tags) allocator.sleep(offload_tags=offload_tags) def resume(self, tags: list[str] | None = None) -> None: diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index 0a203956f0..99ebcd717d 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -7,6 +7,7 @@ from torch.distributed import ProcessGroup import vllm.envs as envs from vllm._aiter_ops import rocm_aiter_ops +from vllm.config import get_current_vllm_config_or_none from vllm.distributed.device_communicators.all_reduce_utils import ( NCCL_SYMM_MEM_ALL_REDUCE_CONFIG, should_nccl_symm_mem_ag_rs, @@ -114,6 +115,10 @@ class CudaCommunicator(DeviceCommunicatorBase): self.aiter_ar_comm: AiterCustomAllreduce | None = None self.use_aiter_ag_rs: bool = False + # cuMem graph buffers cannot be IPC-registered; capture copies them instead. + config = get_current_vllm_config_or_none() + register = config is None or not config.use_cumem_cudagraph_pool + if use_torch_symm_mem and current_platform.is_cuda(): self.symm_mem_comm = SymmMemCommunicator( group=self.cpu_group, @@ -151,6 +156,7 @@ class CudaCommunicator(DeviceCommunicatorBase): symm_mem_enabled=( self.symm_mem_comm is not None and not self.symm_mem_comm.disabled ), + register_graph_buffers=register, ) # AITER custom all-gather/reduce-scatter DP-attention dispatch/combine diff --git a/vllm/distributed/device_communicators/custom_all_reduce.py b/vllm/distributed/device_communicators/custom_all_reduce.py index 9e8477e494..d249b5e2f0 100644 --- a/vllm/distributed/device_communicators/custom_all_reduce.py +++ b/vllm/distributed/device_communicators/custom_all_reduce.py @@ -147,18 +147,22 @@ class CustomAllreduce: _DEFAULT_MNNVL_MULTIMEM_REDUCE_SCATTER_MAX_SIZE ), symm_mem_enabled=False, + *, + register_graph_buffers: bool = True, ) -> None: """Args: group: the process group to work on. If None, it will use the default process group. device: the device to bind the CustomAllreduce to. If None, it will be bound to f"cuda:{local_rank}". + register_graph_buffers: whether graph capture IPC-registers its buffers. It is the caller's responsibility to make sure each communicator is bind to a unique device, and all communicators in this group are in the same node. """ self._IS_CAPTURING = False + self._capture_registered = register_graph_buffers self._ptr = 0 self.disabled = True self.mnnvl_buffer = None @@ -549,7 +553,7 @@ class CustomAllreduce: return None if self._IS_CAPTURING: if torch.cuda.is_current_stream_capturing(): - return self.all_reduce(input, registered=True) + return self.all_reduce(input, registered=self._capture_registered) else: # If warm up, mimic the allocation pattern since custom # allreduce is out-of-place. diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index f4cbcf9012..4f1f417289 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -770,6 +770,7 @@ class EngineArgs: generation_config: str = ModelConfig.generation_config enable_sleep_mode: bool = ModelConfig.enable_sleep_mode + sleep_mode_offload_cudagraph: bool = ModelConfig.sleep_mode_offload_cudagraph sleep_preserve_parameter_names: list[str] = get_field( ModelConfig, "sleep_preserve_parameter_names" ) @@ -1007,6 +1008,10 @@ class EngineArgs: model_group.add_argument( "--enable-sleep-mode", **model_kwargs["enable_sleep_mode"] ) + model_group.add_argument( + "--sleep-mode-offload-cudagraph", + **model_kwargs["sleep_mode_offload_cudagraph"], + ) model_group.add_argument( "--sleep-preserve-parameter-names", **model_kwargs["sleep_preserve_parameter_names"], @@ -1984,6 +1989,7 @@ class EngineArgs: generation_config=self.generation_config, override_generation_config=self.override_generation_config, enable_sleep_mode=self.enable_sleep_mode, + sleep_mode_offload_cudagraph=self.sleep_mode_offload_cudagraph, sleep_preserve_parameter_names=self.sleep_preserve_parameter_names, enable_cumem_allocator=self.enable_cumem_allocator, enable_nccl_comm_suspend=self.enable_nccl_comm_suspend, diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 19f3157cea..05fcf5f98f 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -19,10 +19,13 @@ from vllm.compilation.breakable_cudagraph import ( ) from vllm.compilation.counter import compilation_counter from vllm.compilation.cuda_graph import CUDAGraphStat, CUDAGraphWrapper +from vllm.compilation.cudagraph_pool import ( + capture_outside_cumem_pool, + capture_pool, +) from vllm.compilation.wrapper import TorchCompileWithNoGuardsWrapper from vllm.config import VllmConfig, set_current_vllm_config from vllm.config.compilation import CUDAGraphMode -from vllm.distributed.device_communicators.pynccl_allocator import set_graph_pool_id from vllm.distributed.parallel_state import ( get_pp_group, graph_capture, @@ -481,15 +484,14 @@ class CudaGraphManager: # Sync offloader's copy stream before capture. # Ensure any pre-capture prefetches from offloader are complete. get_offloader().sync_prev_onload() - if self.pool is not None: - set_graph_pool_id(self.pool) - else: - set_graph_pool_id(current_platform.graph_pool_handle()) if self._capture_mem_samples is not None: torch.accelerator.synchronize() free_before = torch.accelerator.get_memory_info()[0] - with torch.cuda.graph( - graph, self.pool, stream=self._capture_stream(desc) + with ( + capture_pool(self.pool, self.vllm_config) as pool, + torch.cuda.graph( + graph, pool, stream=self._capture_stream(desc) + ), ): forward_fn(CUDAGraphMode.NONE) # Join offloader's copy stream after forward to avoid @@ -941,7 +943,8 @@ def profile_cudagraph_memory(runner: "GPUModelRunner") -> int: mem_samples: list[int] = [] manager._capture_mem_samples = mem_samples - measured = int(runner.capture_model(profile_only=True)) + with capture_outside_cumem_pool(): + measured = int(runner.capture_model(profile_only=True)) # The measured delta covers PIECEWISE, encoder and speculator graphs # plus the sampled FULL graphs; swap the sampled FULL cost for the From ce49174247fd84c322f293a31cbe1d0460aaa7fe Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Elvir=20Crn=C4=8Devi=C4=87?= Date: Sun, 4 Oct 2026 04:18:58 +0200 Subject: [PATCH 05/22] [Bugfix] Fix FlashInfer all_reduce backend selection (#56891) --- tests/distributed/test_comm_ops.py | 30 +++++++++- .../flashinfer_all_reduce.py | 59 ++++++++++++------- 2 files changed, 67 insertions(+), 22 deletions(-) diff --git a/tests/distributed/test_comm_ops.py b/tests/distributed/test_comm_ops.py index 25434a51b5..253ec8a188 100644 --- a/tests/distributed/test_comm_ops.py +++ b/tests/distributed/test_comm_ops.py @@ -376,7 +376,6 @@ def test_flashinfer_standalone_workspace_size( "_resolve_fi_ar_backend", Mock(return_value=("mnnvl", False)), ) - monkeypatch.setattr(flashinfer_all_reduce, "get_node_count", lambda: 2) monkeypatch.setattr( flashinfer_all_reduce, "_get_tuned_standalone_max_size", @@ -389,6 +388,34 @@ def test_flashinfer_standalone_workspace_size( assert create_workspace.call_args.args[3] == expected +@pytest.mark.parametrize( + ("node_count", "is_blackwell", "expected"), + [ + (1, False, ("trtllm", False)), + (1, True, ("trtllm", False)), + (2, False, (None, False)), + (2, True, ("mnnvl", False)), + ], +) +def test_flashinfer_auto_backend_uses_group_topology( + monkeypatch: pytest.MonkeyPatch, + node_count: int, + is_blackwell: bool, + expected: tuple[str | None, bool], +) -> None: + monkeypatch.setattr( + flashinfer_all_reduce.envs, "VLLM_FLASHINFER_ALLREDUCE_BACKEND", "auto" + ) + monkeypatch.setattr(flashinfer_all_reduce, "_node_count", lambda _: node_count) + monkeypatch.setattr( + flashinfer_all_reduce.current_platform, + "has_device_capability", + lambda capability: is_blackwell, + ) + + assert flashinfer_all_reduce._resolve_fi_ar_backend(Mock()) == expected + + def test_flashinfer_workspace_failure_is_not_retried( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -402,7 +429,6 @@ def test_flashinfer_workspace_failure_is_not_retried( "_resolve_fi_ar_backend", Mock(return_value=("mnnvl", True)), ) - monkeypatch.setattr(flashinfer_all_reduce, "get_node_count", lambda: 1) monkeypatch.setattr( flashinfer_all_reduce, "_get_tuned_standalone_max_size", Mock(return_value=None) ) diff --git a/vllm/distributed/device_communicators/flashinfer_all_reduce.py b/vllm/distributed/device_communicators/flashinfer_all_reduce.py index 9bb17e1528..d51e9e0c0e 100644 --- a/vllm/distributed/device_communicators/flashinfer_all_reduce.py +++ b/vllm/distributed/device_communicators/flashinfer_all_reduce.py @@ -17,7 +17,7 @@ from vllm.config.compilation import PassConfig from vllm.distributed.device_communicators.all_reduce_utils import ( FI_MNNVL_ALLREDUCE_MAX_SIZE_MB, ) -from vllm.distributed.parallel_state import _node_count, get_node_count +from vllm.distributed.parallel_state import _node_count from vllm.logger import init_logger from vllm.platforms import current_platform from vllm.utils.gpu_sync_debug import gpu_sync_allowed @@ -140,14 +140,13 @@ def _create_workspace( return workspace -def _resolve_fi_ar_backend() -> tuple[str, bool]: +def _resolve_fi_ar_backend(group: ProcessGroup) -> tuple[str | None, bool]: """Resolve the flashinfer allreduce backend for the current setup. Returns: - A ``(backend, allow_trtllm_fallback)`` tuple. ``allow_trtllm_fallback`` - is True only when ``auto`` selects mnnvl for a single node, so that - workspace creation can fall back to trtllm on single-node topologies - without NVSwitch multicast support (where mnnvl is unavailable). + A ``(backend, allow_trtllm_fallback)`` tuple. ``backend`` is ``None`` + when automatic selection has no supported backend for the process + group. """ backend = envs.VLLM_FLASHINFER_ALLREDUCE_BACKEND @@ -155,15 +154,23 @@ def _resolve_fi_ar_backend() -> tuple[str, bool]: logger.debug_once("Using flashinfer allreduce backend: %s", backend) return backend, False - # Default to mnnvl for both single- and multi-node setups. The mnnvl - # cudagraph hang that previously forced single-node to trtllm - # (https://github.com/vllm-project/vllm/issues/35772) was fixed upstream in - # FlashInfer (>= 0.6.12, vLLM pins 0.6.15), so mnnvl is safe here. trtllm - # does not support multi-node allreduce, so mnnvl is required there anyway. - # mnnvl needs NVSwitch multicast; on single-node topologies without it, - # fall back to trtllm so fused allreduce stays enabled. - backend = "mnnvl" - allow_trtllm_fallback = get_node_count() == 1 + node_count = _node_count(group) + if node_count == 1: + # TRTLLM supports single-node allreduce on Hopper and Blackwell. Do not + # create an MNNVL workspace when the process group is node-local. + backend = "trtllm" + allow_trtllm_fallback = False + elif current_platform.has_device_capability(100): + # MNNVL requires Blackwell-class GPUs and is the only supported + # FlashInfer backend for a multi-node allreduce process group. + backend = "mnnvl" + allow_trtllm_fallback = False + else: + logger.debug_once( + "FlashInfer allreduce fusion is disabled: multi-node process groups " + "require an MNNVL-capable Blackwell system." + ) + return None, False logger.debug_once("Auto-selected flashinfer allreduce backend: %s", backend) return backend, allow_trtllm_fallback @@ -187,9 +194,12 @@ def get_fi_ar_workspace( if _fi_ar_workspace is not None or _fi_ar_workspace_failed: return _fi_ar_workspace - backend, allow_trtllm_fallback = _resolve_fi_ar_backend() + backend, allow_trtllm_fallback = _resolve_fi_ar_backend(group) - if get_node_count() > 1 and backend == "trtllm": + if backend is None: + return None + + if backend == "trtllm" and _node_count(group) > 1: raise ValueError( "Flashinfer allreduce is not supported for multi-node allreduce with " "'trtllm' backend. Please use 'mnnvl' backend instead." @@ -253,9 +263,12 @@ def get_fi_ar_quant_workspace( if _fi_ar_quant_workspace is not None or _fi_ar_quant_workspace_failed: return _fi_ar_quant_workspace - backend, allow_trtllm_fallback = _resolve_fi_ar_backend() + backend, allow_trtllm_fallback = _resolve_fi_ar_backend(group) - if get_node_count() > 1 and backend == "trtllm": + if backend is None: + return None + + if backend == "trtllm" and _node_count(group) > 1: raise ValueError( "Flashinfer allreduce quantization fusion is not supported for " "multi-node allreduce with 'trtllm' backend. Please use 'mnnvl' " @@ -397,7 +410,13 @@ class FlashInferAllReduce: self.world_size, ) return - backend, _ = _resolve_fi_ar_backend() + backend, _ = _resolve_fi_ar_backend(self.group) + if backend is None: + logger.info( + "FlashInfer All Reduce is disabled because no supported backend " + "is available for this process group." + ) + return tuned_max_size = _get_tuned_standalone_max_size( self.world_size, backend, From 50b404e71e7cfaec7ac85dcf8304bfc15a5a6bb6 Mon Sep 17 00:00:00 2001 From: Ting SUN Date: Sun, 4 Oct 2026 04:57:26 +0200 Subject: [PATCH 06/22] [Bugfix][Frontend] Preserve abort finish_reason for scale-out token streams (#47933) Signed-off-by: Ting Sun Signed-off-by: Flora Feng <4florafeng@gmail.com> Co-authored-by: Flora Feng <4florafeng@gmail.com> --- .../test_generate_stream.py | 62 ++++++++++++++++--- .../token_in_token_out/test_serving_tokens.py | 18 +++--- .../scale_out/token_in_token_out/serving.py | 28 ++++----- 3 files changed, 76 insertions(+), 32 deletions(-) diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py b/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py index 6b6d732f07..7e7f74de9f 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py @@ -456,11 +456,13 @@ async def test_stream_returns_spec_decode_metrics_on_empty_terminal_output( data_chunks = [chunk for chunk in _parse_sse_chunks(chunks) if chunk != "[DONE]"] assert "metrics" not in data_chunks[0] - assert len(data_chunks) == 1 + int(include_usage) + assert data_chunks[1]["choices"][0]["finish_reason"] == "stop" + assert "metrics" not in data_chunks[1] + assert len(data_chunks) == 2 + int(include_usage) if include_usage: - assert data_chunks[1]["choices"] == [] + assert data_chunks[2]["choices"] == [] assert ( - data_chunks[1]["metrics"]["speculative_decoding"]["num_draft_tokens"] == 3 + data_chunks[2]["metrics"]["speculative_decoding"]["num_draft_tokens"] == 3 ) @@ -823,30 +825,76 @@ async def test_stream_zero_token_completion_still_delivers_prompt_metadata(): @pytest.mark.asyncio -async def test_stream_zero_token_completion_emits_no_chunk_by_default(): +@pytest.mark.parametrize("finish_reason", ["stop", "length", "abort"]) +async def test_stream_emits_terminal_output_without_new_tokens(finish_reason): + """The final output may carry no new tokens (always for an abort); the + client must still see why the choice finished.""" engine = _mock_engine() async def mock_generate(*args, **kwargs): + yield _make_request_output("req-1", token_ids=[10]) yield _make_request_output( - "req-1", token_ids=[], finish_reason="stop", finished=True + "req-1", token_ids=[], finish_reason=finish_reason, finished=True ) engine.generate = MagicMock(side_effect=mock_generate) serving = _build_serving_tokens(engine) request = GenerateRequest( token_ids=[1, 2, 3], - sampling_params=SamplingParams(max_tokens=1), + sampling_params=SamplingParams(max_tokens=10), model=MODEL_NAME, stream=True, ) response = await serving.serve_tokens(request) parsed = _parse_sse_chunks([chunk async for chunk in response]) + choices = [c["choices"][0] for c in parsed[:-1]] - assert not [c for c in parsed if isinstance(c, dict) and c.get("choices")] + assert [c["token_ids"] for c in choices] == [[10], []] + assert [c["finish_reason"] for c in choices] == [None, finish_reason] assert parsed[-1] == "[DONE]" +@pytest.mark.asyncio +async def test_stream_parallel_sampling_first_output_missing_a_choice(): + """With n=2 the first output may carry only one choice; the other must + still stream, and an abort finishes both.""" + engine = _mock_engine() + + async def mock_generate(*args, **kwargs): + yield _make_request_output("req-1", token_ids=[10], index=0) + yield _make_request_output("req-1", token_ids=[11], index=1) + aborted = _make_request_output( + "req-1", token_ids=[], finish_reason="abort", finished=True + ) + aborted.outputs += _make_request_output( + "req-1", token_ids=[], finish_reason="abort", index=1 + ).outputs + yield aborted + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_serving_tokens(engine) + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=10, n=2), + model=MODEL_NAME, + stream=True, + stream_options=StreamOptions(include_usage=True), + ) + + response = await serving.serve_tokens(request) + parsed = _parse_sse_chunks([chunk async for chunk in response]) + choices = [c["choices"][0] for c in parsed[:-2]] + + assert [(c["index"], c["token_ids"], c["finish_reason"]) for c in choices] == [ + (0, [10], None), + (1, [11], None), + (0, [], "abort"), + (1, [], "abort"), + ] + assert parsed[-2]["usage"]["completion_tokens"] == 2 + + @pytest.mark.asyncio async def test_stream_error_mid_generation(): """finish_reason='error' mid-stream yields error chunk then [DONE].""" diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py b/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py index 2921422f01..2ee8f76e0f 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py @@ -710,21 +710,23 @@ async def test_text_mode_rejected_when_tokens_only(client): @pytest.mark.asyncio @pytest.mark.skipif( envs.VLLM_USE_RUST_FRONTEND, - reason="output_mode is not supported by the Rust frontend", + reason="the Rust frontend does not serve /inference/v1/abort_requests", ) -async def test_text_mode_stream_delivers_abort_finish_chunk( - client, tokenizer, messages +@pytest.mark.parametrize("output_mode", ["tokens", "text"]) +async def test_stream_delivers_abort_finish_chunk( + client, tokenizer, messages, output_mode ): - """The final output after an abort has no new token IDs, so it - reaches a text stream only because text mode emits finish-only chunks.""" + """The final output after an abort has no new token IDs, but the + stream must still deliver its finish reason.""" payload = { **_text_mode_payload( _chat_prompt_token_ids(tokenizer, messages), max_tokens=900, ignore_eos=True, ), + "output_mode": output_mode, "stream": True, - "request_id": "text-abort-e2e", + "request_id": f"{output_mode}-abort-e2e", } chunks = [] @@ -736,7 +738,7 @@ async def test_text_mode_stream_delivers_abort_finish_chunk( continue chunk = json.loads(line[len("data: ") :]) chunks.append(chunk) - if not aborted and chunk["choices"] and chunk["choices"][0]["text"]: + if not aborted and chunk["choices"] and chunk["choices"][0]["token_ids"]: abort = await client.post( ABORT_ENDPOINT, json={"request_ids": [chunk["request_id"]]}, @@ -745,7 +747,7 @@ async def test_text_mode_stream_delivers_abort_finish_chunk( aborted = True assert aborted - assert {chunk["output_mode"] for chunk in chunks} == {"text"} + assert {chunk["output_mode"] for chunk in chunks} == {output_mode} final_choice = chunks[-1]["choices"][0] assert final_choice["finish_reason"] == "abort" generated = sum(len(chunk["choices"][0]["token_ids"] or []) for chunk in chunks) diff --git a/vllm/entrypoints/scale_out/token_in_token_out/serving.py b/vllm/entrypoints/scale_out/token_in_token_out/serving.py index a30e519cb2..6b0195b104 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/serving.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/serving.py @@ -516,11 +516,12 @@ class ServingTokens(GenerateBaseServing): request_metadata: RequestResponseMetadata, ) -> AsyncGenerator[str, None]: num_prompt_tokens = 0 - num_generated_tokens: list[int] = [] first_iteration = True prompt_token_ids: list[int] | None = None num_cached_tokens = None sampling_params: SamplingParams = request.sampling_params + # With n > 1, the first output may not include every choice yet. + num_generated_tokens = [0] * sampling_params.n last_res: RequestOutput | None = None text_mode = request.output_mode == "text" tokenizer = self._logprobs_tokenizer(text_mode) @@ -540,7 +541,6 @@ class ServingTokens(GenerateBaseServing): if res.encoder_prompt_token_ids is not None: num_prompt_tokens += len(res.encoder_prompt_token_ids) num_cached_tokens = res.num_cached_tokens - num_generated_tokens = [0] * len(res.outputs) first_iteration = False for output in res.outputs: @@ -551,21 +551,15 @@ class ServingTokens(GenerateBaseServing): finish_reason = output.finish_reason self._raise_if_error(finish_reason, request_id) - if text_mode: - # Text held back for stop string matching is flushed - # with the final output, and the abort output has no - # new token IDs, so text or a finish reason is enough. - emit = bool( - delta_token_ids or output.text or finish_reason is not None - ) - else: - # Still emit a terminal empty chunk while prompt - # metadata is pending, so zero-token completions - # deliver it. - emit = bool(delta_token_ids) or ( - finish_reason is not None and prompt_token_ids is not None - ) - if not emit: + # Terminal outputs are always emitted so the client sees + # the finish reason, e.g. an abort, which has no new + # tokens. Text mode also emits text held back for stop + # string matching, which can arrive without token IDs. + if not ( + delta_token_ids + or finish_reason is not None + or (text_mode and output.text) + ): continue if sampling_params.logprobs is not None: From 6f75cc7d5d1b5adfb399cbf79a8231ae66844fcf Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Sun, 4 Oct 2026 11:01:04 +0800 Subject: [PATCH 07/22] [Bugfix][Responses API] Reuse streamed item ids in final harmony response (#59859) Signed-off-by: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Signed-off-by: sfeng33 <4florafeng@gmail.com> Co-authored-by: Flora Feng <4florafeng@gmail.com> --- .../responses/test_serving_responses.py | 83 ++++++++++++++++++- vllm/entrypoints/openai/responses/context.py | 5 ++ vllm/entrypoints/openai/responses/serving.py | 34 +++++--- .../openai/responses/streaming_events.py | 2 +- vllm/entrypoints/openai/responses/utils.py | 22 +++++ 5 files changed, 134 insertions(+), 12 deletions(-) diff --git a/tests/entrypoints/openai/responses/test_serving_responses.py b/tests/entrypoints/openai/responses/test_serving_responses.py index 591ba5c42a..acab4cb6e2 100644 --- a/tests/entrypoints/openai/responses/test_serving_responses.py +++ b/tests/entrypoints/openai/responses/test_serving_responses.py @@ -67,7 +67,7 @@ from vllm.entrypoints.serve.engine.protocol import ErrorResponse from vllm.inputs import tokens_input from vllm.logprobs import Logprob as SampleLogprob from vllm.outputs import CompletionOutput, RequestOutput -from vllm.parser.harmony import Segment +from vllm.parser.harmony import ChunkResult, HarmonyParser, Segment from vllm.renderers import TokenizeParams from vllm.renderers.online_renderer import ( OnlineRenderer, @@ -2129,6 +2129,87 @@ async def test_stream_completed_response_reuses_streamed_items(monkeypatch): assert function_call.call_id == "chatcmpl-tool-parser-id" +def _harmony_msg(channel: str, text: str, recipient: str | None = None): + msg = OpenAIHarmonyMessage.from_role_and_content(Role.ASSISTANT, text) + msg = msg.with_channel(channel) + return msg.with_recipient(recipient) if recipient else msg + + +def _harmony_segments(msg, *, streamed: bool = True) -> list[Segment]: + """Segments the parser yields for one message: a content delta (absent + for a zero-delta message), then the completed message.""" + delta = [Segment(msg.channel, msg.recipient, msg.content[0].text)] + completed = Segment(msg.channel, msg.recipient, "", completed_message=msg) + return (delta if streamed else []) + [completed] + + +async def _harmony_stream_events(segment_chunks: list[list[Segment]]): + serving = _make_serving_instance() + serving.use_harmony = True + parser = MagicMock(spec=HarmonyParser) + parser.process_chunk.side_effect = [ + ChunkResult(segments=segments, reasoning_token_count=0) + for segments in segment_chunks + ] + context = HarmonyContext([], [], frozenset({"get_weather"}), parser) + + async def result_generator(): + for token_id in range(len(segment_chunks)): + context.append_output(_make_request_output("", [token_id])) + yield context + + return [ + event + async for event in serving.responses_stream_generator( + request=ResponsesRequest(input="hi", stream=True, store=False), + sampling_params=SamplingParams(max_tokens=16), + result_generator=result_generator(), + context=context, + model_name="test-model", + tokenizer=MagicMock(), + request_metadata=RequestResponseMetadata(request_id="req"), + ) + ] + + +@pytest.mark.asyncio +async def test_harmony_stream_completed_response_reuses_streamed_ids(): + """response.completed must carry the id and call_id each item was streamed + with, not ones minted when the harmony messages are converted again.""" + messages = [ + _harmony_msg("analysis", "think"), + _harmony_msg("commentary", '{"city": "Paris"}', "functions.get_weather"), + _harmony_msg("final", "Sunny"), + ] + + events = await _harmony_stream_events([_harmony_segments(msg) for msg in messages]) + + streamed = [e.item for e in events if e.type == "response.output_item.done"] + output = events[-1].response.output + assert [item.type for item in output] == ["reasoning", "function_call", "message"] + assert [item.id for item in output] == [item.id for item in streamed] + assert output[1].call_id == streamed[1].call_id + assert output[2].content[0].text == "Sunny" + + +@pytest.mark.asyncio +async def test_harmony_stream_unstreamed_item_keeps_own_id(): + """An item with no done event (zero-delta) must not take the streamed id + of a later item of the same type.""" + unstreamed = _harmony_msg("analysis", "skipped") + reasoning = _harmony_msg("analysis", "think") + + events = await _harmony_stream_events( + [_harmony_segments(unstreamed, streamed=False), _harmony_segments(reasoning)] + ) + + (streamed,) = [e.item for e in events if e.type == "response.output_item.done"] + first, second = events[-1].response.output + assert first.content[0].text == "skipped" + assert first.id != streamed.id + assert second.id == streamed.id + + @pytest.mark.asyncio @pytest.mark.parametrize(("parallel_tool_calls", "num_calls"), [(False, 1), (None, 3)]) async def test_streaming_parallel_tool_calls( diff --git a/vllm/entrypoints/openai/responses/context.py b/vllm/entrypoints/openai/responses/context.py index e71af75a0e..a68df6dd4d 100644 --- a/vllm/entrypoints/openai/responses/context.py +++ b/vllm/entrypoints/openai/responses/context.py @@ -656,6 +656,11 @@ class HarmonyContext(ConversationContext): self.last_append_segments: list[Segment] = [] self.last_append_flush_status: bool = False + # Each completed message with the items sent for it in + # output_item.done events, so the final response can reuse their ids. + self.streamed_items_by_message: list[ + tuple[Message, list[ResponseOutputItem]] + ] = [] # Turn tracking - replaces multiple individual tracking variables self.current_turn_metrics = TurnMetrics() diff --git a/vllm/entrypoints/openai/responses/serving.py b/vllm/entrypoints/openai/responses/serving.py index bd876e4916..3a417ca7a5 100644 --- a/vllm/entrypoints/openai/responses/serving.py +++ b/vllm/entrypoints/openai/responses/serving.py @@ -12,6 +12,7 @@ from typing import Any, Final, cast from fastapi import Request from openai.types.responses import ( ResponseOutputItem, + ResponseOutputItemDoneEvent, ResponseOutputMessage, ResponseOutputText, ResponseStatus, @@ -68,6 +69,7 @@ from vllm.entrypoints.openai.responses.utils import ( build_response_output_items, extract_function_tool_names, extract_tool_types, + reuse_streamed_item_ids, ) from vllm.entrypoints.serve.engine.protocol import ErrorResponse from vllm.entrypoints.serve.utils.api_utils import get_max_tokens @@ -781,17 +783,21 @@ class OpenAIServingResponses(GenerateBaseServing): assert isinstance(context, HarmonyContext) output = [] harmony_msgs = context.messages[context.num_init_messages :] - if harmony_msgs: - fn_names = context.function_tool_names - for msg in harmony_msgs[:-1]: - output.extend(harmony_to_response_output(msg, fn_names)) - output.extend( - harmony_to_response_output( - harmony_msgs[-1], - fn_names, - incomplete=context.last_append_flush_status, - ) + # Streamed messages are a subsequence of harmony_msgs, in order. + streamed = iter(context.streamed_items_by_message) + next_streamed = next(streamed, None) + for i, msg in enumerate(harmony_msgs): + items = harmony_to_response_output( + msg, + context.function_tool_names, + incomplete=( + i == len(harmony_msgs) - 1 and context.last_append_flush_status + ), ) + if next_streamed is not None and next_streamed[0] is msg: + reuse_streamed_item_ids(items, next_streamed[1]) + next_streamed = next(streamed, None) + output.extend(items) if request.enable_response_messages: input_messages = context.messages[: context.num_init_messages] @@ -1291,16 +1297,24 @@ class OpenAIServingResponses(GenerateBaseServing): yield _increment_sequence_number_and_return(event) elif completed_message := segment.completed_message: + done_items: list[ResponseOutputItem] = [] # TODO: Fix browser emitted as MCP calls for event in emit_previous_item_done_events( completed_message, state, ctx.function_tool_names ): + if isinstance(event, ResponseOutputItemDoneEvent): + done_items.append(event.item) yield _increment_sequence_number_and_return(event) for event in emit_tool_action_events( completed_message, state, self.tool_server ): + if isinstance(event, ResponseOutputItemDoneEvent): + done_items.append(event.item) yield _increment_sequence_number_and_return(event) + ctx.streamed_items_by_message.append( + (completed_message, done_items) + ) state.reset_for_new_item() async def responses_stream_generator( diff --git a/vllm/entrypoints/openai/responses/streaming_events.py b/vllm/entrypoints/openai/responses/streaming_events.py index b4731dc687..2f5b151a1b 100644 --- a/vllm/entrypoints/openai/responses/streaming_events.py +++ b/vllm/entrypoints/openai/responses/streaming_events.py @@ -202,7 +202,7 @@ def emit_reasoning_delta_events( events: list[StreamingResponsesResponse] = [] if not state.sent_output_item_added: state.sent_output_item_added = True - state.current_item_id = f"msg_{random_uuid()}" + state.current_item_id = f"rs_{random_uuid()}" events.append( ResponseOutputItemAddedEvent( type="response.output_item.added", diff --git a/vllm/entrypoints/openai/responses/utils.py b/vllm/entrypoints/openai/responses/utils.py index 3caff4b056..6c8507840f 100644 --- a/vllm/entrypoints/openai/responses/utils.py +++ b/vllm/entrypoints/openai/responses/utils.py @@ -112,6 +112,28 @@ def build_response_output_items( return outputs +def reuse_streamed_item_ids( + items: list[ResponseOutputItem], + streamed_items: list[ResponseOutputItem], +) -> None: + """Give rebuilt output items the ids already streamed for the same message. + + Items are paired by type in order; an item without a streamed counterpart + keeps its own id. + """ + remaining = list(streamed_items) + for item in items: + for i, streamed in enumerate(remaining): + if streamed.type != item.type: + continue + item.id = streamed.id + if isinstance(item, ResponseFunctionToolCall): + assert isinstance(streamed, ResponseFunctionToolCall) + item.call_id = streamed.call_id + del remaining[i] + break + + def should_continue_final_message( request_input: str | list[ResponseInputOutputItem], ) -> bool: From 18f8f96025b556071eb627076f94df560fbd3a22 Mon Sep 17 00:00:00 2001 From: Zupeng Wang Date: Sun, 4 Oct 2026 12:30:17 +0800 Subject: [PATCH 08/22] [Bugfix][KV Offload] Release pending CPU lookup pins on cache reset (#59862) Signed-off-by: Wang Zupeng --- tests/v1/simple_kv_offload/test_scheduler.py | 25 ++++++++++++++++++++ vllm/v1/simple_kv_offload/manager.py | 4 ++++ 2 files changed, 29 insertions(+) diff --git a/tests/v1/simple_kv_offload/test_scheduler.py b/tests/v1/simple_kv_offload/test_scheduler.py index 1c0cd74565..0f2a98e4a6 100644 --- a/tests/v1/simple_kv_offload/test_scheduler.py +++ b/tests/v1/simple_kv_offload/test_scheduler.py @@ -1864,6 +1864,31 @@ def test_reset_pending_lazy_stores() -> None: assert hit_tokens == 0, "CPU cache should be empty after reset" +@pytest.mark.parametrize("num_groups", [1, 2]) +@pytest.mark.parametrize("lazy", [False, True]) +def test_reset_releases_pending_cpu_hits(num_groups: int, lazy: bool) -> None: + """A lookup awaiting GPU allocation must not prevent a cache reset.""" + fix = make_scheduler(num_groups=num_groups, lazy=lazy) + sched = fix.scheduler + req = make_request(num_blocks=2) + cpu_blocks = [] + for group_id in range(num_groups): + cpu_blocks.extend( + _allocate_gpu_blocks(sched.cpu_block_pool, req, 2, group_id=group_id) + ) + sched.cpu_block_pool.free_blocks(cpu_blocks) + + assert sched.get_num_new_matched_tokens(req, 0) == (2 * BLOCK_SIZE, True) + assert all(block.ref_cnt == 1 for block in cpu_blocks) + assert not sched._reqs_to_load + + assert sched.reset() + assert all(block.ref_cnt == 0 for block in cpu_blocks) + assert not sched._pending_cpu_hits + assert sched.get_num_new_matched_tokens(req, 0) == (0, False) + assert sched.reset() + + # --------------------------------------------------------------------------- # Test 14: Reset with pending loads waits for completion # --------------------------------------------------------------------------- diff --git a/vllm/v1/simple_kv_offload/manager.py b/vllm/v1/simple_kv_offload/manager.py index c91537375f..3d57e73566 100644 --- a/vllm/v1/simple_kv_offload/manager.py +++ b/vllm/v1/simple_kv_offload/manager.py @@ -1395,6 +1395,10 @@ class SimpleCPUOffloadScheduler: the transfer finished, then release refs without caching abandoned store results. """ + for pending in self._pending_cpu_hits.values(): + self._free_pending_cpu_hit(pending) + self._pending_cpu_hits.clear() + self._abandoned_store_event_to_blocks.update(self._store_event_to_blocks) for transfer in self._pending_finished_stores: self._release_transfer_refs(transfer) From 0872ddf4b96d57a0745893478deef595b1994f2e Mon Sep 17 00:00:00 2001 From: Ylean <61880605+Woolgathererer@users.noreply.github.com> Date: Sun, 4 Oct 2026 14:11:41 +0800 Subject: [PATCH 09/22] [Bugfix][Multimodal] Normalize single-channel audio to 1D (#56691) Signed-off-by: Ylean <61880605+Woolgathererer@users.noreply.github.com> Co-authored-by: Codex Co-authored-by: Cyrus Leung --- tests/multimodal/test_audio.py | 32 +++++++++++++++++++------------- vllm/multimodal/audio.py | 2 +- 2 files changed, 20 insertions(+), 14 deletions(-) diff --git a/tests/multimodal/test_audio.py b/tests/multimodal/test_audio.py index 24d925aa02..c15ef9adb7 100644 --- a/tests/multimodal/test_audio.py +++ b/tests/multimodal/test_audio.py @@ -216,17 +216,19 @@ class TestNormalizeAudio: assert result.ndim == 1 torch.testing.assert_close(result, torch.tensor([0.0, 1.0])) - def test_mono_passthrough_for_1d_numpy(self): - """1D numpy array should pass through unchanged with mono spec.""" - mono = np.array([1.0, 2.0, 3.0], dtype=np.float32) - result = normalize_audio(mono, MONO_AUDIO_SPEC) + @pytest.mark.parametrize("shape", [(3,), (1, 3), (3, 1), (1, 1)]) + def test_mono_passthrough_for_numpy(self, shape): + """Mono output is 1D, including audio with an explicit channel axis.""" + mono = np.arange(max(shape), dtype=np.float32) + result = normalize_audio(mono.reshape(shape), MONO_AUDIO_SPEC) assert result.ndim == 1 np.testing.assert_array_equal(result, mono) - def test_mono_passthrough_for_1d_torch(self): - """1D torch tensor should pass through unchanged with mono spec.""" - mono = torch.tensor([1.0, 2.0, 3.0]) - result = normalize_audio(mono, MONO_AUDIO_SPEC) + @pytest.mark.parametrize("shape", [(3,), (1, 3), (3, 1), (1, 1)]) + def test_mono_passthrough_for_torch(self, shape): + """Mono output is 1D, including audio with an explicit channel axis.""" + mono = torch.arange(max(shape), dtype=torch.float32) + result = normalize_audio(mono.reshape(shape), MONO_AUDIO_SPEC) assert result.ndim == 1 torch.testing.assert_close(result, mono) @@ -410,8 +412,9 @@ class TestMultiModalDataParserChannelNormalization: # When target_channels=None, stereo audio should be preserved assert audio_item.ndim == 2, f"Expected 2D stereo audio, got {audio_item.ndim}D" - def test_parser_mono_passthrough_when_target_channels_1(self): - """Parser should pass through mono audio unchanged when target_channels=1.""" + @pytest.mark.parametrize("shape", [(16000,), (1, 16000), (16000, 1)]) + def test_parser_mono_passthrough_when_target_channels_1(self, shape): + """Mono layouts must preserve samples and report the waveform length.""" from vllm.multimodal.parse import MultiModalDataParser # Create parser with mono normalization enabled @@ -420,17 +423,20 @@ class TestMultiModalDataParserChannelNormalization: target_channels=1, ) - # Create mono audio (already 1D) - mono_audio = np.random.randn(16000).astype(np.float32) + mono_audio = np.linspace(-1.0, 1.0, 16000, dtype=np.float32) # Parse audio data - result = parser._parse_audio_data((mono_audio, 16000)) + result = parser.parse_mm_data({"audio": (mono_audio.reshape(shape), 16000)})[ + "audio" + ] # Check that result is still mono (1D) assert result is not None audio_item = result.get(0) assert audio_item.ndim == 1 assert audio_item.shape == (16000,) + assert len(audio_item) == 16000 + np.testing.assert_array_equal(audio_item, mono_audio) def test_parser_with_target_channels_2(self): """Parser should reduce 6-channel to 2-channel when target_channels=2.""" diff --git a/vllm/multimodal/audio.py b/vllm/multimodal/audio.py index be026810bf..991909ccbc 100644 --- a/vllm/multimodal/audio.py +++ b/vllm/multimodal/audio.py @@ -146,7 +146,7 @@ def normalize_audio( # No reduction needed if already at target if num_channels == spec.target_channels: - return audio + return audio[0] if num_channels == 1 else audio # Cannot expand channels if num_channels < spec.target_channels: From 155d23cb008f793d4609d2bd5885cd9882ba38ee Mon Sep 17 00:00:00 2001 From: Yongqi Wang Date: Sun, 4 Oct 2026 18:31:15 +0800 Subject: [PATCH 10/22] [Bugfix][XPU] Make DeepSeek V4 FP8 sparse decode graph-capturable (#59159) Signed-off-by: Yongqi Wang Co-authored-by: Kunshang Ji --- .../deepseek_v4/xpu/xpu_sparse_decode_fp8.py | 125 +++++------------- 1 file changed, 36 insertions(+), 89 deletions(-) diff --git a/vllm/models/deepseek_v4/xpu/xpu_sparse_decode_fp8.py b/vllm/models/deepseek_v4/xpu/xpu_sparse_decode_fp8.py index 7c0b808de2..4fb8c4059b 100644 --- a/vllm/models/deepseek_v4/xpu/xpu_sparse_decode_fp8.py +++ b/vllm/models/deepseek_v4/xpu/xpu_sparse_decode_fp8.py @@ -166,8 +166,7 @@ def xpu_sparse_decode_fp8( num_tokens = q.shape[0] device = q.device - # Determine max topk and swa widths - if not swa_only and topk_indices is not None: + if not swa_only and topk_indices is not None and kv_cache is not None: topk_idx_2d = ( topk_indices.squeeze(1) if topk_indices.dim() == 3 else topk_indices ) @@ -176,109 +175,57 @@ def xpu_sparse_decode_fp8( topk_idx_2d = None max_topk = 0 + use_topk = topk_idx_2d is not None swa_idx_2d = swa_indices.squeeze(1) if swa_indices.dim() == 3 else swa_indices max_swa = swa_idx_2d.shape[1] - K_total = max_topk + max_swa - # Allocate flat workspace: [num_tokens * K_total, 512] bf16 + # Topk rows precede SWA rows in the flat workspace. + n_topk_rows = num_tokens * max_topk workspace = torch.empty( (num_tokens * K_total, OUTPUT_DIM), dtype=torch.bfloat16, device=device ) - ws_3d = workspace.view(num_tokens, K_total, OUTPUT_DIM) - # Dequant+gather topk slots from compressed cache - if not swa_only and topk_idx_2d is not None and kv_cache is not None: - topk_flat = topk_idx_2d.reshape(-1).to(torch.int32) - topk_buf = torch.empty( - (num_tokens * max_topk, OUTPUT_DIM), dtype=torch.bfloat16, device=device + if topk_idx_2d is not None and kv_cache is not None and n_topk_rows > 0: + dequant_gather_slots( + workspace[:n_topk_rows], + kv_cache, + topk_idx_2d.reshape(-1).to(torch.int32), + kv_cache.shape[1], ) - compressed_block_size = kv_cache.shape[1] - dequant_gather_slots(topk_buf, kv_cache, topk_flat, compressed_block_size) - ws_3d[:, :max_topk, :] = topk_buf.view(num_tokens, max_topk, OUTPUT_DIM) - # Dequant+gather SWA slots - swa_flat = swa_idx_2d.reshape(-1).to(torch.int32) - swa_buf = torch.empty( - (num_tokens * max_swa, OUTPUT_DIM), dtype=torch.bfloat16, device=device + dequant_gather_slots( + workspace[n_topk_rows:], + swa_kv_cache, + swa_idx_2d.reshape(-1).to(torch.int32), + swa_kv_cache.shape[1], ) - swa_block_size = swa_kv_cache.shape[1] - dequant_gather_slots(swa_buf, swa_kv_cache, swa_flat, swa_block_size) - ws_3d[:, max_topk:, :] = swa_buf.view(num_tokens, max_swa, OUTPUT_DIM) - - # Build combined indices into the flat workspace and combined lengths. - # Workspace layout per token t: [topk_0..topk_{max_topk-1}, swa_0..swa_{max_swa-1}] - # Flat index for token t, position p = t * K_total + p - # - # IMPORTANT: The attention kernel uses combined_lens as a position cutoff — - # it only reads indices[0..combined_lens-1]. So indices must be PACKED - # contiguously: [valid_topk_indices..., valid_swa_indices..., -1 padding...] - if not swa_only and topk_lens is not None: - combined_lens = (topk_lens + swa_lens).to(torch.int32) - else: - combined_lens = swa_lens.to(torch.int32) - - max_combined = int(combined_lens.max().item()) if combined_lens.numel() > 0 else 0 - # Round up to BLOCK_N=16 alignment for kernel efficiency + # Use a shape-derived width so replay accepts changing lens without host sync. _BLOCK_N = 16 - max_combined_padded = ((max_combined + _BLOCK_N - 1) // _BLOCK_N) * _BLOCK_N + width = max(_BLOCK_N, ((K_total + _BLOCK_N - 1) // _BLOCK_N) * _BLOCK_N) - # Build packed index table: [num_tokens, max_combined_padded] - # Each token t: [topk_0..topk_{tlen-1}, swa_0..swa_{slen-1}, -1 padding] - # Vectorized: for each token, topk indices are t*K_total + 0..tlen-1, - # swa indices are t*K_total + max_topk + 0..slen-1 - combined_indices = torch.full( - (num_tokens, max_combined_padded), - fill_value=-1, - dtype=torch.int32, - device=device, + pos = torch.arange(width, device=device, dtype=torch.int32).unsqueeze(0) + tok = torch.arange(num_tokens, device=device, dtype=torch.int32).unsqueeze(1) + + slen = swa_lens.reshape(-1, 1).to(torch.int32).clamp(min=0, max=max_swa) + if use_topk and topk_lens is not None: + tlen = topk_lens.reshape(-1, 1).to(torch.int32).clamp(min=0, max=max_topk) + else: + tlen = torch.zeros_like(slen) + + in_topk = pos < tlen + in_swa = (pos >= tlen) & (pos < tlen + slen) + + topk_ws = tok * max_topk + pos + swa_ws = n_topk_rows + tok * max_swa + (pos - tlen) + + combined_indices = torch.where( + in_topk, + topk_ws, + torch.where(in_swa, swa_ws, torch.full_like(topk_ws, -1)), ) - token_offsets = ( - torch.arange(num_tokens, device=device, dtype=torch.int32) * K_total - ) # [B] - - if not swa_only and topk_lens is not None: - # Pack topk: for each token, write t*K_total + 0..tlen-1 at positions 0..tlen-1 - max_tlen = int(topk_lens.max().item()) - topk_range = torch.arange(max_tlen, device=device, dtype=torch.int32).unsqueeze( - 0 - ) - topk_valid = topk_range < topk_lens.unsqueeze(1) - topk_ws_indices = token_offsets.unsqueeze(1) + topk_range - combined_indices[:, :max_tlen] = torch.where( - topk_valid, - topk_ws_indices, - torch.tensor(-1, dtype=torch.int32, device=device), - ) - # Pack swa after topk: positions tlen..tlen+slen-1 - # Since tlen varies per token, we need per-token offset - swa_range = torch.arange(max_swa, device=device, dtype=torch.int32).unsqueeze(0) - swa_valid = swa_range < swa_lens.unsqueeze(1) - swa_ws_indices = token_offsets.unsqueeze(1) + max_topk + swa_range - # Write at position topk_lens[t] + swa_pos for each token - for t_idx in range(num_tokens): - tlen = int(topk_lens[t_idx].item()) - slen = int(swa_lens[t_idx].item()) - combined_indices[t_idx, tlen : tlen + slen] = swa_ws_indices[t_idx, :slen] - else: - # SWA-only: pack swa indices at positions 0..slen-1 - # Use min(max_swa, max_combined_padded) because combined_indices only - # has max_combined_padded columns, and all valid entries fit within it. - effective_swa = min(max_swa, max_combined_padded) - swa_range = torch.arange( - effective_swa, device=device, dtype=torch.int32 - ).unsqueeze(0) - swa_valid = swa_range < swa_lens.unsqueeze(1) - swa_ws_indices = token_offsets.unsqueeze(1) + swa_range # max_topk=0 - combined_indices[:, :effective_swa] = torch.where( - swa_valid, - swa_ws_indices, - torch.tensor(-1, dtype=torch.int32, device=device), - ) - - # Call BF16 sparse MLA kernel out_attn, _, _ = triton_bf16_mla_sparse_interface( q=q, kv=workspace.unsqueeze(1), From d61081dc3d3f1740a5d8bf82608b62974393c2de Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Sun, 4 Oct 2026 11:53:07 +0100 Subject: [PATCH 11/22] [CI] Bump peft to satisfy transformers 5.18 minimum (#59932) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 --- requirements/test/cpu.txt | 2 +- requirements/test/cuda.in | 2 +- requirements/test/cuda.txt | 2 +- requirements/test/rocm.in | 2 +- requirements/test/rocm.txt | 2 +- 5 files changed, 5 insertions(+), 5 deletions(-) diff --git a/requirements/test/cpu.txt b/requirements/test/cpu.txt index 91fac06966..45725d903b 100644 --- a/requirements/test/cpu.txt +++ b/requirements/test/cpu.txt @@ -713,7 +713,7 @@ pathvalidate==3.2.1 # via pytablewriter patsy==1.0.1 # via statsmodels -peft==0.19.1 +peft==0.21.1 # via -r requirements/test/cuda.in perceptron==0.1.4 # via -r requirements/test/cuda.in diff --git a/requirements/test/cuda.in b/requirements/test/cuda.in index 4732f3a505..b40d80413e 100644 --- a/requirements/test/cuda.in +++ b/requirements/test/cuda.in @@ -20,7 +20,7 @@ httpx librosa # required for audio tests vector_quantize_pytorch # required for minicpmo_26 test vocos # required for minicpmo_26 test -peft>=0.19.1 # required for phi-4-mm test +peft>=0.20.0 # required for phi-4-mm test pqdm ray[cgraph,default]>=2.56.0,<2.57 # Ray Compiled Graph, required by pipeline parallelism tests; sharded-RDT weight transfer needs Ray Direct Transport (>=2.56) sentence-transformers>=5.2.0 # required for embedding tests diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt index ad185844be..1f64120e31 100644 --- a/requirements/test/cuda.txt +++ b/requirements/test/cuda.txt @@ -801,7 +801,7 @@ pathvalidate==3.2.1 # via pytablewriter patsy==1.0.1 # via statsmodels -peft==0.19.1 +peft==0.21.1 # via -r requirements/test/cuda.in perceptron==0.1.4 # via -r requirements/test/cuda.in diff --git a/requirements/test/rocm.in b/requirements/test/rocm.in index 753fad0810..7c0e4848e9 100644 --- a/requirements/test/rocm.in +++ b/requirements/test/rocm.in @@ -19,7 +19,7 @@ httpx librosa # required for audio tests vector_quantize_pytorch # required for minicpmo_26 test vocos # required for minicpmo_26 test -peft>=0.19.1 # required for phi-4-mm test +peft>=0.20.0 # required for phi-4-mm test pqdm ray[cgraph,default]==2.55.1 # Includes worker startup getenv/setenv race fix sentence-transformers>=5.2.0 # required for embedding tests diff --git a/requirements/test/rocm.txt b/requirements/test/rocm.txt index e6c137ebf5..b2919a6e83 100644 --- a/requirements/test/rocm.txt +++ b/requirements/test/rocm.txt @@ -773,7 +773,7 @@ pathvalidate==3.3.1 # via pytablewriter patsy==1.0.2 # via statsmodels -peft==0.19.1 +peft==0.21.1 # via -r requirements/test/rocm.in perceptron==0.1.4 # via -r requirements/test/rocm.in From d64f6cd08a4eb390664ed5230e6b6ae48b73bfae Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Sun, 4 Oct 2026 08:26:10 -0500 Subject: [PATCH 12/22] [CI] Auto-label pooling PRs and issues (#59913) Signed-off-by: Taneem Ibrahim Co-authored-by: Lucas Wilkinson --- .github/mergify.yml | 26 ++++++++++++++++++++++++++ .github/workflows/issue_autolabel.yml | 21 +++++++++++++++++++++ 2 files changed, 47 insertions(+) diff --git a/.github/mergify.yml b/.github/mergify.yml index 73ea58aed6..0798654d34 100644 --- a/.github/mergify.yml +++ b/.github/mergify.yml @@ -112,6 +112,32 @@ pull_request_rules: add: - frontend +- name: label-pooling + description: Automatically apply pooling label + conditions: + - -closed + - or: + - files~=^vllm/entrypoints/pooling/ + - files~=^vllm/model_executor/layers/pooler/ + - files=vllm/model_executor/models/transformers/pooling.py + - files=vllm/model_executor/models/adapters.py + - files=vllm/config/pooler.py + - files=vllm/pooling_params.py + - files~=^vllm/v1/pool/ + - files~=^vllm/v1/worker/gpu/pool/ + - files~=^tests/entrypoints/pooling/ + - files~=^tests/models/(language|multimodal)/pooling(_mteb_test)?/ + - files~=^tests/model_executor/layers/test_pooler_.*\.py$ + - files=tests/test_pooling_params.py + - files~=^tests/v1/e2e/general/test_pooling_.*\.py$ + - files~=^docs/models/pooling_models/ + - files~=^examples/pooling/ + - title~=(?i)\bpooling\b + actions: + label: + add: + - pooling + - name: label-rust description: Automatically apply rust label conditions: diff --git a/.github/workflows/issue_autolabel.yml b/.github/workflows/issue_autolabel.yml index fc6e4921ef..c3190f8f18 100644 --- a/.github/workflows/issue_autolabel.yml +++ b/.github/workflows/issue_autolabel.yml @@ -255,6 +255,27 @@ jobs: // phrases go in `substrings`: they are specific enough on their // own, and a boundary match would miss ordinary variations -- // "structured outputs", "tool parsers", "prefix caching". + pooling: { + keywords: [ + { term: "pooling", searchIn: "title" }, + { term: "pooler", searchIn: "title" }, + { term: "rerank", searchIn: "title" }, + { term: "reranker", searchIn: "title" }, + { term: "reranking", searchIn: "title" }, + ], + substrings: [ + { term: "embedding model", searchIn: "title" }, + { term: "embedding API", searchIn: "title" }, + { term: "embeddings API", searchIn: "title" }, + { term: "sentence embedding", searchIn: "title" }, + { term: "sequence classification", searchIn: "title" }, + { term: "token classification", searchIn: "title" }, + { term: "cross-encoder", searchIn: "title" }, + { term: "late interaction", searchIn: "title" }, + { term: "late-interaction", searchIn: "title" }, + { term: "reward model", searchIn: "title" }, + ], + }, "speculative-decoding": { keywords: [ { term: "eagle", searchIn: "title" }, From f98d1fc49361cb50457f6109c3c01a5c77e415b4 Mon Sep 17 00:00:00 2001 From: 00CC <83797617+tangzzycc@users.noreply.github.com> Date: Sun, 4 Oct 2026 21:48:46 +0800 Subject: [PATCH 13/22] [ROCm][Bugfix] Fix ROCM_ATTN sliding-window boundary (#59550) Signed-off-by: tangzzycc <3081129260@qq.com> --- .../kernels/attention/test_prefix_prefill.py | 58 +++++++++++++++++++ vllm/v1/attention/backends/rocm_attn.py | 4 +- 2 files changed, 61 insertions(+), 1 deletion(-) diff --git a/tests/kernels/attention/test_prefix_prefill.py b/tests/kernels/attention/test_prefix_prefill.py index 290d645e07..b2efa76da8 100644 --- a/tests/kernels/attention/test_prefix_prefill.py +++ b/tests/kernels/attention/test_prefix_prefill.py @@ -5,6 +5,7 @@ import math import random import time from collections.abc import Callable +from types import SimpleNamespace import pytest import torch @@ -323,6 +324,63 @@ def test_contexted_kv_attention( torch.testing.assert_close(output, output_ref, atol=atol, rtol=0) +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm backend") +@pytest.mark.parametrize("sliding_window", [None, 1, 4]) +@pytest.mark.parametrize("query_len", [1, 3]) +@torch.inference_mode() +def test_rocm_backend_sliding_window_includes_boundary_token( + sliding_window: int | None, query_len: int +) -> None: + """A window of W must include the current token and W - 1 preceding tokens.""" + from vllm.v1.attention.backends.rocm_attn import RocmAttentionImpl + from vllm.v1.attention.ops.paged_attn import PagedAttention + + device, dtype = "cuda:0", torch.float16 + num_heads, head_size, seq_len = 4, 128, 8 + query = torch.zeros(query_len, num_heads, head_size, dtype=dtype, device=device) + key = torch.zeros(seq_len, 1, head_size, dtype=dtype, device=device) + values = torch.arange(1, seq_len + 1, dtype=dtype, device=device) + value = values[:, None, None].expand_as(key) + kv_cache = torch.zeros(2, 1, 16, head_size, dtype=dtype, device=device) + _, value_cache = PagedAttention.split_kv_cache(kv_cache, 1, head_size) + value_cache[..., :seq_len] = value.permute(1, 2, 0) + scale = torch.tensor(1.0, device=device) + metadata = SimpleNamespace( + use_cascade=False, + num_actual_tokens=query_len, + query_start_loc=torch.tensor([0, query_len], dtype=torch.int32, device=device), + seq_lens=torch.tensor([seq_len], dtype=torch.int32, device=device), + max_query_len=query_len, + max_seq_len=seq_len, + block_table=torch.zeros(1, 1, dtype=torch.int32, device=device), + causal=True, + ) + impl = RocmAttentionImpl( + num_heads=num_heads, + head_size=head_size, + scale=head_size**-0.5, + num_kv_heads=1, + alibi_slopes=None, + sliding_window=sliding_window, + kv_cache_dtype="auto", + ) + output = torch.empty_like(query) + impl.forward( + SimpleNamespace(_k_scale=scale, _v_scale=scale), + query, + key[-query_len:], + value[-query_len:], + kv_cache.transpose(0, 1), + metadata, + output, + ) + mask = create_causal_attention_mask_for_sdpa( + [query_len], [seq_len], sliding_window or 0, device=device, dtype=dtype + ) + expected = (mask.float().softmax(dim=-1) @ values.float())[:, None, None] + torch.testing.assert_close(output, expected.to(dtype).expand_as(output)) + + @pytest.mark.parametrize("num_heads", NUM_HEADS) @pytest.mark.parametrize("num_queries_per_kv", NUM_QUERIES_PER_KV) @pytest.mark.parametrize("head_size", HEAD_SIZES) diff --git a/vllm/v1/attention/backends/rocm_attn.py b/vllm/v1/attention/backends/rocm_attn.py index 453f49b468..3ccc1e5d57 100644 --- a/vllm/v1/attention/backends/rocm_attn.py +++ b/vllm/v1/attention/backends/rocm_attn.py @@ -481,7 +481,9 @@ class RocmAttentionImpl(AttentionImpl): k_scale=layer._k_scale, v_scale=layer._v_scale, alibi_slopes=self.alibi_slopes, - sliding_window=self.sliding_window[0], + # self.sliding_window[0] is the FlashAttention-style left span (W - 1). + # chunked_prefill_paged_decode expects the full window length W. + sliding_window=1 + self.sliding_window[0], sm_scale=self.scale, output_scale=output_scale, sinks=self.sinks, From bdd31c31defd0e47a2779eb46ac1a6d88198a41b Mon Sep 17 00:00:00 2001 From: Aarushi Jain <142941703+aarushjain29@users.noreply.github.com> Date: Sun, 4 Oct 2026 09:14:27 -0500 Subject: [PATCH 14/22] [CI] Stop requiring both API servers to record weight sync metrics (#59521) Signed-off-by: aarushjain29 Co-authored-by: Andreas Karatzas --- .../dev/rlhf/test_weight_operation_metrics.py | 29 ++----------------- 1 file changed, 3 insertions(+), 26 deletions(-) diff --git a/tests/entrypoints/serve/dev/rlhf/test_weight_operation_metrics.py b/tests/entrypoints/serve/dev/rlhf/test_weight_operation_metrics.py index 063b307767..d7758011f9 100644 --- a/tests/entrypoints/serve/dev/rlhf/test_weight_operation_metrics.py +++ b/tests/entrypoints/serve/dev/rlhf/test_weight_operation_metrics.py @@ -2,11 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Weight-transfer metrics scraped from a real two-API-server deployment.""" -from pathlib import Path - import requests import torch -from prometheus_client.multiprocess import MultiProcessCollector from prometheus_client.parser import text_string_to_metric_families from transformers import AutoModelForCausalLM @@ -21,10 +18,7 @@ from vllm.distributed.weight_transfer.ipc_engine import IPCTrainerInitInfo DURATION = "vllm:rl_weight_update_operation_duration_seconds" IN_FLIGHT = "vllm:rl_weight_update_operations_in_flight" -# The API servers accept from one shared listening socket, so the kernel may -# hand one process most connections. Sync until both have recorded, within a cap. -MIN_SYNCS = 8 -MAX_SYNCS = 64 +SYNCS = 8 def scrape(url: str) -> dict[tuple[str, str], float]: @@ -37,17 +31,6 @@ def scrape(url: str) -> dict[tuple[str, str], float]: } -def pids_that_recorded(multiproc_dir: Path) -> set[str]: - return { - path.stem.rsplit("_", 1)[1] - for path in multiproc_dir.glob("histogram_*.db") - if any( - metric.name == DURATION - for metric in MultiProcessCollector.merge([str(path)]) - ) - } - - def test_weight_sync_metrics_aggregate_across_api_servers(tmp_path): args = [ "--load-format", @@ -77,17 +60,11 @@ def test_weight_sync_metrics_aggregate_across_api_servers(tmp_path): client=HTTPVLLMWeightSyncClient(url), source=ModuleSource(trainer), ) - syncs = 0 - while syncs < MIN_SYNCS or ( - len(pids_that_recorded(tmp_path)) < 2 and syncs < MAX_SYNCS - ): + for _ in range(SYNCS): engine.send_weights() - syncs += 1 metrics = scrape(url) - calls = {"init": 1, "start": syncs, "update": syncs, "finish": syncs} + calls = {"init": 1, "start": SYNCS, "update": SYNCS, "finish": SYNCS} for operation, count in calls.items(): assert metrics[(f"{DURATION}_count", operation)] == count, operation assert metrics[(IN_FLIGHT, operation)] == 0, operation - # The counts above are sums over processes only if both servers recorded. - assert len(pids_that_recorded(tmp_path)) == 2 From bd42276e1e25f540f019381c65994120158aa2a1 Mon Sep 17 00:00:00 2001 From: Divy Date: Sun, 4 Oct 2026 08:13:20 -0700 Subject: [PATCH 15/22] [Misc][Platform] Check aligned KV block sizes against every attention backend (#58457) Signed-off-by: Divy Co-authored-by: Claude Opus 5.5 (1M context) Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> --- tests/v1/worker/test_gpu_model_runner.py | 61 ++++++++++++++++++++++++ vllm/platforms/cpu.py | 2 + vllm/platforms/interface.py | 31 ++++++++++++ vllm/platforms/xpu.py | 5 ++ 4 files changed, 99 insertions(+) diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index fb4b7d1e58..5cf53899b8 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -486,6 +486,67 @@ def test_preferred_block_size_rejects_backends_with_no_common_size(): Platform._preferred_block_size_for_backends(classes, 16, None) +@pytest.mark.parametrize("cpu", [False, True]) +def test_alignment_rejected_by_a_sibling_backend_raises(monkeypatch, cpu): + # Alignment that lands on a size an exact-size sibling rejects must fail + # here, not later in select_common_block_size(). + backends = [_mock_backend([MultipleOf(16)]), _mock_backend([16], exact=True)] + monkeypatch.setattr( + Platform, "_find_non_ssm_backends", classmethod(lambda cls, c: backends) + ) + + def bump(cls, vllm_config, backend_cls): + vllm_config.cache_config.block_size = 32 + + monkeypatch.setattr(Platform, "_align_hybrid_block_size", classmethod(bump)) + cache_config = SimpleNamespace( + block_size=16, user_specified_block_size=True, kv_cache_dtype_skip_layers=None + ) + vllm_config = SimpleNamespace( + cache_config=cache_config, model_config=SimpleNamespace(is_hybrid=True) + ) + from vllm.platforms.cpu import CpuPlatform # overrides the method; must check too + + platform = CpuPlatform if cpu else Platform + with pytest.raises(ValueError, match="MOCK_EXACT"): + platform.update_block_size_for_backend(vllm_config) + + +def test_xpu_gdn_rounding_rejected_by_a_sibling_backend_raises(monkeypatch): + # XPU's GDN rounding runs after super()'s check, so it must check again. + try: + from vllm.platforms.xpu import XPUPlatform + except ImportError: + pytest.skip("vllm_xpu_kernels not importable outside an XPU stack") + + backends = [_mock_backend([MultipleOf(16)]), _mock_backend([16], exact=True)] + monkeypatch.setattr( + Platform, "update_block_size_for_backend", classmethod(lambda cls, c: None) + ) + monkeypatch.setattr( + Platform, "_find_non_ssm_backends", classmethod(lambda cls, c: backends) + ) + gdn_layer = SimpleNamespace( + get_attn_backend=lambda: SimpleNamespace(get_name=lambda: "GDN_ATTN") + ) + monkeypatch.setattr( + "vllm.config.vllm.get_layers_from_vllm_config", + lambda vllm_config, layer_type: {"gdn": gdn_layer}, + ) + cache_config = SimpleNamespace( + block_size=16, + mamba_cache_mode="none", + mamba_page_size_padded=None, + user_specified_block_size=True, + kv_cache_dtype_skip_layers=None, + ) + vllm_config = SimpleNamespace( + cache_config=cache_config, model_config=SimpleNamespace(is_hybrid=True) + ) + with pytest.raises(ValueError, match="MOCK_EXACT"): + XPUPlatform.update_block_size_for_backend(vllm_config) + + def test_set_active_mm_loras_builds_tower_and_connector_mappings(): model = Mock() model.get_mm_lora_token_counts.side_effect = ( diff --git a/vllm/platforms/cpu.py b/vllm/platforms/cpu.py index c9d57fd5d3..25d28daf93 100644 --- a/vllm/platforms/cpu.py +++ b/vllm/platforms/cpu.py @@ -523,7 +523,9 @@ class CpuPlatform(Platform): if not backend_classes: return + pre_block_size = vllm_config.cache_config.block_size cls._align_hybrid_block_size(vllm_config, backend_classes[0]) + cls._check_aligned_block_size(vllm_config, backend_classes, pre_block_size) @classmethod def discover_numa_topology(cls) -> list[list[int]]: diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index 769ffe5e8e..dc7178cf84 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -698,6 +698,7 @@ class Platform: # Phase 2: Align block/mamba sizes for hybrid models # (may override user settings). + pre_block_size = cache_config.block_size if model_config.is_hybrid: cls._align_hybrid_block_size(vllm_config, backend_classes[0]) @@ -706,6 +707,36 @@ class Platform: # May override the user's --block-size. if cache_config.kv_cache_dtype_skip_layers: cls._align_heterogeneous_kv_block_size(vllm_config, backend_classes[0]) + cls._check_aligned_block_size(vllm_config, backend_classes, pre_block_size) + + @classmethod + def _check_aligned_block_size( + cls, + vllm_config: "VllmConfig", + backend_classes: "list[type[AttentionBackend]]", + pre_block_size: int, + ) -> None: + """Fail now, not at worker setup, if alignment left a rejected size. + + Alignment rounds by the first backend's kernel block only, so it can + land on a size a sibling backend rejects. + """ + from vllm.config.vllm import set_current_vllm_config + + block_size = vllm_config.cache_config.block_size + if block_size == pre_block_size: + return + with set_current_vllm_config(vllm_config): + rejecting = [ + b.get_name() + for b in backend_classes + if not b.supports_block_size(block_size) + ] + if rejecting: + raise ValueError( + f"Aligned KV cache block size {block_size} is not supported by " + f"attention backend(s): {', '.join(rejecting)}." + ) @classmethod def _align_heterogeneous_kv_block_size( diff --git a/vllm/platforms/xpu.py b/vllm/platforms/xpu.py index b3399bd53e..6b7db6d0dc 100644 --- a/vllm/platforms/xpu.py +++ b/vllm/platforms/xpu.py @@ -452,6 +452,7 @@ class XPUPlatform(Platform): if new_block_size == cache_config.block_size: return + pre_block_size = cache_config.block_size if cache_config.mamba_cache_mode == "align": cache_config.mamba_block_size = new_block_size original_mamba_page_size_padded = cache_config.mamba_page_size_padded @@ -474,6 +475,10 @@ class XPUPlatform(Platform): original_mamba_page_size_padded, cache_config.mamba_page_size_padded, ) + # This rounding runs after super()'s check, so check its result too. + cls._check_aligned_block_size( + vllm_config, cls._find_non_ssm_backends(vllm_config), pre_block_size + ) @classmethod def support_hybrid_kv_cache(cls) -> bool: From 89439db7276cf165377c088496e20098c8137727 Mon Sep 17 00:00:00 2001 From: Lio Einaudi Date: Sun, 4 Oct 2026 23:55:18 +0800 Subject: [PATCH 16/22] [CI][Docs] Fix decode/prefill consistency test prefix construction; add GLM-4-9B and Phi-4 to batch-invariant tested models (#53692) Signed-off-by: Yifan Da <1363818765@qq.com> Signed-off-by: LioEinaudi Signed-off-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: Yifan Da <1363818765@qq.com> Co-authored-by: CodeBuddy Co-authored-by: Claude Fable 5.1 Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> --- docs/features/batch_invariance.md | 3 +- tests/v1/determinism/test_batch_invariance.py | 51 ++++++++----------- 2 files changed, 22 insertions(+), 32 deletions(-) diff --git a/docs/features/batch_invariance.md b/docs/features/batch_invariance.md index a88c83ee57..07d6d803a7 100644 --- a/docs/features/batch_invariance.md +++ b/docs/features/batch_invariance.md @@ -130,7 +130,8 @@ Batch invariance has been tested and verified on the following models: - **Llama 3**: Llama3.1 and 3.2 series, `meta-llama/Llama-3.2-1B-Instruct`, `meta-llama/Llama-3.2-3B-Instruct` for example - **GPT-OSS**: `openai/gpt-oss-20b`, `openai/gpt-oss-120b` - **Mistral**: `mistralai/Mistral-7B-v0.3` -- **Phi series**: `microsoft/Phi-3.5-mini-instruct` +- **GLM-4 series**: `zai-org/glm-4-9b-chat` +- **Phi series**: `microsoft/Phi-3.5-mini-instruct`, `microsoft/phi-4` - **Granite 3.1 (MoE)**: `ibm-granite/granite-3.1-1b-a400m-instruct`, `ibm-granite/granite-3.1-3b-a800m-instruct` - **Granite 3.1 (Dense)**: `ibm-granite/granite-3.1-2b-instruct`, `ibm-granite/granite-3.1-8b-instruct` - **EXAONE 4.0 series**: `LGAI-EXAONE/EXAONE-4.0-1.2B`, `LGAI-EXAONE/EXAONE-4.0.1-32B`, `LGAI-EXAONE/EXAONE-4.0-32B` diff --git a/tests/v1/determinism/test_batch_invariance.py b/tests/v1/determinism/test_batch_invariance.py index f49d171f6f..4523627151 100644 --- a/tests/v1/determinism/test_batch_invariance.py +++ b/tests/v1/determinism/test_batch_invariance.py @@ -18,6 +18,7 @@ from utils import ( import vllm.envs as envs from vllm import LLM, SamplingParams +from vllm.inputs import TokensPrompt @skip_unsupported @@ -778,32 +779,16 @@ def test_decode_logprobs_match_prefill_logprobs( print(f"[Prompt {prompt_idx}] Generated {len(token_ids)} tokens: {token_ids}") print(f"[Prompt {prompt_idx}] Decode logprobs: {decode_logprobs.tolist()}") - # Step 2: For each token position, run prefill and compare + # Step 2: For each token position, run prefill and compare. + prompt_token_ids = list(decode_output.prompt_token_ids or []) + print(f"\n[Prompt {prompt_idx}] Verifying each token via prefill...") for token_idx in range(len(token_ids)): - # Construct the prefix up to (but not including) this token + # Construct the prefix up to (but not including) this token using + # token ids directly. current_token = token_ids[token_idx] - - # We need to detokenize to get the text prefix - # For this, we'll use the tokenizer from the LLM - # However, the LLM API doesn't expose tokenizer easily, so we'll - # construct the prefix by decoding from the original prompt - - # Get text up to this point by using the output text - # This is approximate but should work for verification - if token_idx == 0: - prefix_prompt = prompt - else: - # No per-token text, so regenerate with max_tokens = token_idx - # to get the prefix text - prefix_sp = SamplingParams( - temperature=0.0, - max_tokens=token_idx, - logprobs=1, - ) - prefix_output = llm.generate([prompt], prefix_sp, use_tqdm=False)[0] - prefix_prompt = prompt + prefix_output.outputs[0].text + prefix_token_ids = prompt_token_ids + list(token_ids[:token_idx]) # Now run prefill with max_tokens=1 to get the logprob of the next token prefill_sp = SamplingParams( @@ -814,19 +799,23 @@ def test_decode_logprobs_match_prefill_logprobs( print( f" [Token {token_idx}] Running prefill for prefix " - f"(len={len(prefix_prompt)})..." + f"(num_tokens={len(prefix_token_ids)})..." + ) + prefill_output = llm.generate( + [TokensPrompt(prompt_token_ids=prefix_token_ids)], + prefill_sp, + use_tqdm=False, + )[0] + prefill_logprobs, prefill_token_ids_out = _extract_step_logprobs( + prefill_output ) - prefill_output = llm.generate([prefix_prompt], prefill_sp, use_tqdm=False)[ - 0 - ] - prefill_logprobs, prefill_token_ids = _extract_step_logprobs(prefill_output) if prefill_logprobs is None: print(f" [Token {token_idx}] Warning: No prefill logprobs available") continue # The first token from prefill should match the current token - prefill_token = prefill_token_ids[0] + prefill_token = prefill_token_ids_out[0] prefill_logprob = prefill_logprobs[0].item() decode_logprob = decode_logprobs[token_idx].item() @@ -851,7 +840,7 @@ def test_decode_logprobs_match_prefill_logprobs( "decode_logprob": decode_logprob, "prefill_logprob": prefill_logprob, "prompt_text": prompt[:100], - "prefix_text": prefix_prompt[:100], + "prefix_token_ids": prefix_token_ids[:50], } ) print(f" [Token {token_idx}] ✗ TOKEN MISMATCH!") @@ -871,7 +860,7 @@ def test_decode_logprobs_match_prefill_logprobs( "prefill_logprob": prefill_logprob, "diff": diff, "prompt_text": prompt[:100], - "prefix_text": prefix_prompt[:100], + "prefix_token_ids": prefix_token_ids[:50], "decode_all_tokens": token_ids, "decode_all_logprobs": decode_logprobs.tolist(), } @@ -908,7 +897,7 @@ def test_decode_logprobs_match_prefill_logprobs( for i, fail in enumerate(failures[:5]): # Show first 5 failures per prompt print(f"\n [Failure {i + 1}] Token position {fail['token_idx']}:") print(f" Reason: {fail['reason']}") - print(f" Prefix text: '{fail['prefix_text']}...'") + print(f" Prefix token ids: {fail['prefix_token_ids']}...") print( f" Decode: token={fail['decode_token']}, " f"logprob={fail['decode_logprob']:.10f}" From b0eb87fe4943316781913d99a8ef8413e425decb Mon Sep 17 00:00:00 2001 From: Stefan Wang <1fannnw@gmail.com> Date: Sun, 4 Oct 2026 09:42:49 -0700 Subject: [PATCH 17/22] [Bugfix][Frontend] Return streaming errors before the first token (#40986) Signed-off-by: 1fanwang <1fannnw@gmail.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> --- .../openai/chat_completion/test_chat_error.py | 64 +++++++++++++++- .../completion/test_completion_error.py | 76 +++++++++++++++++-- .../openai/chat_completion/serving.py | 6 +- vllm/entrypoints/openai/completion/serving.py | 5 +- 4 files changed, 138 insertions(+), 13 deletions(-) diff --git a/tests/entrypoints/openai/chat_completion/test_chat_error.py b/tests/entrypoints/openai/chat_completion/test_chat_error.py index a4126ba5f7..baa8bcec59 100644 --- a/tests/entrypoints/openai/chat_completion/test_chat_error.py +++ b/tests/entrypoints/openai/chat_completion/test_chat_error.py @@ -248,7 +248,8 @@ async def test_renderer_only_chat_request_skips_mm_cache(): @pytest.mark.asyncio -async def test_chat_error_stream(): +@pytest.mark.parametrize("include_usage", [False, True]) +async def test_chat_error_stream(include_usage: bool): """Test finish_reason='error' returns 500 InternalServerError (streaming).""" mock_engine = MagicMock(spec=AsyncLLM) mock_engine.errored = False @@ -313,6 +314,67 @@ async def test_chat_error_stream(): messages=[{"role": "user", "content": "Test prompt"}], max_tokens=10, stream=True, + stream_options={"include_usage": include_usage}, + ) + + response = await serving_chat.create_chat_completion(request) + + chunks = [] + async for chunk in response: + chunks.append(chunk) + + assert len(chunks) >= 2 + assert any("Internal server error" in chunk for chunk in chunks), ( + f"Expected error message in chunks: {chunks}" + ) + assert chunks[-1] == "data: [DONE]\n\n" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("include_usage", [False, True]) +async def test_chat_error_stream_first_chunk(include_usage: bool): + """An error before the first token must reach the streaming client.""" + mock_engine = MagicMock(spec=AsyncLLM) + mock_engine.errored = False + mock_engine.model_config = MockModelConfig() + mock_engine.input_processor = MagicMock() + mock_engine.renderer = _build_renderer(mock_engine.model_config) + + serving_chat = _build_serving_chat(mock_engine) + + completion_output = CompletionOutput( + index=0, + text="", + token_ids=[], + cumulative_logprob=None, + logprobs=None, + finish_reason="error", + ) + + request_output = RequestOutput( + request_id="test-id", + prompt="Test prompt", + prompt_token_ids=[1, 2, 3], + prompt_logprobs=None, + outputs=[completion_output], + finished=True, + metrics=None, + lora_request=None, + encoder_prompt=None, + encoder_prompt_token_ids=None, + ) + + async def mock_generate(*args, **kwargs): + yield request_output + + mock_engine.generate = MagicMock(side_effect=mock_generate) + + request = ChatCompletionRequest( + model=MODEL_NAME, + messages=[{"role": "user", "content": "Test prompt"}], + max_tokens=10, + stream=True, + stream_options={"include_usage": include_usage}, ) response = await serving_chat.create_chat_completion(request) diff --git a/tests/entrypoints/openai/completion/test_completion_error.py b/tests/entrypoints/openai/completion/test_completion_error.py index 05811555a9..4056ea49d6 100644 --- a/tests/entrypoints/openai/completion/test_completion_error.py +++ b/tests/entrypoints/openai/completion/test_completion_error.py @@ -383,10 +383,14 @@ async def test_completion_suffix_supports_text_prompt_list(): assert result == [{"ok": True}] assert online_renderer.preprocess_completion.call_args.kwargs["prompt_input"] == [ - f"{DEEPSEEK_V4_FIM_BEGIN}prefix A" - f"{DEEPSEEK_V4_FIM_HOLE} suffix{DEEPSEEK_V4_FIM_END}", - f"{DEEPSEEK_V4_FIM_BEGIN}prefix B" - f"{DEEPSEEK_V4_FIM_HOLE} suffix{DEEPSEEK_V4_FIM_END}", + ( + f"{DEEPSEEK_V4_FIM_BEGIN}prefix A" + f"{DEEPSEEK_V4_FIM_HOLE} suffix{DEEPSEEK_V4_FIM_END}" + ), + ( + f"{DEEPSEEK_V4_FIM_BEGIN}prefix B" + f"{DEEPSEEK_V4_FIM_HOLE} suffix{DEEPSEEK_V4_FIM_END}" + ), ] @@ -619,7 +623,8 @@ async def test_renderer_only_completion_request_skips_mm_cache(): @pytest.mark.asyncio -async def test_completion_error_stream(): +@pytest.mark.parametrize("include_usage", [False, True]) +async def test_completion_error_stream(include_usage: bool): """Test finish_reason='error' returns 500 InternalServerError (streaming).""" mock_engine = MagicMock(spec=AsyncLLM) mock_engine.errored = False @@ -684,6 +689,67 @@ async def test_completion_error_stream(): prompt="Test prompt", max_tokens=10, stream=True, + stream_options={"include_usage": include_usage}, + ) + + response = await serving_completion.create_completion(request) + + chunks = [] + async for chunk in response: + chunks.append(chunk) + + assert len(chunks) >= 2 + assert any("Internal server error" in chunk for chunk in chunks), ( + f"Expected error message in chunks: {chunks}" + ) + assert chunks[-1] == "data: [DONE]\n\n" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("include_usage", [False, True]) +async def test_completion_error_stream_first_chunk(include_usage: bool): + """An error before the first token must reach the streaming client.""" + mock_engine = MagicMock(spec=AsyncLLM) + mock_engine.errored = False + mock_engine.model_config = MockModelConfig() + mock_engine.input_processor = MagicMock() + mock_engine.renderer = _build_renderer(mock_engine.model_config) + + serving_completion = _build_serving_completion(mock_engine) + + completion_output = CompletionOutput( + index=0, + text="", + token_ids=[], + cumulative_logprob=None, + logprobs=None, + finish_reason="error", + ) + + request_output = RequestOutput( + request_id="test-id", + prompt="Test prompt", + prompt_token_ids=[1, 2, 3], + prompt_logprobs=None, + outputs=[completion_output], + finished=True, + metrics=None, + lora_request=None, + encoder_prompt=None, + encoder_prompt_token_ids=None, + ) + + async def mock_generate(*args, **kwargs): + yield request_output + + mock_engine.generate = MagicMock(side_effect=mock_generate) + + request = CompletionRequest( + model=MODEL_NAME, + prompt="Test prompt", + max_tokens=10, + stream=True, + stream_options={"include_usage": include_usage}, ) response = await serving_completion.create_completion(request) diff --git a/vllm/entrypoints/openai/chat_completion/serving.py b/vllm/entrypoints/openai/chat_completion/serving.py index 060d22b031..400089c963 100644 --- a/vllm/entrypoints/openai/chat_completion/serving.py +++ b/vllm/entrypoints/openai/chat_completion/serving.py @@ -631,6 +631,8 @@ class OpenAIServingChat(GenerateBaseServing): if finish_reason_sent[i]: continue + self._raise_if_error(output.finish_reason, request_id) + if request.logprobs and ( request.top_logprobs is not None or request.logprob_token_ids ): @@ -758,10 +760,6 @@ class OpenAIServingChat(GenerateBaseServing): # if the model is finished generating else: - # check for error finish reason and abort streaming - # finish_reason='error' indicates a retryable error - self._raise_if_error(output.finish_reason, request_id) - # Send the finish response for each request.n only once # In OpenAI's API, when a tool is called, the # finish_reason is: diff --git a/vllm/entrypoints/openai/completion/serving.py b/vllm/entrypoints/openai/completion/serving.py index d7e235da1b..81dd00f0b3 100644 --- a/vllm/entrypoints/openai/completion/serving.py +++ b/vllm/entrypoints/openai/completion/serving.py @@ -326,6 +326,8 @@ class OpenAIServingCompletion(GenerateBaseServing): for output in res.outputs: i = output.index + prompt_idx * num_choices + finish_reason = output.finish_reason + self._raise_if_error(finish_reason, request_id) # Useful when request.return_token_ids is True # Returning prompt token IDs shares the same logic @@ -392,11 +394,8 @@ class OpenAIServingCompletion(GenerateBaseServing): previous_text_lens[i] += len(delta_text) previous_num_tokens[i] += len(output.token_ids) - finish_reason = output.finish_reason stop_reason = output.stop_reason - self._raise_if_error(finish_reason, request_id) - chunk = CompletionStreamResponse( id=request_id, object="text_completion", From 155488d853a0bc42df227dbfc74005b3fd488e94 Mon Sep 17 00:00:00 2001 From: Prudhvi Vuda <53619858+Prudhvivuda@users.noreply.github.com> Date: Sun, 4 Oct 2026 13:07:46 -0400 Subject: [PATCH 18/22] [Bugfix][Determinism] Preserve output dtype in batch-invariant mean (#59106) Signed-off-by: Prudhvivuda Signed-off-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- vllm/model_executor/determinism/batch_invariant.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/model_executor/determinism/batch_invariant.py b/vllm/model_executor/determinism/batch_invariant.py index 6490e69a60..403f475702 100644 --- a/vllm/model_executor/determinism/batch_invariant.py +++ b/vllm/model_executor/determinism/batch_invariant.py @@ -918,7 +918,7 @@ def mean_batch_invariant(input, dim, keepdim=False, dtype: torch.dtype | None = for d in sorted_dims: result = result.squeeze(d) - return result + return result if dtype is not None else result.to(input.dtype) @triton.jit From 7867d6c52d4542c1b5da641f8ca79124be497ee3 Mon Sep 17 00:00:00 2001 From: Tyrone <71038642+TyroneNel@users.noreply.github.com> Date: Sun, 4 Oct 2026 22:45:51 +0200 Subject: [PATCH 19/22] [Frontend] Name the served models in the model-not-found 404 (#59889) Signed-off-by: TyroneNel <71038642+TyroneNel@users.noreply.github.com> Signed-off-by: TyroneNel Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- vllm/entrypoints/serve/engine/serving.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/vllm/entrypoints/serve/engine/serving.py b/vllm/entrypoints/serve/engine/serving.py index 86ba524946..8b42449552 100644 --- a/vllm/entrypoints/serve/engine/serving.py +++ b/vllm/entrypoints/serve/engine/serving.py @@ -63,8 +63,12 @@ class BaseServing: ): error_response = load_result + served_names = ", ".join(model.name for model in self.models.base_model_paths) return error_response or self.create_error_response( - message=f"The model `{request.model}` does not exist.", + message=( + f"The model `{request.model}` does not exist. " + f"Valid aliases: {served_names}." + ), err_type="NotFoundError", status_code=HTTPStatus.NOT_FOUND, param="model", From 138810056093301f4881050fcf2b1786939da387 Mon Sep 17 00:00:00 2001 From: Shenglei Fu <117230642+ShengleiFu@users.noreply.github.com> Date: Sun, 4 Oct 2026 15:46:16 -0700 Subject: [PATCH 20/22] [Kernel][LoRA] Add deterministic split-K=8 LoRA shrink for batch invariance (#59377) Signed-off-by: Shenglei Fu <117230642+ShengleiFu@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> --- .buildkite/test_areas/misc.yaml | 4 + docs/features/batch_invariance.md | 1 + .../test_lora_shrink_batch_invariant.py | 99 +++++++++++ vllm/lora/ops/triton_ops/kernel_utils.py | 26 ++- vllm/lora/ops/triton_ops/lora_shrink_op.py | 155 +++++++++++++++++- 5 files changed, 275 insertions(+), 10 deletions(-) create mode 100644 tests/v1/determinism/test_lora_shrink_batch_invariant.py diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 602491e0e0..cbcfaeaea7 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -655,12 +655,14 @@ steps: source_file_dependencies: - vllm/v1/attention - vllm/model_executor/layers + - vllm/lora/ops/triton_ops - tests/v1/determinism/ commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pip install pytest-timeout pytest-forked - pytest -v -s v1/determinism/test_batch_invariance.py - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py + - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k FLASH_ATTN mirror: @@ -690,12 +692,14 @@ steps: source_file_dependencies: - vllm/v1/attention - vllm/model_executor/layers + - vllm/lora/ops/triton_ops - tests/v1/determinism/ commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pip install pytest-timeout pytest-forked - pytest -v -s v1/determinism/test_batch_invariance.py - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py + - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py - pytest -v -s v1/determinism/test_batch_invariance_vlm.py - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k FLASH_ATTN diff --git a/docs/features/batch_invariance.md b/docs/features/batch_invariance.md index 07d6d803a7..e82336d241 100644 --- a/docs/features/batch_invariance.md +++ b/docs/features/batch_invariance.md @@ -153,6 +153,7 @@ When batch invariance is enabled, vLLM: 5. On CUDA devices with tuned matmul table entries for the model's bf16 unquantized forward linear layers (Ada, Hopper, Blackwell), runs without `torch.compile` using breakable CUDA graphs so tile configs follow the runtime batch size; set `VLLM_USE_BREAKABLE_CUDAGRAPH=0` to opt out (not applied when sequence parallelism / async TP are enabled, since those are `torch.compile` passes). +6. Splits the LoRA shrink reduction into 8 fixed K chunks that are summed in a fixed order instead of with atomics !!! warning Batch invariance under tensor parallelism is not yet supported for all-reduces whose size isn't a multiple of 16 bytes, for example a hidden size that isn't a multiple of 8 in fp16/bf16. Such tensors can switch between NCCL and custom all-reduce depending on batch size. All validated models meet this requirement. diff --git a/tests/v1/determinism/test_lora_shrink_batch_invariant.py b/tests/v1/determinism/test_lora_shrink_batch_invariant.py new file mode 100644 index 0000000000..2a7029a603 --- /dev/null +++ b/tests/v1/determinism/test_lora_shrink_batch_invariant.py @@ -0,0 +1,99 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Batch-invariant LoRA shrink: fixed split-K partials, fixed-order reduction.""" + +import pytest +import torch +from utils import skip_if_not_cuda + +import vllm.lora.ops.triton_ops.lora_shrink_op as shrink_op +import vllm.lora.ops.triton_ops.utils as lora_utils +from vllm.lora.ops.triton_ops import LoRAKernelMeta, lora_shrink +from vllm.lora.ops.triton_ops.utils import _LORA_A_PTR_DICT +from vllm.utils.torch_utils import set_random_seed + +DEVICE = "cuda" + + +@pytest.fixture(autouse=True) +def batch_invariant(monkeypatch): + monkeypatch.setattr(lora_utils, "is_batch_invariant", True) + monkeypatch.setattr(shrink_op, "is_batch_invariant", True) + lora_utils.get_lora_op_configs.cache_clear() + yield + lora_utils.get_lora_op_configs.cache_clear() + + +def _shrink(inputs, weights, mapping, scaling): + num_tokens = inputs.size(0) + meta = LoRAKernelMeta.make(weights[0].size(0), num_tokens, device=DEVICE) + meta.prepare_tensors(mapping) + out = torch.zeros( + (len(weights), num_tokens, weights[0].size(1)), + dtype=torch.float32, + device=DEVICE, + ) + _LORA_A_PTR_DICT.clear() + args = meta.meta_args(token_nums=num_tokens, specialize_active_lora=False) + lora_shrink(inputs, weights, out, *args, scaling) + return out + + +@skip_if_not_cuda +@pytest.mark.parametrize( + "hidden_size,rank,nslices,dtype", + [ + (2560, 8, 3, torch.bfloat16), + (9728, 16, 1, torch.float16), + (4097, 64, 1, torch.bfloat16), # K tail block + (1024, 8, 1, torch.float16), # K < 8 * BLOCK_K: some splits are empty + ], +) +def test_lora_shrink_matches_reference(hidden_size, rank, nslices, dtype): + set_random_seed(0) + inputs = torch.randn((33, hidden_size), dtype=dtype, device=DEVICE) + weights = [ + torch.randn((2, rank, hidden_size), dtype=dtype, device=DEVICE) + for _ in range(nslices) + ] + mapping = torch.randint(-1, 2, (33,), dtype=torch.int32, device=DEVICE) + out = _shrink(inputs, weights, mapping, 0.5).double().cpu() + + ids, x = mapping.cpu(), inputs.double().cpu() + ref = torch.zeros_like(out) + for s, weight in enumerate(weights): + for lora_id in (0, 1): + rows = ids == lora_id + ref[s, rows] = 0.5 * x[rows] @ weight[lora_id].double().cpu().T + torch.testing.assert_close(out, ref, rtol=5e-3, atol=5e-3) + + +@skip_if_not_cuda +@pytest.mark.parametrize("nslices", [1, 3]) +def test_lora_shrink_invariant_across_batch_sizes(monkeypatch, nslices): + reduce_kernel, reduce_grids = shrink_op._lora_shrink_reduce_kernel, [] + + class RecordReduce: + def __getitem__(self, grid): + reduce_grids.append(grid) + return reduce_kernel[grid] + + monkeypatch.setattr(shrink_op, "_lora_shrink_reduce_kernel", RecordReduce()) + set_random_seed(0) + # K = 2560 is ten BLOCK_K=256 blocks, so all 8 splits contribute. + dtype, hidden_size = torch.bfloat16, 2560 + weights = [ + torch.randn((3, 16, hidden_size), dtype=dtype, device=DEVICE) + for _ in range(nslices) + ] + target = torch.randn((1, hidden_size), dtype=dtype, device=DEVICE) + target_id = torch.zeros(1, dtype=torch.int32, device=DEVICE) + expected = _shrink(target, weights, target_id, 0.7)[:, 0] + for num_tokens in (31, 32, 33, 127, 128, 129): + inputs = torch.randn((num_tokens, hidden_size), dtype=dtype, device=DEVICE) + mapping = torch.randint(-1, 3, (num_tokens,), dtype=torch.int32, device=DEVICE) + pos = num_tokens // 2 + inputs[pos], mapping[pos] = target[0], 0 + out = _shrink(inputs, weights, mapping, 0.7) + assert torch.equal(out[:, pos], expected), num_tokens + assert len(reduce_grids) == 7 diff --git a/vllm/lora/ops/triton_ops/kernel_utils.py b/vllm/lora/ops/triton_ops/kernel_utils.py index fd6ae44428..bc07817ae1 100644 --- a/vllm/lora/ops/triton_ops/kernel_utils.py +++ b/vllm/lora/ops/triton_ops/kernel_utils.py @@ -257,6 +257,7 @@ def do_shrink_kernel( output_d0_stride, output_d1_stride, output_d2_stride, + output_split_stride, scaling, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, @@ -265,6 +266,7 @@ def do_shrink_kernel( SPLIT_K: tl.constexpr, SLICE_NUM: tl.constexpr, USE_GDC: tl.constexpr, + STORE_PARTIALS: tl.constexpr, ): """Given an array of integers that identifies the rows of A, ram, a lora index that identifies which LoRA to use from lora_ptr, lora_index, @@ -324,17 +326,31 @@ def do_shrink_kernel( # Identify the C output pointers to store the results of the accumulator. offset_cn = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N offset_cm = tl.arange(0, BLOCK_M) - cur_out_ptr = out_ptr if SLICE_NUM == 1 else out_ptr + slice_id * output_d0_stride + if STORE_PARTIALS: + cur_out_ptr = ( + out_ptr + + slice_id.to(tl.int64) * output_d0_stride + + pid_sk.to(tl.int64) * output_split_stride + ) + else: + cur_out_ptr = ( + out_ptr if SLICE_NUM == 1 else out_ptr + slice_id * output_d0_stride + ) c_ptr = ( cur_out_ptr + ram[:, None] * output_d1_stride + offset_cn[None, :] * output_d2_stride ) c_mask = (offset_cm[:, None] < M_LEN) & (offset_cn[None, :] < N) - accumulator *= scaling - # handles write-back with reduction-splitting - if SPLIT_K == 1: + if STORE_PARTIALS: + # Unscaled FP32 partial per K split; the reducer adds the splits in a + # fixed order and applies scaling once. tl.store(c_ptr, accumulator, mask=c_mask) else: - tl.atomic_add(c_ptr, accumulator, mask=c_mask, sem="relaxed") + accumulator *= scaling + # handles write-back with reduction-splitting + if SPLIT_K == 1: + tl.store(c_ptr, accumulator, mask=c_mask) + else: + tl.atomic_add(c_ptr, accumulator, mask=c_mask, sem="relaxed") diff --git a/vllm/lora/ops/triton_ops/lora_shrink_op.py b/vllm/lora/ops/triton_ops/lora_shrink_op.py index ec29065b9a..abe63bd41b 100644 --- a/vllm/lora/ops/triton_ops/lora_shrink_op.py +++ b/vllm/lora/ops/triton_ops/lora_shrink_op.py @@ -13,11 +13,90 @@ from vllm.lora.ops.triton_ops.kernel_utils import do_shrink_kernel from vllm.lora.ops.triton_ops.utils import ( _get_lora_a_ptr, get_lora_op_configs, + is_batch_invariant, supports_pdl, ) from vllm.triton_utils import tl, triton from vllm.utils.torch_utils import direct_register_custom_op +# Batch-invariant shrink splits K into this many fixed partials (no atomics). +_BI_SPLIT_K = 8 + + +@triton.jit +def _lora_shrink_reduce_kernel( + partial_ptr, + out_ptr, + M, + N: tl.constexpr, + token_indices_sorted_by_lora_ids, + num_tokens_per_lora, + lora_token_start_loc, + lora_ids, + scaling, + partial_d0_stride, + partial_d1_stride, + partial_d2_stride, + partial_d3_stride, + output_d0_stride, + output_d1_stride, + output_d2_stride, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + SPLIT_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, +): + cta_n_num = tl.cdiv(N, BLOCK_N) + cta_m_num = tl.cdiv(M, BLOCK_M) + pid_m_n = tl.program_id(axis=0) + num_pid_in_group = GROUP_SIZE_M * cta_n_num + group_id = pid_m_n // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(cta_m_num - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + ((pid_m_n % num_pid_in_group) % group_size_m) + pid_n = (pid_m_n % num_pid_in_group) // group_size_m + + slice_id = tl.program_id(axis=1) + lora_idx = tl.program_id(axis=2) + lora_id = tl.load(lora_ids + lora_idx) + if lora_id == -1: + return + + lora_m_size = tl.load(num_tokens_per_lora + lora_idx) + cta_m_offset = pid_m * BLOCK_M + if cta_m_offset >= lora_m_size: + return + + cta_m_len = min(BLOCK_M, lora_m_size - cta_m_offset) + lora_m_indices_start = tl.load(lora_token_start_loc + lora_idx) + cta_lora_seq_indices = ( + token_indices_sorted_by_lora_ids + lora_m_indices_start + cta_m_offset + ) + offset_m = tl.arange(0, BLOCK_M) % cta_m_len + ram = tl.load(cta_lora_seq_indices + offset_m) + offset_n = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N + offset_cm = tl.arange(0, BLOCK_M) + mask = (offset_cm[:, None] < cta_m_len) & (offset_n[None, :] < N) + + accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + for split_id in range(SPLIT_K): + partial_in = ( + partial_ptr + + slice_id.to(tl.int64) * partial_d0_stride + + tl.cast(split_id, tl.int64) * partial_d1_stride + + ram[:, None] * partial_d2_stride + + offset_n[None, :] * partial_d3_stride + ) + accumulator += tl.load(partial_in, mask=mask, other=0.0) + + output = ( + out_ptr + + slice_id * output_d0_stride + + ram[:, None] * output_d1_stride + + offset_n[None, :] * output_d2_stride + ) + tl.store(output, accumulator * scaling, mask=mask) + @triton.jit def _lora_shrink_kernel( @@ -49,7 +128,15 @@ def _lora_shrink_kernel( SLICE_NUM: tl.constexpr, USE_GDC: tl.constexpr, launch_pdl: tl.constexpr, + STORE_PARTIALS: tl.constexpr = False, + output_split_stride=0, + PARTIAL_N: tl.constexpr = 0, + PARTIAL_K: tl.constexpr = 0, ): + if STORE_PARTIALS: + # The split-K partial stage specializes on N and K. + N = PARTIAL_N + K = PARTIAL_K cta_n_num = tl.cdiv(N, BLOCK_N) cta_m_num = tl.cdiv(M, BLOCK_M) @@ -116,6 +203,7 @@ def _lora_shrink_kernel( output_d0_stride, output_d1_stride, output_d2_stride, + output_split_stride, scaling, BLOCK_M, BLOCK_N, @@ -124,6 +212,7 @@ def _lora_shrink_kernel( SPLIT_K, SLICE_NUM, USE_GDC, + STORE_PARTIALS, ) @@ -167,6 +256,23 @@ def _lora_shrink( """ assert no_lora_flag_cpu.numel() == 1 + if inputs.size(0) == 0: + # An empty first call has no work to do. + return + partials: torch.Tensor | None = None + if is_batch_invariant: + # Allocated before the no-LoRA return so profiling sees the scratch. + # Per-call allocation keeps each CUDA graph's scratch in its own pool. + partials = torch.empty( + ( + output_tensor.size(0), + _BI_SPLIT_K, + inputs.size(0), + output_tensor.size(-1), + ), + dtype=torch.float32, + device=output_tensor.device, + ) if no_lora_flag_cpu.item(): # None of the inputs require LoRA. return @@ -213,6 +319,11 @@ def _lora_shrink( NUM_STAGES = kernel_config["num_stages"] NUM_CTAS = kernel_config["num_ctas"] GROUP_SIZE_M = kernel_config.get("group_size_m", 8) + + # Every M must keep the same K partition for batch invariance. + if partials is not None: + BLOCK_K = 256 + SPLIT_K = _BI_SPLIT_K EVEN_K = K % (BLOCK_K * SPLIT_K) == 0 # type: ignore # TODO (varun): This grid formulation maximizes parallelization at the @@ -225,11 +336,24 @@ def _lora_shrink( ) # PDL only works when dual-stream is being used. - use_gdc = supports_pdl(inputs.device) and envs.VLLM_LORA_ENABLE_DUAL_STREAM + use_gdc = ( + supports_pdl(inputs.device) + and envs.VLLM_LORA_ENABLE_DUAL_STREAM + and partials is None + ) + first_out = output_tensor if partials is None else partials + partial_args = {} + if partials is not None: + partial_args = dict( + STORE_PARTIALS=True, + output_split_stride=partials.stride(1), + PARTIAL_N=N, + PARTIAL_K=K, + ) _lora_shrink_kernel[grid]( inputs, lora_ptr_tensor, - output_tensor, + first_out, M, N, K, @@ -243,9 +367,9 @@ def _lora_shrink( lora_strides_d0, lora_strides_d1, lora_strides_d2, - output_tensor.stride(0), - output_tensor.stride(1), - output_tensor.stride(2), + first_out.stride(0), + first_out.stride(-2), + first_out.stride(-1), BLOCK_M, BLOCK_N, BLOCK_K, @@ -258,7 +382,28 @@ def _lora_shrink( num_ctas=NUM_CTAS, num_stages=NUM_STAGES, launch_pdl=use_gdc, + **partial_args, ) + if partials is not None: + _lora_shrink_reduce_kernel[(grid[0] // SPLIT_K, *grid[1:])]( + partials, + output_tensor, + M, + N, + token_indices_sorted_by_lora_ids, + num_tokens_per_lora, + lora_token_start_loc, + lora_ids, + scaling, + *partials.stride(), + *output_tensor.stride(), + BLOCK_M, + BLOCK_N, + SPLIT_K, + GROUP_SIZE_M, + num_warps=NUM_WARPS, + num_stages=NUM_STAGES, + ) return From 49e0f4797853da0d1da7fdfe794f327fb5d5843d Mon Sep 17 00:00:00 2001 From: bluemelov <73552292+bluemelov1@users.noreply.github.com> Date: Mon, 5 Oct 2026 03:49:41 +0200 Subject: [PATCH 21/22] [Bugfix][DSv4.1] Keep the compressor ring out of the null block (#58560) Signed-off-by: bluemelov Signed-off-by: Yifan Qiao Co-authored-by: bluemelov Co-authored-by: Cursor Co-authored-by: Yifan Qiao Co-authored-by: Claude Opus 5.5 --- tests/kernels/test_compressor_kv_cache.py | 28 ++++++++++++----------- vllm/models/deepseek_v41/compressor.py | 6 ++++- 2 files changed, 20 insertions(+), 14 deletions(-) diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index 3e40a3882f..020dd607b6 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -711,7 +711,8 @@ def test_v41_rope_insert_plain_row(compress_ratio: int, store_fp8: bool): def test_v41_compressor_metadata_maps_tokens_to_their_ring(): """The ring group's generic slot mapping is disabled (all PAD), so the builder must map every real token to ``ring_block * capacity + pos % - capacity`` and keep padding tokens at PAD.""" + capacity``. Padding tokens and requests on the null block (block 0, as in + dummy batches) must stay at PAD.""" from unittest.mock import MagicMock from vllm.models.deepseek_v41.compressor import CompressorMetadataBuilder @@ -731,29 +732,30 @@ def test_v41_compressor_metadata_maps_tokens_to_their_ring(): device = torch.device("cuda") builder = CompressorMetadataBuilder(spec, ["state"], vllm_config, device) - # Two requests: 3 tokens at positions 13..15 on ring block 5, then 2 - # tokens at positions 7..8 on ring block 2; three padding tokens. - query_start_loc = torch.tensor([0, 3, 5], dtype=torch.int32, device=device) - positions = torch.tensor([13, 14, 15, 7, 8, 0, 0, 0], device=device) - block_table = torch.tensor([[5], [2]], dtype=torch.int32, device=device) + # Three requests: 3 tokens at positions 13..15 on ring block 5, 2 tokens + # at positions 7..8 on ring block 2, then 2 tokens on the null block; + # three padding tokens. + query_start_loc = torch.tensor([0, 3, 5, 7], dtype=torch.int32, device=device) + positions = torch.tensor([13, 14, 15, 7, 8, 0, 1, 0, 0, 0], device=device) + block_table = torch.tensor([[5], [2], [0]], dtype=torch.int32, device=device) common = CommonAttentionMetadata( query_start_loc=query_start_loc, query_start_loc_cpu=query_start_loc.cpu(), - seq_lens=torch.tensor([16, 9], dtype=torch.int32, device=device), - num_reqs=2, - num_actual_tokens=5, + seq_lens=torch.tensor([16, 9, 2], dtype=torch.int32, device=device), + num_reqs=3, + num_actual_tokens=7, max_query_len=3, max_seq_len=16, block_table_tensor=block_table, - slot_mapping=torch.full((8,), -1, dtype=torch.int64, device=device), + slot_mapping=torch.full((10,), -1, dtype=torch.int64, device=device), positions=positions, ) metadata = builder.build(0, common) - expected = [5 * 8 + 5, 5 * 8 + 6, 5 * 8 + 7, 2 * 8 + 7, 2 * 8 + 0, -1, -1, -1] - assert metadata.slot_mapping.tolist() == expected + ring = [5 * 8 + 5, 5 * 8 + 6, 5 * 8 + 7, 2 * 8 + 7, 2 * 8 + 0] + assert metadata.slot_mapping.tolist() == ring + [-1] * 5 assert metadata.query_start_loc is query_start_loc - assert metadata.token_to_req_indices.tolist() == [0, 0, 0, 1, 1] + assert metadata.token_to_req_indices.tolist() == [0, 0, 0, 1, 1, 2, 2] @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA stream coverage") diff --git a/vllm/models/deepseek_v41/compressor.py b/vllm/models/deepseek_v41/compressor.py index 54dc9c4a13..6f56ad8d5b 100644 --- a/vllm/models/deepseek_v41/compressor.py +++ b/vllm/models/deepseek_v41/compressor.py @@ -75,8 +75,12 @@ def _ring_slot_mapping_kernel( block = tl.load(block_table_ptr + req * block_table_stride, mask=valid, other=0) pos = tl.load(positions_ptr + offsets, mask=valid, other=0) slot = block.to(tl.int64) * CAPACITY + pos % CAPACITY + # Block 0 is the null block. A request without a ring block (dummy or padding + # rows, which the runners fill with the null block) must not write. tl.store( - slot_mapping_ptr + offsets, tl.where(valid, slot, -1), mask=offsets < num_tokens + slot_mapping_ptr + offsets, + tl.where(valid & (block != 0), slot, -1), + mask=offsets < num_tokens, ) From 4f52fa35efc7c56b77221460346e0e9eea28638b Mon Sep 17 00:00:00 2001 From: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com> Date: Mon, 5 Oct 2026 05:52:06 +0400 Subject: [PATCH 22/22] [Bugfix] Fix Qwen4Exp PLE embedding rejecting INC (AutoRound) checkpoints (#59990) Signed-off-by: Vadim Gimpelson Co-authored-by: Claude --- tests/models/qwen4_exp/test_ple.py | 19 +++++++++++++++++++ .../qwen4_exp/common/ngram_embedding.py | 6 ++++++ 2 files changed, 25 insertions(+) diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py index fb191c1b4f..85c9122c9d 100644 --- a/tests/models/qwen4_exp/test_ple.py +++ b/tests/models/qwen4_exp/test_ple.py @@ -20,6 +20,7 @@ from vllm.config import SpeculativeConfig, VllmConfig from vllm.config.compilation import CUDAGraphMode from vllm.config.quantization import QuantizationConfigArgs from vllm.model_executor.layers.quantization.fp8 import Fp8Config +from vllm.model_executor.layers.quantization.inc import INCConfig from vllm.model_executor.layers.quantization.modelopt import ( ModelOptMixedPrecisionConfig, ModelOptNvFp4Config, @@ -566,6 +567,24 @@ def test_ple_embedding_respects_modelopt_exclusion() -> None: ) +def test_ple_embedding_respects_inc_layer_config() -> None: + prefix = "model.layers.1.ple.ple_embedding.ngram_embedding" + quant_config = INCConfig( + weight_bits=4, + group_size=128, + block_name_to_quantize="model.layers", + extra_config={".*ple.*": {"bits": 16, "data_type": "float"}}, + ) + assert isinstance( + Qwen4ExpPLEEmbeddingMethod.from_quant_config(quant_config, prefix), + Qwen4ExpPLEUnquantizedEmbeddingMethod, + ) + + quant_config.extra_config = None + with pytest.raises(NotImplementedError, match="INCConfig"): + Qwen4ExpPLEEmbeddingMethod.from_quant_config(quant_config, prefix) + + def test_ple_embedding_dtype_overrides_modelopt_exclusion() -> None: prefix = "model.layers.1.ple.ple_embedding.ngram_embedding" quant_config = ModelOptNvFp4Config(exclude_modules=[prefix]) diff --git a/vllm/models/qwen4_exp/common/ngram_embedding.py b/vllm/models/qwen4_exp/common/ngram_embedding.py index 9e6eef03df..6b4ea9e98f 100644 --- a/vllm/models/qwen4_exp/common/ngram_embedding.py +++ b/vllm/models/qwen4_exp/common/ngram_embedding.py @@ -27,6 +27,7 @@ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tenso CompressedTensorsConfig, ) from vllm.model_executor.layers.quantization.fp8 import Fp8Config +from vllm.model_executor.layers.quantization.inc import INCConfig from vllm.model_executor.layers.quantization.modelopt import ( ModelOptMixedPrecisionConfig, ModelOptQuantConfigBase, @@ -190,6 +191,11 @@ class Qwen4ExpPLEEmbeddingMethod(QuantizeMethodBase): and quant_config.get_scheme_dict(None, layer_name=prefix) is None ): return Qwen4ExpPLEUnquantizedEmbeddingMethod() + if ( + isinstance(quant_config, INCConfig) + and not quant_config.config_parser.resolve(None, prefix).quantized + ): + return Qwen4ExpPLEUnquantizedEmbeddingMethod() if not isinstance(quant_config, Fp8Config): raise NotImplementedError( "Qwen4Exp PLE embedding does not support quantization config "