karylab_agents/vllm
1
0
Fork 0
forked from Karylab-cklius/vllm
Code Pull Requests 2 Actions 1 Packages Activity

[Feat][DSV4] Fuse q pad into deepseek v4 fused kernel (#43162)

Browse Source 
(cherry picked from commit 6ab6ffb428)
This commit is contained in:
Yongye Zhu
2026-05-27 00:36:50 -07:00
committed by khluu
parent 260b528b1e
commit fb83f09e8d
11 changed files with 338 additions and 150 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.buildkite/test_areas/kernels.yaml
+22
View File
@@ -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
csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
+183 -82
View File
@@ -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 <typename scalar_t_in>
// 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 <typename scalar_t_in, int kNumHeadsQPadded>
__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<scalar_t_in>;
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<int64_t>(tokenIdx) * kNumHeadsQPadded + slotIdx) *
kHeadDim +
dim_base;
uint4 const zero4 = {0u, 0u, 0u, 0u};
*reinterpret_cast<uint4*>(dst) = zero4;
*reinterpret_cast<uint4*>(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<typename Converter::packed_hip_type*>(&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<int64_t>(tokenIdx) * num_heads_q + slotIdx) * kHeadDim +
q_out +
(static_cast<int64_t>(tokenIdx) * kNumHeadsQPadded + slotIdx) *
kHeadDim +
dim_base;
*reinterpret_cast<uint4*>(dst) = out0;
*reinterpret_cast<uint4*>(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 <typename scalar_t_in>
template <typename scalar_t_in, int kNumHeadsQPadded>
__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<int64_t>(tokenIdx) * kHeadDim + dim_base;
} else {
int64_t const q_row_offset =
(static_cast<int64_t>(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<int64_t>(tokenIdx) * kHeadDim + dim_base;
} else {
int64_t const q_row_offset =
(static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) *
kHeadDim +
dim_base;
src_ptr = q_in + q_row_offset;
}
v0 = *reinterpret_cast<uint4 const*>(src_ptr);
v1 = *reinterpret_cast<uint4 const*>(src_ptr + 8);
}
uint4 const v0 = *reinterpret_cast<uint4 const*>(src_ptr);
uint4 const v1 = *reinterpret_cast<uint4 const*>(src_ptr + 8);
processDeepseekV4Slot<scalar_t_in>(
v0, v1, tokenIdx, slotIdx, dim_base, laneId, num_heads_q, eps, q_inout,
processDeepseekV4Slot<scalar_t_in, kNumHeadsQPadded>(
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 <typename scalar_t_in>
template <typename scalar_t_in, int kNumHeadsQPadded>
__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<int64_t>(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<int64_t>(tokenIdx) * kHeadDim + dim_base;
} else {
src = q_in +
(static_cast<int64_t>(tokenIdx) * num_heads_q +
static_cast<int64_t>(s)) *
kHeadDim +
dim_base;
}
return q_inout +
(static_cast<int64_t>(tokenIdx) * num_heads_q +
static_cast<int64_t>(s)) *
kHeadDim +
dim_base;
va = *reinterpret_cast<uint4 const*>(src);
vb = *reinterpret_cast<uint4 const*>(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<uint4 const*>(src_curr);
uint4 v1_curr = *reinterpret_cast<uint4 const*>(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<uint4 const*>(src_next);
v1_next = *reinterpret_cast<uint4 const*>(src_next + 8);
load_slot(next_slot, v0_next, v1_next);
}
processDeepseekV4Slot<scalar_t_in>(
processDeepseekV4Slot<scalar_t_in, kNumHeadsQPadded>(
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 <typename scalar_t_in>
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 <typename scalar_t_in, int kNumHeadsQPadded>
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<int64_t>(num_tokens_full) * (num_heads_q + 1);
static_cast<int64_t>(num_tokens_full) * (kNumHeadsQPadded + 1);
int const grid =
static_cast<int>((total_warps + kWarpsPerBlock - 1) / kWarpsPerBlock);
@@ -532,46 +582,85 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
if (num_tokens_full < NUM_TOKEN_CUTOFF) {
cudaLaunchKernelEx(
&config, fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<scalar_t_in>,
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<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 {
config.gridDim = dim3(num_tokens_full);
cudaLaunchKernelEx(
&config,
fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid<scalar_t_in>,
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<scalar_t_in>
fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<scalar_t_in, kNumHeadsQPadded>
<<<grid, kBlockSize, 0, stream>>>(
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 <typename scalar_t_in>
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<scalar_t_in, N>( \
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<int>(q.size(0));
int const num_tokens_full = static_cast<int>(q_in.size(0));
int const num_tokens_insert = static_cast<int>(slot_mapping.size(0));
TORCH_CHECK(static_cast<int>(kv.size(0)) == num_tokens_full &&
static_cast<int>(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<int>(q.size(1));
int const num_heads_q = static_cast<int>(q_in.size(1));
int const num_heads_q_padded = static_cast<int>(q_head_padded);
int const cache_block_size_i = static_cast<int>(cache_block_size);
int const kv_block_stride = static_cast<int>(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<qkv_scalar_t>(
reinterpret_cast<qkv_scalar_t*>(q.data_ptr()),
reinterpret_cast<qkv_scalar_t const*>(q_in.data_ptr()),
reinterpret_cast<qkv_scalar_t*>(q_out.data_ptr()),
reinterpret_cast<qkv_scalar_t const*>(kv.data_ptr()),
reinterpret_cast<uint8_t*>(k_cache.data_ptr()),
reinterpret_cast<int64_t const*>(slot_mapping.data_ptr()),
reinterpret_cast<int64_t const*>(position_ids.data_ptr()),
cos_sin_cache.data_ptr<float>(), static_cast<float>(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);
});
}
return q_out;
}
csrc/ops.h
+4 -3
View File
@@ -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,
csrc/torch_bindings.cpp
+2 -2
View File
@@ -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);
tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
+72 -17
View File
@@ -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)
tests/kernels/test_fused_inv_rope_fp8_quant.py
+2 -2
View File
@@ -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.
"""
vllm/models/deepseek_v4/amd/model.py
+1 -1
View File
@@ -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,
vllm/models/deepseek_v4/amd/rocm.py
+5 -1
View File
@@ -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,
vllm/models/deepseek_v4/nvidia/ops/attention.py → vllm/models/deepseek_v4/attention.py
+23 -40
View File
@@ -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
vllm/models/deepseek_v4/nvidia/flashmla.py
+23 -1
View File
@@ -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,
vllm/models/deepseek_v4/nvidia/model.py
+1 -1
View File
@@ -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,