diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 70437d0ecfd..fc7fc25ff29 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -38,6 +38,28 @@ steps: commands: - pytest -v -s kernels/core/test_minimax_reduce_rms.py +- label: Deepseek V4 Kernel Test (H100) + key: deepseek-v4-kernel-test-h100 + timeout_in_minutes: 15 + device: h100 + source_file_dependencies: + - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu + - vllm/models/deepseek_v4/common/ops/ + - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py + commands: + - pytest -v -s kernels/test_fused_deepseek_v4_*.py + +- label: Deepseek V4 Kernel Test (B200) + key: deepseek-v4-kernel-test-b200 + timeout_in_minutes: 15 + device: b200-k8s + source_file_dependencies: + - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu + - vllm/models/deepseek_v4/common/ops/ + - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py + commands: + - pytest -v -s kernels/test_fused_deepseek_v4_*.py + - label: Kernels Attention Test %N key: kernels-attention-test timeout_in_minutes: 35 diff --git a/csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu b/csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu index 589e5f7bac0..e4d432cac97 100644 --- a/csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu +++ b/csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu @@ -122,18 +122,45 @@ __device__ __forceinline__ float warpSum(float val) { // Per-slot inner pipeline // ──────────────────────────────────────────────────────────────────────────── // Shared by both kernel variants: 1 CTA per (token, head) pair vs. 1 CTA per -// token -template +// token. Templated on `kNumHeadsQPadded` so the KV-sentinel comparison and +// q_out stride fold to compile-time constants. +// +// Slot layout (per token): +// slot < num_heads_q → live-Q (RMSNorm + RoPE, +// read q_in → +// write q_out) +// num_heads_q <= slot < kNumHeadsQPadded → pad-Q (zero-fill q_out; +// v0/v1 unused) +// slot == kNumHeadsQPadded → KV (RoPE + UE8M0 quant +// + paged-cache +// insert) +template __device__ __forceinline__ void processDeepseekV4Slot( uint4 v0, uint4 v1, int const tokenIdx, int const slotIdx, int const dim_base, int const laneId, int const num_heads_q, - float const eps, scalar_t_in* __restrict__ q_inout, + float const eps, scalar_t_in* __restrict__ q_out, uint8_t* __restrict__ k_cache, int64_t const* __restrict__ slot_mapping, int64_t const* __restrict__ position_ids, float const* __restrict__ cos_sin_cache, int const cache_block_size, int const kv_block_stride) { using Converter = vllm::_typeConvert; - bool const isKV = (slotIdx == num_heads_q); + bool const isKV = (slotIdx == kNumHeadsQPadded); + bool const isPadQ = !isKV && (slotIdx >= num_heads_q); + + // ── Pad-Q branch: write 32 B of zeros and exit. ───────────────────────── + // FlashMLA reads these slots; bf16 +0.0 is bit pattern 0x0000, so a uint4 + // zero literal is correct. Matches the live-Q branch's vectorized store. + if (isPadQ) { + scalar_t_in* dst = + q_out + + (static_cast(tokenIdx) * kNumHeadsQPadded + slotIdx) * + kHeadDim + + dim_base; + uint4 const zero4 = {0u, 0u, 0u, 0u}; + *reinterpret_cast(dst) = zero4; + *reinterpret_cast(dst + 8) = zero4; + return; + } // ── Decode the bf16 → 16 fp32 registers ───────────────────────────── float elements[kElemsPerLane]; @@ -207,7 +234,7 @@ __device__ __forceinline__ void processDeepseekV4Slot( // triggering and per-iteration buffer rotation. // ═══════════════════════════════════════════════════════════════════ if (!isKV) { - // ── Q: cast back to bf16 and store. ──────────────────────────── + // ── Live-Q: cast back to bf16 and store into the padded q_out. ───── uint4 out0, out1; typename Converter::packed_hip_type* po0 = reinterpret_cast(&out0); @@ -224,8 +251,9 @@ __device__ __forceinline__ void processDeepseekV4Slot( make_float2(elements[8 + 2 * i], elements[8 + 2 * i + 1])); } scalar_t_in* dst = - q_inout + - (static_cast(tokenIdx) * num_heads_q + slotIdx) * kHeadDim + + q_out + + (static_cast(tokenIdx) * kNumHeadsQPadded + slotIdx) * + kHeadDim + dim_base; *reinterpret_cast(dst) = out0; *reinterpret_cast(dst + 8) = out1; @@ -313,20 +341,32 @@ __device__ __forceinline__ void processDeepseekV4Slot( // Kernel // ──────────────────────────────────────────────────────────────────────────── // -// Grid: 1D, gridDim.x = ceil(num_tokens_full * (num_heads_q + 1) / +// Grid: 1D, gridDim.x = ceil(num_tokens_full * (kNumHeadsQPadded + 1) / // warps_per_block) Block: blockDim.x = 256 threads (8 warps per block) Each -// warp handles one (token, head_slot) pair. head_slot < num_heads_q → -// Q branch (RMSNorm + RoPE, in place) head_slot == num_heads_q → KV -// branch (RoPE + UE8M0 quant + insert) +// warp handles one (token, head_slot) pair. +// slot < num_heads_q → live-Q branch +// (RMSNorm + RoPE, +// read q_in → write q_out) +// num_heads_q <= slot < kNumHeadsQPadded → pad-Q branch +// (zero-fill q_out) +// slot == kNumHeadsQPadded → KV branch +// (RoPE + UE8M0 quant + +// paged-cache insert) +// +// `kNumHeadsQPadded` is a template parameter (compile-time constant) so the +// divisions in the grid math and the KV-sentinel comparison fold to fast +// constant operations. The launch wrapper dispatches the runtime value to +// the matching instantiation. // // With DP padding, q/kv/position_ids can have more rows than slot_mapping. -// The Q branch covers all `num_tokens_full` rows (downstream attention uses -// them). The KV branch only inserts the first `num_tokens_insert` tokens -// (= slot_mapping length) into the paged cache. +// The live-Q and pad-Q branches cover all `num_tokens_full` rows (downstream +// attention uses them). The KV branch only inserts the first +// `num_tokens_insert` tokens (= slot_mapping length) into the paged cache. // -template +template __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel( - scalar_t_in* __restrict__ q_inout, // [N, H, 512] bf16, in place + scalar_t_in const* __restrict__ q_in, // [N, num_heads_q, 512] + scalar_t_in* __restrict__ q_out, // [N, kNumHeadsQPadded, 512] scalar_t_in const* __restrict__ kv_in, // [N, 512] bf16 uint8_t* __restrict__ k_cache, // [num_blocks, block_stride] int64_t const* __restrict__ slot_mapping, // [num_tokens_insert] i64 @@ -335,7 +375,7 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel( float const eps, int const num_tokens_full, // = q.size(0) = kv.size(0) int const num_tokens_insert, // = slot_mapping.size(0), ≤ num_tokens_full - int const num_heads_q, // H + int const num_heads_q, // live Q heads (input layout) int const cache_block_size, // tokens per paged-cache block int const kv_block_stride) { // bytes per paged-cache block #if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM) @@ -351,12 +391,13 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel( int const laneId = threadIdx.x % 32; int const globalWarpIdx = blockIdx.x * warpsPerBlock + warpId; - int const total_slots_per_token = num_heads_q + 1; - int const tokenIdx = globalWarpIdx / total_slots_per_token; - int const slotIdx = globalWarpIdx % total_slots_per_token; + constexpr int kTotalSlotsPerToken = kNumHeadsQPadded + 1; + int const tokenIdx = globalWarpIdx / kTotalSlotsPerToken; + int const slotIdx = globalWarpIdx % kTotalSlotsPerToken; if (tokenIdx >= num_tokens_full) return; - bool const isKV = (slotIdx == num_heads_q); + bool const isKV = (slotIdx == kNumHeadsQPadded); + bool const isPadQ = !isKV && (slotIdx >= num_heads_q); // KV branch: skip DP-padded tokens (no slot reserved for them). if (isKV && tokenIdx >= num_tokens_insert) return; @@ -371,22 +412,26 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel( // Dim range this lane owns within the 512-wide head. int const dim_base = laneId * kElemsPerLane; // in [0, 512) step 16 - // Two 16-byte loads per thread (8 bf16 each). Use uint4 as the vector - // type; the shared per-slot helper bitcasts to scalar_t_in packed pairs. - scalar_t_in const* src_ptr; - if (isKV) { - src_ptr = kv_in + static_cast(tokenIdx) * kHeadDim + dim_base; - } else { - int64_t const q_row_offset = - (static_cast(tokenIdx) * num_heads_q + slotIdx) * kHeadDim + - dim_base; - src_ptr = q_inout + q_row_offset; + // Load only for live-Q and KV slots; pad-Q skips the read (q_in beyond + // num_heads_q is out of bounds) and the helper zero-fills its output. + uint4 v0, v1; + if (!isPadQ) { + scalar_t_in const* src_ptr; + if (isKV) { + src_ptr = kv_in + static_cast(tokenIdx) * kHeadDim + dim_base; + } else { + int64_t const q_row_offset = + (static_cast(tokenIdx) * num_heads_q + slotIdx) * + kHeadDim + + dim_base; + src_ptr = q_in + q_row_offset; + } + v0 = *reinterpret_cast(src_ptr); + v1 = *reinterpret_cast(src_ptr + 8); } - uint4 const v0 = *reinterpret_cast(src_ptr); - uint4 const v1 = *reinterpret_cast(src_ptr + 8); - processDeepseekV4Slot( - v0, v1, tokenIdx, slotIdx, dim_base, laneId, num_heads_q, eps, q_inout, + processDeepseekV4Slot( + v0, v1, tokenIdx, slotIdx, dim_base, laneId, num_heads_q, eps, q_out, k_cache, slot_mapping, position_ids, cos_sin_cache, cache_block_size, kv_block_stride); @@ -408,9 +453,9 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel( // Q branch (RMSNorm + RoPE, in place) head_slot == num_heads_q // KV branch (RoPE + UE8M0 quant + insert) // -template +template __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid( - scalar_t_in* __restrict__ q_inout, // [N, H, 512] bf16, in place + scalar_t_in const* __restrict__ q_in, scalar_t_in* __restrict__ q_out, scalar_t_in const* __restrict__ kv_in, uint8_t* __restrict__ k_cache, int64_t const* __restrict__ slot_mapping, int64_t const* __restrict__ position_ids, @@ -435,25 +480,32 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid( #endif int const dim_base = laneId * kElemsPerLane; // in [0, 512) step 16 - int const slot_end = - (tokenIdx >= num_tokens_insert) ? num_heads_q : (num_heads_q + 1); + // Slot enumeration: live-Q + pad-Q + (KV if this token has a slot). + int const slot_end = (tokenIdx >= num_tokens_insert) + ? kNumHeadsQPadded + : (kNumHeadsQPadded + 1); - auto src_for_slot = [&](int s) -> scalar_t_in const* { - if (s == num_heads_q) { - return kv_in + static_cast(tokenIdx) * kHeadDim + dim_base; + auto load_slot = [&](int s, uint4& va, uint4& vb) { + // pad-Q slots skip the load — q_in beyond num_heads_q is OOB. + if (s >= num_heads_q && s < kNumHeadsQPadded) return; + scalar_t_in const* src; + if (s == kNumHeadsQPadded) { + src = kv_in + static_cast(tokenIdx) * kHeadDim + dim_base; + } else { + src = q_in + + (static_cast(tokenIdx) * num_heads_q + + static_cast(s)) * + kHeadDim + + dim_base; } - return q_inout + - (static_cast(tokenIdx) * num_heads_q + - static_cast(s)) * - kHeadDim + - dim_base; + va = *reinterpret_cast(src); + vb = *reinterpret_cast(src + 8); }; if (warpId < slot_end) { int curr_slot = warpId; - scalar_t_in const* src_curr = src_for_slot(curr_slot); - uint4 v0_curr = *reinterpret_cast(src_curr); - uint4 v1_curr = *reinterpret_cast(src_curr + 8); + uint4 v0_curr, v1_curr; + load_slot(curr_slot, v0_curr, v1_curr); while (curr_slot < slot_end) { int const next_slot = curr_slot + warpsPerBlock; @@ -462,14 +514,12 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid( // Prefetch src for the next slot uint4 v0_next, v1_next; if (has_next) { - scalar_t_in const* src_next = src_for_slot(next_slot); - v0_next = *reinterpret_cast(src_next); - v1_next = *reinterpret_cast(src_next + 8); + load_slot(next_slot, v0_next, v1_next); } - processDeepseekV4Slot( + processDeepseekV4Slot( v0_curr, v1_curr, tokenIdx, curr_slot, dim_base, laneId, - num_heads_q, eps, q_inout, k_cache, slot_mapping, position_ids, + num_heads_q, eps, q_out, k_cache, slot_mapping, position_ids, cos_sin_cache, cache_block_size, kv_block_stride); // ── Buffer rotation: hand the prefetched LDGs to the next iter. @@ -490,10 +540,10 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid( // ──────────────────────────────────────────────────────────────────────────── // Launch wrapper // ──────────────────────────────────────────────────────────────────────────── -template -void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert( - scalar_t_in* q_inout, scalar_t_in const* kv_in, uint8_t* k_cache, - int64_t const* slot_mapping, int64_t const* position_ids, +template +static void launchFusedDeepseekV4Templated( + scalar_t_in const* q_in, scalar_t_in* q_out, scalar_t_in const* kv_in, + uint8_t* k_cache, int64_t const* slot_mapping, int64_t const* position_ids, float const* cos_sin_cache, float const eps, int const num_tokens_full, int const num_tokens_insert, int const num_heads_q, int const cache_block_size, int const kv_block_stride, @@ -501,7 +551,7 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert( constexpr int kBlockSize = 256; constexpr int kWarpsPerBlock = kBlockSize / 32; int64_t const total_warps = - static_cast(num_tokens_full) * (num_heads_q + 1); + static_cast(num_tokens_full) * (kNumHeadsQPadded + 1); int const grid = static_cast((total_warps + kWarpsPerBlock - 1) / kWarpsPerBlock); @@ -532,46 +582,85 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert( if (num_tokens_full < NUM_TOKEN_CUTOFF) { cudaLaunchKernelEx( - &config, fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel, - q_inout, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, eps, - num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size, + &config, + fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel, + q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, + eps, num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size, kv_block_stride); } else { config.gridDim = dim3(num_tokens_full); cudaLaunchKernelEx( &config, - fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid, - q_inout, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, eps, - num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size, + fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid< + scalar_t_in, kNumHeadsQPadded>, + q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, + eps, num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size, kv_block_stride); } - #else // ROCm: use standard kernel launch syntax (no PDL/stream serialization) // clang-format off - fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel + fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel <<>>( - q_inout, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, - eps, num_tokens_full, num_tokens_insert, num_heads_q, + q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, + cos_sin_cache, eps, num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size, kv_block_stride); #endif } +// Runtime dispatch into one of the precompiled `kNumHeadsQPadded` +// instantiations. Supported padded head counts: 8, 16, 32, 64, 128. +template +void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert( + scalar_t_in const* q_in, scalar_t_in* q_out, scalar_t_in const* kv_in, + uint8_t* k_cache, int64_t const* slot_mapping, + int64_t const* position_ids, float const* cos_sin_cache, float const eps, + int const num_tokens_full, int const num_tokens_insert, + int const num_heads_q, int const num_heads_q_padded, + int const cache_block_size, int const kv_block_stride, + cudaStream_t stream) { +#define DISPATCH(N) \ + case N: \ + launchFusedDeepseekV4Templated( \ + q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, \ + cos_sin_cache, eps, num_tokens_full, num_tokens_insert, num_heads_q, \ + cache_block_size, kv_block_stride, stream); \ + return; + + switch (num_heads_q_padded) { + DISPATCH(8) + DISPATCH(16) + DISPATCH(32) + DISPATCH(64) + DISPATCH(128) + default: + TORCH_CHECK(false, + "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert: " + "unsupported num_heads_q_padded=", + num_heads_q_padded, + " (compiled instantiations: 8, 16, 32, 64, 128)."); + } +#undef DISPATCH +} + } // namespace deepseek_v4_fused_ops } // namespace vllm // ──────────────────────────────────────────────────────────────────────────── // Torch op wrapper // ──────────────────────────────────────────────────────────────────────────── -void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( - torch::Tensor& q, // [N, H, 512] bf16, in place +torch::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + torch::Tensor const& q_in, // [N, num_heads_q, 512] bf16 torch::Tensor const& kv, // [N, 512] bf16 (read-only) torch::Tensor& k_cache, // [num_blocks, block_bytes] uint8 torch::Tensor const& slot_mapping, // [N] int64 torch::Tensor const& position_ids, // [N] int64 torch::Tensor const& cos_sin_cache, // [max_pos, rope_dim] bf16 + int64_t q_head_padded, // padded Q head count for output double eps, int64_t cache_block_size) { - TORCH_CHECK(q.is_cuda() && q.is_contiguous(), "q must be contiguous CUDA"); + TORCH_CHECK(q_in.is_cuda() && q_in.is_contiguous(), + "q_in must be contiguous CUDA"); TORCH_CHECK(kv.is_cuda() && kv.is_contiguous(), "kv must be contiguous CUDA"); TORCH_CHECK(k_cache.is_cuda(), "k_cache must be CUDA"); TORCH_CHECK(slot_mapping.is_cuda() && slot_mapping.dtype() == torch::kInt64, @@ -579,9 +668,12 @@ void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( TORCH_CHECK(position_ids.is_cuda() && position_ids.dtype() == torch::kInt64, "position_ids must be int64 CUDA"); TORCH_CHECK(cos_sin_cache.is_cuda(), "cos_sin_cache must be CUDA"); - TORCH_CHECK(q.dim() == 3 && q.size(2) == 512, "q shape [N, H, 512]"); + TORCH_CHECK(q_in.dim() == 3 && q_in.size(2) == 512, + "q_in shape [N, num_heads_q, 512]"); TORCH_CHECK(kv.dim() == 2 && kv.size(1) == 512, "kv shape [N, 512]"); - TORCH_CHECK(q.dtype() == kv.dtype(), "q and kv dtype must match"); + TORCH_CHECK(q_in.dtype() == kv.dtype(), "q_in and kv dtype must match"); + TORCH_CHECK(q_head_padded >= q_in.size(1), + "q_head_padded must be >= q_in.size(1) (num_heads_q)"); TORCH_CHECK(k_cache.dtype() == torch::kUInt8, "k_cache must be uint8"); TORCH_CHECK(cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == 64, "cos_sin_cache shape [max_pos, 64]"); @@ -591,32 +683,41 @@ void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( // With DP padding, slot_mapping can be shorter than q/kv/positions. // Q-norm+RoPE runs on all q.size(0) rows (downstream attention uses them); // KV quant+insert runs only on the first slot_mapping.size(0) rows. - int const num_tokens_full = static_cast(q.size(0)); + int const num_tokens_full = static_cast(q_in.size(0)); int const num_tokens_insert = static_cast(slot_mapping.size(0)); TORCH_CHECK(static_cast(kv.size(0)) == num_tokens_full && static_cast(position_ids.size(0)) == num_tokens_full, "q/kv/position_ids row counts must match"); TORCH_CHECK(num_tokens_insert <= num_tokens_full, "slot_mapping must not exceed q row count"); - int const num_heads_q = static_cast(q.size(1)); + int const num_heads_q = static_cast(q_in.size(1)); + int const num_heads_q_padded = static_cast(q_head_padded); int const cache_block_size_i = static_cast(cache_block_size); int const kv_block_stride = static_cast(k_cache.stride(0)); - at::cuda::OptionalCUDAGuard device_guard(device_of(q)); + at::cuda::OptionalCUDAGuard device_guard(device_of(q_in)); auto stream = at::cuda::getCurrentCUDAStream(); + // Allocate the padded q output. The kernel writes every element (live + // region gets RMSNorm+RoPE; pad region gets zeros), so `empty` is safe. + torch::Tensor q_out = torch::empty( + {q_in.size(0), q_head_padded, q_in.size(2)}, q_in.options()); + VLLM_DISPATCH_HALF_TYPES( - q.scalar_type(), "fused_deepseek_v4_qnorm_rope_kv_insert", [&] { + q_in.scalar_type(), "fused_deepseek_v4_qnorm_rope_kv_insert", [&] { using qkv_scalar_t = scalar_t; vllm::deepseek_v4_fused_ops:: launchFusedDeepseekV4QNormRopeKVRopeQuantInsert( - reinterpret_cast(q.data_ptr()), + reinterpret_cast(q_in.data_ptr()), + reinterpret_cast(q_out.data_ptr()), reinterpret_cast(kv.data_ptr()), reinterpret_cast(k_cache.data_ptr()), reinterpret_cast(slot_mapping.data_ptr()), reinterpret_cast(position_ids.data_ptr()), cos_sin_cache.data_ptr(), static_cast(eps), num_tokens_full, num_tokens_insert, num_heads_q, - cache_block_size_i, kv_block_stride, stream); + num_heads_q_padded, cache_block_size_i, kv_block_stride, + stream); }); -} \ No newline at end of file + return q_out; +} diff --git a/csrc/ops.h b/csrc/ops.h index 3e2faac5b7a..3c6b2b7b9bc 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -70,10 +70,11 @@ void rms_norm(torch::Tensor& out, torch::Tensor& input, torch::Tensor& weight, void fused_add_rms_norm(torch::Tensor& input, torch::Tensor& residual, torch::Tensor& weight, double epsilon); -void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( - torch::Tensor& q, torch::Tensor const& kv, torch::Tensor& k_cache, +torch::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + torch::Tensor const& q_in, torch::Tensor const& kv, torch::Tensor& k_cache, torch::Tensor const& slot_mapping, torch::Tensor const& position_ids, - torch::Tensor const& cos_sin_cache, double eps, int64_t cache_block_size); + torch::Tensor const& cos_sin_cache, int64_t q_head_padded, double eps, + int64_t cache_block_size); void apply_repetition_penalties_(torch::Tensor& logits, const torch::Tensor& prompt_mask, diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index 78c2875644c..487ff8df9dd 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -99,9 +99,9 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { // kernel launch. ops.def( "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(" - "Tensor! q, Tensor kv, Tensor! k_cache, " + "Tensor q_in, Tensor kv, Tensor! k_cache, " "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, " - "float eps, int cache_block_size) -> ()"); + "int q_head_padded, float eps, int cache_block_size) -> Tensor"); ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert", torch::kCUDA, &fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert); diff --git a/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py b/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py index 13010540d97..a49ea498e5e 100644 --- a/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py +++ b/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py @@ -120,9 +120,19 @@ pytestmark = pytest.mark.skipif( ) -def _call_fused(q, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs): - torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( - q, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs +def _call_fused( + q_in, q_head_padded, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs +): + return torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + q_in, + kv, + k_cache, + slot_mapping, + positions, + cos_sin_cache, + q_head_padded, + eps, + bs, ) @@ -130,8 +140,23 @@ def _call_fused(q, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs) @pytest.mark.parametrize("num_tokens", [1, 4, 17, 64, 2048]) -@pytest.mark.parametrize("n_heads", [8, 64]) -def test_q_path_matches_reference(num_tokens: int, n_heads: int): +@pytest.mark.parametrize( + "n_heads,padded_heads", + [ + # Each supported padded_heads instantiation: padded (n_heads < + # padded_heads) and unpadded (n_heads == padded_heads). + (1, 8), + (8, 8), + (8, 16), + (16, 16), + (16, 32), + (32, 32), + (8, 64), + (64, 64), + (64, 128), + ], +) +def test_q_path_matches_reference(num_tokens: int, n_heads: int, padded_heads: int): torch.manual_seed(0) device = "cuda" dtype = torch.bfloat16 @@ -156,10 +181,16 @@ def test_q_path_matches_reference(num_tokens: int, n_heads: int): num_blocks, bs, HEAD_BYTES, dtype=torch.uint8, device=device ).view(num_blocks, -1) slot_mapping = torch.full((num_tokens,), -1, dtype=torch.int64, device=device) - q_fused = q.clone() - _call_fused(q_fused, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs) + q_out = _call_fused( + q, padded_heads, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs + ) - torch.testing.assert_close(q_fused, q_ref, rtol=1e-2, atol=1e-2) + torch.testing.assert_close(q_out[:, :n_heads], q_ref, rtol=1e-2, atol=1e-2) + if n_heads < padded_heads: + pad_region = q_out[:, n_heads:padded_heads] + assert pad_region.abs().max().item() == 0.0, ( + "padded head slots must be exact zero" + ) # ── Test 2: KV path round-trip byte/value parity ───────────────────────────── @@ -201,11 +232,12 @@ def test_kv_path_matches_reference(num_tokens: int, block_size: int): kv_ref, k_cache_ref, slot_mapping, block_size=block_size ) - # ── Fused path (dummy q, single head) ────────────────────────────────── + # ── Fused path (dummy q, padded to FlashMLA's min head count 64) ─────── k_cache_fused = torch.zeros_like(k_cache_ref) q_dummy = torch.zeros(num_tokens, 1, HEAD_DIM, dtype=dtype, device=device) - _call_fused( + _ = _call_fused( q_dummy, + 64, kv, k_cache_fused, slot_mapping, @@ -298,8 +330,9 @@ def test_kv_path_with_dp_padding(num_tokens: int, pad: int, block_size: int): # Fused: pass full-sized q/kv/positions, shorter slot_mapping. q_dummy = torch.zeros(total, 1, HEAD_DIM, dtype=dtype, device=device) k_cache_fused = torch.zeros_like(k_cache_ref) - _call_fused( + _ = _call_fused( q_dummy, + 64, kv, k_cache_fused, slot_mapping, @@ -316,9 +349,26 @@ def test_kv_path_with_dp_padding(num_tokens: int, pad: int, block_size: int): @pytest.mark.parametrize("num_tokens", [1, 4, 17, 2048]) -@pytest.mark.parametrize("n_heads", [8, 64]) +@pytest.mark.parametrize( + "n_heads,padded_heads", + [ + # Each supported padded_heads instantiation: padded (n_heads < + # padded_heads) and unpadded (n_heads == padded_heads). + (1, 8), + (8, 8), + (8, 16), + (16, 16), + (16, 32), + (32, 32), + (8, 64), + (64, 64), + (64, 128), + ], +) @pytest.mark.parametrize("block_size", [16, 64]) -def test_combined_q_and_kv(num_tokens: int, n_heads: int, block_size: int): +def test_combined_q_and_kv( + num_tokens: int, n_heads: int, padded_heads: int, block_size: int +): torch.manual_seed(2) device = "cuda" dtype = torch.bfloat16 @@ -345,10 +395,10 @@ def test_combined_q_and_kv(num_tokens: int, n_heads: int, block_size: int): ) # Fused single call. - q_fused = q.clone() k_cache_fused = torch.zeros_like(k_cache_ref) - _call_fused( - q_fused, + q_out = _call_fused( + q, + padded_heads, kv, k_cache_fused, slot_mapping, @@ -358,5 +408,10 @@ def test_combined_q_and_kv(num_tokens: int, n_heads: int, block_size: int): block_size, ) - torch.testing.assert_close(q_fused, q_ref, rtol=1e-2, atol=1e-2) + torch.testing.assert_close(q_out[:, :n_heads], q_ref, rtol=1e-2, atol=1e-2) + if n_heads < padded_heads: + pad_region = q_out[:, n_heads:padded_heads] + assert pad_region.abs().max().item() == 0.0, ( + "padded head slots must be exact zero" + ) torch.testing.assert_close(k_cache_fused, k_cache_ref, rtol=0, atol=0) diff --git a/tests/kernels/test_fused_inv_rope_fp8_quant.py b/tests/kernels/test_fused_inv_rope_fp8_quant.py index 563661e441c..b5b81efa772 100644 --- a/tests/kernels/test_fused_inv_rope_fp8_quant.py +++ b/tests/kernels/test_fused_inv_rope_fp8_quant.py @@ -8,7 +8,7 @@ the existing separate operations (inverse RoPE via rotate_neox + FP8 quant via per_token_group_quant_fp8). The reference faithfully reproduces the exact flow in -deepseek_v4/nvidia/ops/attention.py:295-310: +deepseek_v4/attention.py:295-310: 1. Apply inverse RoPE (NeoX style, last rope_dim=64 dims of each head) 2. Reshape [T, H, head_dim] -> [T, G, D] 3. Transpose+flatten to [G*T, D], quantize, reshape back @@ -668,7 +668,7 @@ def _unfused_inv_rope_fp8_quant( nope_dim: int = NOPE_DIM, rope_dim: int = ROPE_DIM, ) -> tuple[torch.Tensor, torch.Tensor]: - """Unfused path matching deepseek_v4/nvidia/ops/attention.py:295-310. + """Unfused path matching deepseek_v4/attention.py:295-310. Uses the production CUDA RoPE kernel + per_token_group_quant_fp8. """ diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index 1540667d1a4..3e4f401d32c 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -53,7 +53,7 @@ from vllm.model_executor.models.utils import ( maybe_prefix, ) from vllm.model_executor.utils import set_weight_attrs -from vllm.models.deepseek_v4.nvidia.ops.attention import ( +from vllm.models.deepseek_v4.attention import ( DeepseekV4Indexer, DeepseekV4MLAModules, DeepseekV4MultiHeadLatentAttentionWrapper, diff --git a/vllm/models/deepseek_v4/amd/rocm.py b/vllm/models/deepseek_v4/amd/rocm.py index 24a58a51b54..2af93fba31e 100644 --- a/vllm/models/deepseek_v4/amd/rocm.py +++ b/vllm/models/deepseek_v4/amd/rocm.py @@ -32,7 +32,7 @@ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( from vllm.v1.worker.workspace import current_workspace_manager if TYPE_CHECKING: - from vllm.models.deepseek_v4.nvidia.ops.attention import ( + from vllm.models.deepseek_v4.attention import ( DeepseekV4MLAAttention, ) @@ -592,6 +592,10 @@ class DeepseekV4ROCMAiterMLASparseImpl(DeepseekV4SparseMLAAttentionImpl): backend_cls = DeepseekV4ROCMAiterMLASparseBackend + @classmethod + def get_padded_num_q_heads(cls, num_heads: int) -> int: + return num_heads + @classmethod def forward_mqa( # type: ignore[override] cls, diff --git a/vllm/models/deepseek_v4/nvidia/ops/attention.py b/vllm/models/deepseek_v4/attention.py similarity index 95% rename from vllm/models/deepseek_v4/nvidia/ops/attention.py rename to vllm/models/deepseek_v4/attention.py index 5a983737885..d052b41fc54 100644 --- a/vllm/models/deepseek_v4/nvidia/ops/attention.py +++ b/vllm/models/deepseek_v4/attention.py @@ -156,18 +156,6 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer): self.head_dim = head_dim self.scale = scale - # FlashMLA sparse kernel only supports 64 or 128 heads; pad up to the - # next supported size. Must match DeepseekV4MLAAttention.padded_heads. - if num_heads <= 64: - self.padded_heads = 64 - elif num_heads <= 128: - self.padded_heads = 128 - else: - raise ValueError( - f"DeepseekV4 attention does not support {num_heads} heads " - "(must be <= 128)." - ) - self.q_lora_rank = q_lora_rank self.kv_lora_rank = kv_lora_rank self.window_size = window_size @@ -263,6 +251,9 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer): indexer=self.indexer, topk_indices_buffer=self.topk_indices_buffer, ) + # Mirror the inner layer's padded head count (single source of truth). + self.padded_heads = self.mla_attn.padded_heads + # Register this layer in the compilation config's static forward context # This allows the custom op to retrieve the layer during execution compilation_config = mla_modules.vllm_config.compilation_config @@ -450,7 +441,7 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer): def wq_b_kv_insert() -> torch.Tensor: q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim) - self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata) + q = self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata) return q # 3-way overlap (matches TRT-LLM PR #14142 Level 1): default runs @@ -484,7 +475,7 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer): def wq_b_kv_insert() -> torch.Tensor: q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim) - self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata) + q = self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata) return q q, _ = maybe_execute_in_parallel( @@ -497,12 +488,7 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer): else: # SWA-only layer: no compressor, no overlap. q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim) - self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata) - - # Pad q to FlashMLA-required head count (64 or 128) - if self.n_local_heads < self.padded_heads: - pad_size = self.padded_heads - self.n_local_heads - q = F.pad(q, (0, 0, 0, pad_size), value=0.0) + q = self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata) # MLA attention writes into the pre-allocated `out` buffer # ([num_tokens, padded_heads, head_dim]). @@ -516,9 +502,17 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer): attn_metadata: ( dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None ), - ) -> None: + ) -> torch.Tensor: if not isinstance(attn_metadata, dict): - return + # Profile run: kernel doesn't fire; produce a padded tensor so + # downstream FlashMLA gets the right shape. + if self.n_local_heads < self.padded_heads: + return F.pad( + q, + (0, 0, 0, self.padded_heads - self.n_local_heads), + value=0.0, + ) + return q swa_metadata = cast( "DeepseekSparseSWAMetadata | None", @@ -530,16 +524,19 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer): swa_kv_cache_2d = swa_kv_cache.view(swa_kv_cache.shape[0], -1) # Horizontally fused: - # Q side: q_head_norm (per-head RMSNorm, no weight) + GPT-J RoPE + # Q side: q_head_norm (per-head RMSNorm, no weight) + GPT-J RoPE, + # with zero-fill for the padding head slots. The kernel + # allocates and returns the padded q tensor. # KV side: GPT-J RoPE + UE8M0 FP8 quant + paged cache insert # kv is unchanged; mla_attn reads kv solely via swa_kv_cache. - torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + return torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( q, kv, swa_kv_cache_2d, swa_metadata.slot_mapping, positions.to(torch.int64), self.rotary_emb.cos_sin_cache, + self.padded_heads, self.eps, swa_metadata.block_size, ) @@ -607,9 +604,6 @@ direct_register_custom_op( class DeepseekV4MLAAttention(nn.Module, AttentionLayerBase): - # FlashMLA FP8 sparse only supports 64 or 128 heads - SUPPORTED_HEAD_COUNTS = (64, 128) - def __init__( self, num_heads: int, @@ -655,19 +649,8 @@ class DeepseekV4MLAAttention(nn.Module, AttentionLayerBase): self.aux_stream = aux_stream self.ln_events = [torch.cuda.Event(), torch.cuda.Event()] - # Determine padded head count for FlashMLA - if num_heads not in self.SUPPORTED_HEAD_COUNTS: - if num_heads < 64: - self.padded_heads = 64 - elif num_heads < 128: - self.padded_heads = 128 - else: - raise ValueError( - f"DeepseekV4MLAAttention does not support {num_heads} heads. " - f"Supported: <= 128 (will be padded to 64 or 128)" - ) - else: - self.padded_heads = num_heads + # Padded Q head count is dictated by the selected impl. + self.padded_heads = self.impl_cls.get_padded_num_q_heads(num_heads) # Store attention sink assert attn_sink is not None diff --git a/vllm/models/deepseek_v4/nvidia/flashmla.py b/vllm/models/deepseek_v4/nvidia/flashmla.py index 41630456561..5c8b08d4c12 100644 --- a/vllm/models/deepseek_v4/nvidia/flashmla.py +++ b/vllm/models/deepseek_v4/nvidia/flashmla.py @@ -28,7 +28,7 @@ from vllm.v1.attention.ops.flashmla import ( from vllm.v1.worker.workspace import current_workspace_manager if TYPE_CHECKING: - from vllm.models.deepseek_v4.nvidia.ops.attention import ( + from vllm.models.deepseek_v4.attention import ( DeepseekV4MLAAttention, ) from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata @@ -63,6 +63,18 @@ class DeepseekV4SparseMLAAttentionImpl(SparseMLAAttentionImpl[FlashMLASparseMeta ) -> None: raise NotImplementedError + @classmethod + @abstractmethod + def get_padded_num_q_heads(cls, num_heads: int) -> int: + """Q head count the backend wants q allocated at. + + The MLA wrapper allocates the q/output buffers at + ``[N, get_padded_num_q_heads(n_local_heads), head_dim]``. Must + satisfy ``result >= num_heads``. Backends with no padding constraint + return ``num_heads``. + """ + raise NotImplementedError + class DeepseekV4FlashMLASparseBackend(FlashMLASparseBackend): @staticmethod @@ -104,6 +116,16 @@ class DeepseekV4FlashMLASparseImpl(DeepseekV4SparseMLAAttentionImpl): backend_cls = DeepseekV4FlashMLASparseBackend + @classmethod + def get_padded_num_q_heads(cls, num_heads: int) -> int: + # FP8 decode kernel only supports h_q = 64 or 128. + if num_heads > 128: + raise ValueError( + f"DeepseekV4 FlashMLA does not support {num_heads} heads " + "(FP8 decode kernel requires h_q in {64, 128})." + ) + return 64 if num_heads <= 64 else 128 + @classmethod def forward_mqa( # type: ignore[override] cls, diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index 0c0cee2a01e..25f0a730fdb 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -54,7 +54,7 @@ from vllm.model_executor.models.utils import ( maybe_prefix, ) from vllm.model_executor.utils import set_weight_attrs -from vllm.models.deepseek_v4.nvidia.ops.attention import ( +from vllm.models.deepseek_v4.attention import ( DeepseekV4Indexer, DeepseekV4MLAModules, DeepseekV4MultiHeadLatentAttentionWrapper,