iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Merge remote-tracking branch 'upstream/main' into patches

Browse Source
This commit is contained in:
MOVZX
2026-10-05 08:52:31 +07:00
parent 85c31d47ad 4f52fa35ef
commit f83ecfac16
85 changed files with 3146 additions and 769 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.buildkite/test_areas/misc.yaml
+4
View File
@@ -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
.github/mergify.yml
+26
View File
@@ -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:
.github/workflows/issue_autolabel.yml
+21
View File
@@ -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" },
csrc/rocm/moe_q_gemm_rdna3.cu
+52 -30
View File
@@ -57,16 +57,6 @@ __forceinline__ __device__ bf16_t tzero<bf16_t>() {
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<unsigned long long*>(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<T, half>::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<T, half>::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<T, half>::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<const half*>(&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<const half2*>(&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<T, half>::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));
csrc/rocm/q_gemm_rdna3.cu
+460 -122
View File
@@ -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 <cstdint>
#include <cstdio>
@@ -92,40 +103,16 @@ __forceinline__ __device__ bf16_t tzero<bf16_t>() {
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 <typename T, int M_COUNT>
__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<T>(sc_row, n, scales);
if constexpr (std::is_same<T, half>::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<T, half>::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<T>(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<T, half>::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<T, half>::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<T, half>::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<T, half>::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<T, half>::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<const half*>(&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<const half2*>(&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<T, half>::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<T, half>::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<T, half>::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<T, half>::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 <typename T, int M_COUNT>
__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 <typename T, int M_COUNT>
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<T, M_COUNT>
<<<grid, block, 0, stream>>>(a, b_q_weight, b_qzeros, b_scales, c, size_m,
size_n, size_k, groups, zero_offset);
gemm_q4_kernel_rdna3<T, M_COUNT><<<grid, block, 0, stream>>>(
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 <typename T>
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<T, 1>(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<T, 2>(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<T, 4>(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<T, 8>(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<T, 1>(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<T, 2>(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<T, 4>(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 <typename T>
__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<T, half>::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<T, 8>(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 <typename T>
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<float>();
int* tickets_ptr = nullptr;
if constexpr (std::is_same<T, half>::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<int>(),
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<T><<<blocks, threads, 0, stream>>>(
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<half>(
vllm::gptq_rdna3::launch_gemm_q4_deterministic<half>(
(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::bf16_t>(
vllm::gptq_rdna3::launch_gemm_q4_deterministic<vllm::gptq_rdna3::bf16_t>(
(const vllm::gptq_rdna3::bf16_t*)a.data_ptr(),
(const uint32_t*)b_q_weight.data_ptr(),
(const uint32_t*)b_qzeros.data_ptr(),
csrc/rocm/q_gemm_rdna3_wmma.cu
+367 -68
View File
@@ -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 <cstdint>
@@ -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<T>::elem;
using V16 = typename WmmaNative<T>::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 <typename T>
__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 <typename T>
__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<T, half>::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 <typename T>
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<T><<<blocks, threads, 0, stream>>>(
partials.data_ptr<float>(), c, k_split, size_m, size_n);
}
template <typename T>
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<T>
<<<grid, block, 0, stream>>>(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<float>();
}
gemm_q4_wmma_kernel_16x16_1w<T><<<grid, block, 0, stream>>>(
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<T>(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<T>::elem;
using V16 = typename WmmaNative<T>::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 <typename T>
__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 <typename T>
@@ -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<T>
<<<grid, block, 0, stream>>>(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<float>();
}
gemm_q4_wmma_kernel_32x16_2w<T><<<grid, block, 0, stream>>>(
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<T>(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<T>::elem;
using V16 = typename WmmaNative<T>::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 <typename T>
__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 <typename T>
@@ -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<T>
<<<grid, block, 0, stream>>>(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<float>();
}
gemm_q4_wmma_kernel_64x16_4w<T><<<grid, block, 0, stream>>>(
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<T>(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<T>::elem;
using V16 = typename WmmaNative<T>::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 <typename T>
__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 <typename T>
@@ -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<T>
<<<grid, block, 0, stream>>>(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<float>();
}
gemm_q4_wmma_kernel_64x32_4w<T><<<grid, block, 0, stream>>>(
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<T>(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<T>::elem;
using V16 = typename WmmaNative<T>::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<T>::elem;
using V16 = typename WmmaNative<T>::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<T>::elem;
using V16 = typename WmmaNative<T>::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 <typename T>
__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 <typename T>
__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 <typename T>
__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 <typename T>
@@ -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<T><<<grid, block, 0, stream>>>(
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<T><<<grid, block, 0, stream>>>(
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<T><<<grid, block, 0, stream>>>(
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<T><<<grid, block, 0, stream>>>(
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<float>();
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<T><<<grid, block, 0, stream>>>(
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<T><<<grid, block, 0, stream>>>(
a_t, b_q_weight, b_qzeros, b_scales, c_t, rows, size_n, size_k,
groups, zero_offset, partials_ptr);
}
launch_wmma_reduce<T>(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<T>
<<<grid, block, 0, stream>>>(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<T><<<grid, block, 0, stream>>>(
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<float>();
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<T><<<grid, block, 0, stream>>>(
a_t, b_q_weight, b_qzeros, b_scales, c_t, rows, size_n, size_k, groups,
zero_offset, partials_ptr);
launch_wmma_reduce<T>(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;
csrc/rocm/qdq_4_rdna3.cuh
+60 -46
View File
@@ -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;
}
// ---------------------------------------------------------------------------
docs/features/batch_invariance.md
+3 -1
View File
@@ -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`
@@ -152,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.
docs/features/sleep_mode.md
+35
View File
@@ -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 <model> --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 <model> --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.
requirements/common.txt
+1 -1
View File
@@ -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
requirements/cuda.txt
+1 -1
View File
@@ -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
requirements/test/cpu.txt
+2 -2
View File
@@ -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
@@ -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
requirements/test/cuda.in
+2 -2
View File
@@ -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
@@ -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
requirements/test/cuda.txt
+2 -2
View File
@@ -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
@@ -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
requirements/test/nightly-torch.txt
+1 -1
View File
@@ -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
requirements/test/rocm.in
+2 -2
View File
@@ -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
@@ -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
requirements/test/rocm.txt
+2 -2
View File
@@ -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
@@ -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
requirements/test/xpu.in
+1 -1
View File
@@ -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
requirements/test/xpu.txt
+1 -1
View File
@@ -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
tests/basic_correctness/memory/cumem/test_cumem.py
+55
View File
@@ -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))
tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py
+74 -1
View File
@@ -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():
tests/distributed/test_comm_ops.py
+28 -2
View File
@@ -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)
)
tests/distributed/test_engram_dp_shard.py
+3 -1
View File
@@ -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()
tests/entrypoints/openai/chat_completion/test_chat_error.py
+63 -1
View File
@@ -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)
tests/entrypoints/openai/completion/test_completion_error.py
+71 -5
View File
@@ -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)
tests/entrypoints/openai/responses/test_serving_responses.py
+82 -1
View File
@@ -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(
tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py
+55 -7
View File
@@ -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]."""
tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py
+10 -8
View File
@@ -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)
tests/entrypoints/serve/dev/rlhf/test_weight_operation_metrics.py
+3 -26
View File
@@ -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
tests/kernels/attention/test_prefix_prefill.py
+58
View File
@@ -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)
tests/kernels/quantization/test_rdna3_w4a16_determinism.py
+585
View File
@@ -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)
tests/kernels/test_compressor_kv_cache.py
+15 -13
View File
@@ -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")
tests/models/qwen4_exp/test_ple.py
+19
View File
@@ -26,6 +26,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,
@@ -573,6 +574,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])
tests/models/registry.py
+1 -5
View File
@@ -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(
tests/multimodal/test_audio.py
+19 -13
View File
@@ -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."""
tests/v1/attention/test_mla_backends.py
+146 -2
View File
@@ -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"])
tests/v1/cudagraph/test_cudagraph_manager.py
+1
View File
@@ -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
tests/v1/determinism/test_batch_invariance.py
+20 -31
View File
@@ -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}"
tests/v1/determinism/test_lora_shrink_batch_invariant.py
+99
View File
@@ -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
tests/v1/simple_kv_offload/test_scheduler.py
+25
View File
@@ -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
# ---------------------------------------------------------------------------
tests/v1/worker/test_gpu_model_runner.py
+61
View File
@@ -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 = (
vllm/compilation/breakable_cudagraph.py
+5 -8
View File
@@ -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
vllm/compilation/cuda_graph.py
+5 -6
View File
@@ -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
vllm/compilation/cudagraph_pool.py
+45
View File
@@ -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
vllm/config/model.py
+5
View File
@@ -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",
}
@@ -376,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
vllm/config/speculative.py
+1 -1
View File
@@ -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(
vllm/config/vllm.py
+23
View File
@@ -704,6 +704,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:
@@ -2328,6 +2343,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()
vllm/device_allocator/cumem.py
+17
View File
@@ -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
vllm/device_allocator/sleep_mode_backend.py
+5 -2
View File
@@ -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:
vllm/distributed/device_communicators/cuda_communicator.py
+6
View File
@@ -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
vllm/distributed/device_communicators/custom_all_reduce.py
+5 -1
View File
@@ -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.
vllm/distributed/device_communicators/flashinfer_all_reduce.py
+39 -20
View File
@@ -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,
vllm/engine/arg_utils.py
+6
View File
@@ -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,
vllm/entrypoints/openai/chat_completion/serving.py
+2 -4
View File
@@ -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:
vllm/entrypoints/openai/completion/serving.py
+2 -3
View File
@@ -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",
vllm/entrypoints/openai/responses/context.py
+5
View File
@@ -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()
vllm/entrypoints/openai/responses/serving.py
+24 -10
View File
@@ -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(
vllm/entrypoints/openai/responses/streaming_events.py
+1 -1
View File
@@ -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",
vllm/entrypoints/openai/responses/utils.py
+22
View File
@@ -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:
vllm/entrypoints/scale_out/token_in_token_out/serving.py
+11 -17
View File
@@ -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:
vllm/entrypoints/serve/engine/serving.py
+5 -1
View File
@@ -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",
vllm/lora/ops/triton_ops/kernel_utils.py
+21 -5
View File
@@ -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")
vllm/lora/ops/triton_ops/lora_shrink_op.py
+150 -5
View File
@@ -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
vllm/model_executor/determinism/batch_invariant.py
+1 -1
View File
@@ -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
vllm/model_executor/models/lightonocr.py
+13
View File
@@ -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],
vllm/model_executor/models/longcat_flash.py
+15 -151
View File
@@ -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
vllm/model_executor/models/longcat_flash_mtp.py
+6 -7
View File
@@ -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
vllm/model_executor/models/longcat_flash_ngram.py
+6 -16
View File
@@ -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
vllm/models/deepseek_v4/xpu/xpu_sparse_decode_fp8.py
+36 -89
View File
@@ -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),
vllm/models/deepseek_v41/compressor.py
+5 -1
View File
@@ -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,
)
vllm/models/hy_v4/nvidia/attention.py
+7 -1
View File
@@ -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+)(?:\.|$)")
vllm/models/qwen4_exp/amd/model.py
+7 -3
View File
@@ -80,6 +80,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,
@@ -208,7 +212,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(
@@ -290,7 +294,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,
@@ -381,7 +385,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)
vllm/models/qwen4_exp/common/ngram_embedding.py
+6
View File
@@ -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 "
vllm/models/qwen4_exp/nvidia/model.py
+7 -3
View File
@@ -82,6 +82,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,
@@ -212,7 +216,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(
@@ -296,7 +300,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,
@@ -376,7 +380,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)
vllm/multimodal/audio.py
+1 -1
View File
@@ -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:
vllm/platforms/cpu.py
+2
View File
@@ -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]]:
vllm/platforms/interface.py
+31
View File
@@ -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(
vllm/platforms/xpu.py
+5
View File
@@ -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:
vllm/transformers_utils/config.py
+22
View File
@@ -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(
vllm/transformers_utils/model_arch_config_convertor.py
+7 -1
View File
@@ -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,
vllm/v1/attention/backends/mla/tokenspeed_mla.py
+9
View File
@@ -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
vllm/v1/attention/backends/rocm_attn.py
+3 -1
View File
@@ -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,
vllm/v1/simple_kv_offload/manager.py
+4
View File
@@ -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)
vllm/v1/worker/gpu/attn_utils.py
+1 -4
View File
@@ -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
vllm/v1/worker/gpu/cudagraph_utils.py
+11 -8
View File
@@ -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