@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 6ab6ffb428
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 2, 19:30 UTC

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

Yongye Zhu committed May 26, 2026 at 11:12 UTC 6ab6ffb428be5ed3c5d6c3cc029c5afe7ff59eee
11 files changed +339 -151
.buildkite/test_areas/kernels.yaml
+22
@@ -38,6 +38,28 @@ steps:
38 commands:
39 - pytest -v -s kernels/core/test_minimax_reduce_rms.py
40
41 +- label: Deepseek V4 Kernel Test (H100)
42 + key: deepseek-v4-kernel-test-h100
43 + timeout_in_minutes: 15
44 + device: h100
45 + source_file_dependencies:
46 + - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
47 + - vllm/models/deepseek_v4/common/ops/
48 + - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
49 + commands:
50 + - pytest -v -s kernels/test_fused_deepseek_v4_*.py
51 +
52 +- label: Deepseek V4 Kernel Test (B200)
53 + key: deepseek-v4-kernel-test-b200
54 + timeout_in_minutes: 15
55 + device: b200-k8s
56 + source_file_dependencies:
57 + - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
58 + - vllm/models/deepseek_v4/common/ops/
59 + - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
60 + commands:
61 + - pytest -v -s kernels/test_fused_deepseek_v4_*.py
62 +
63 - label: Kernels Attention Test %N
64 key: kernels-attention-test
65 timeout_in_minutes: 35
csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
+184 -83
@@ -122,18 +122,45 @@ __device__ __forceinline__ float warpSum(float val) {
122 // Per-slot inner pipeline
123 // ────────────────────────────────────────────────────────────────────────────
124 // Shared by both kernel variants: 1 CTA per (token, head) pair vs. 1 CTA per
125 -// token
126 -template <typename scalar_t_in>
125 +// token. Templated on `kNumHeadsQPadded` so the KV-sentinel comparison and
126 +// q_out stride fold to compile-time constants.
127 +//
128 +// Slot layout (per token):
129 +// slot < num_heads_q → live-Q (RMSNorm + RoPE,
130 +// read q_in →
131 +// write q_out)
132 +// num_heads_q <= slot < kNumHeadsQPadded → pad-Q (zero-fill q_out;
133 +// v0/v1 unused)
134 +// slot == kNumHeadsQPadded → KV (RoPE + UE8M0 quant
135 +// + paged-cache
136 +// insert)
137 +template <typename scalar_t_in, int kNumHeadsQPadded>
138 __device__ __forceinline__ void processDeepseekV4Slot(
139 uint4 v0, uint4 v1, int const tokenIdx, int const slotIdx,
140 int const dim_base, int const laneId, int const num_heads_q,
130 - float const eps, scalar_t_in* __restrict__ q_inout,
141 + float const eps, scalar_t_in* __restrict__ q_out,
142 uint8_t* __restrict__ k_cache, int64_t const* __restrict__ slot_mapping,
143 int64_t const* __restrict__ position_ids,
144 float const* __restrict__ cos_sin_cache, int const cache_block_size,
145 int const kv_block_stride) {
146 using Converter = vllm::_typeConvert<scalar_t_in>;
136 - bool const isKV = (slotIdx == num_heads_q);
147 + bool const isKV = (slotIdx == kNumHeadsQPadded);
148 + bool const isPadQ = !isKV && (slotIdx >= num_heads_q);
149 +
150 + // ── Pad-Q branch: write 32 B of zeros and exit. ─────────────────────────
151 + // FlashMLA reads these slots; bf16 +0.0 is bit pattern 0x0000, so a uint4
152 + // zero literal is correct. Matches the live-Q branch's vectorized store.
153 + if (isPadQ) {
154 + scalar_t_in* dst =
155 + q_out +
156 + (static_cast<int64_t>(tokenIdx) * kNumHeadsQPadded + slotIdx) *
157 + kHeadDim +
158 + dim_base;
159 + uint4 const zero4 = {0u, 0u, 0u, 0u};
160 + *reinterpret_cast<uint4*>(dst) = zero4;
161 + *reinterpret_cast<uint4*>(dst + 8) = zero4;
162 + return;
163 + }
164
165 // ── Decode the bf16 → 16 fp32 registers ─────────────────────────────
166 float elements[kElemsPerLane];
@@ -207,7 +234,7 @@ __device__ __forceinline__ void processDeepseekV4Slot(
234 // triggering and per-iteration buffer rotation.
235 // ═══════════════════════════════════════════════════════════════════
236 if (!isKV) {
210 - // ── Q: cast back to bf16 and store. ────────────────────────────
237 + // ── Live-Q: cast back to bf16 and store into the padded q_out. ─────
238 uint4 out0, out1;
239 typename Converter::packed_hip_type* po0 =
240 reinterpret_cast<typename Converter::packed_hip_type*>(&out0);
@@ -224,8 +251,9 @@ __device__ __forceinline__ void processDeepseekV4Slot(
251 make_float2(elements[8 + 2 * i], elements[8 + 2 * i + 1]));
252 }
253 scalar_t_in* dst =
227 - q_inout +
228 - (static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) * kHeadDim +
254 + q_out +
255 + (static_cast<int64_t>(tokenIdx) * kNumHeadsQPadded + slotIdx) *
256 + kHeadDim +
257 dim_base;
258 *reinterpret_cast<uint4*>(dst) = out0;
259 *reinterpret_cast<uint4*>(dst + 8) = out1;
@@ -313,20 +341,32 @@ __device__ __forceinline__ void processDeepseekV4Slot(
341 // Kernel
342 // ────────────────────────────────────────────────────────────────────────────
343 //
316 -// Grid: 1D, gridDim.x = ceil(num_tokens_full * (num_heads_q + 1) /
344 +// Grid: 1D, gridDim.x = ceil(num_tokens_full * (kNumHeadsQPadded + 1) /
345 // warps_per_block) Block: blockDim.x = 256 threads (8 warps per block) Each
318 -// warp handles one (token, head_slot) pair. head_slot < num_heads_q →
319 -// Q branch (RMSNorm + RoPE, in place) head_slot == num_heads_q → KV
320 -// branch (RoPE + UE8M0 quant + insert)
346 +// warp handles one (token, head_slot) pair.
347 +// slot < num_heads_q → live-Q branch
348 +// (RMSNorm + RoPE,
349 +// read q_in → write q_out)
350 +// num_heads_q <= slot < kNumHeadsQPadded → pad-Q branch
351 +// (zero-fill q_out)
352 +// slot == kNumHeadsQPadded → KV branch
353 +// (RoPE + UE8M0 quant +
354 +// paged-cache insert)
355 +//
356 +// `kNumHeadsQPadded` is a template parameter (compile-time constant) so the
357 +// divisions in the grid math and the KV-sentinel comparison fold to fast
358 +// constant operations. The launch wrapper dispatches the runtime value to
359 +// the matching instantiation.
360 //
361 // With DP padding, q/kv/position_ids can have more rows than slot_mapping.
323 -// The Q branch covers all `num_tokens_full` rows (downstream attention uses
324 -// them). The KV branch only inserts the first `num_tokens_insert` tokens
325 -// (= slot_mapping length) into the paged cache.
362 +// The live-Q and pad-Q branches cover all `num_tokens_full` rows (downstream
363 +// attention uses them). The KV branch only inserts the first
364 +// `num_tokens_insert` tokens (= slot_mapping length) into the paged cache.
365 //
327 -template <typename scalar_t_in>
366 +template <typename scalar_t_in, int kNumHeadsQPadded>
367 __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
329 - scalar_t_in* __restrict__ q_inout, // [N, H, 512] bf16, in place
368 + scalar_t_in const* __restrict__ q_in, // [N, num_heads_q, 512]
369 + scalar_t_in* __restrict__ q_out, // [N, kNumHeadsQPadded, 512]
370 scalar_t_in const* __restrict__ kv_in, // [N, 512] bf16
371 uint8_t* __restrict__ k_cache, // [num_blocks, block_stride]
372 int64_t const* __restrict__ slot_mapping, // [num_tokens_insert] i64
@@ -335,7 +375,7 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
375 float const eps,
376 int const num_tokens_full, // = q.size(0) = kv.size(0)
377 int const num_tokens_insert, // = slot_mapping.size(0), ≤ num_tokens_full
338 - int const num_heads_q, // H
378 + int const num_heads_q, // live Q heads (input layout)
379 int const cache_block_size, // tokens per paged-cache block
380 int const kv_block_stride) { // bytes per paged-cache block
381 #if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
@@ -351,12 +391,13 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
391 int const laneId = threadIdx.x % 32;
392 int const globalWarpIdx = blockIdx.x * warpsPerBlock + warpId;
393
354 - int const total_slots_per_token = num_heads_q + 1;
355 - int const tokenIdx = globalWarpIdx / total_slots_per_token;
356 - int const slotIdx = globalWarpIdx % total_slots_per_token;
394 + constexpr int kTotalSlotsPerToken = kNumHeadsQPadded + 1;
395 + int const tokenIdx = globalWarpIdx / kTotalSlotsPerToken;
396 + int const slotIdx = globalWarpIdx % kTotalSlotsPerToken;
397 if (tokenIdx >= num_tokens_full) return;
398
359 - bool const isKV = (slotIdx == num_heads_q);
399 + bool const isKV = (slotIdx == kNumHeadsQPadded);
400 + bool const isPadQ = !isKV && (slotIdx >= num_heads_q);
401 // KV branch: skip DP-padded tokens (no slot reserved for them).
402 if (isKV && tokenIdx >= num_tokens_insert) return;
403
@@ -371,22 +412,26 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
412 // Dim range this lane owns within the 512-wide head.
413 int const dim_base = laneId * kElemsPerLane; // in [0, 512) step 16
414
374 - // Two 16-byte loads per thread (8 bf16 each). Use uint4 as the vector
375 - // type; the shared per-slot helper bitcasts to scalar_t_in packed pairs.
376 - scalar_t_in const* src_ptr;
377 - if (isKV) {
378 - src_ptr = kv_in + static_cast<int64_t>(tokenIdx) * kHeadDim + dim_base;
379 - } else {
380 - int64_t const q_row_offset =
381 - (static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) * kHeadDim +
382 - dim_base;
383 - src_ptr = q_inout + q_row_offset;
415 + // Load only for live-Q and KV slots; pad-Q skips the read (q_in beyond
416 + // num_heads_q is out of bounds) and the helper zero-fills its output.
417 + uint4 v0, v1;
418 + if (!isPadQ) {
419 + scalar_t_in const* src_ptr;
420 + if (isKV) {
421 + src_ptr = kv_in + static_cast<int64_t>(tokenIdx) * kHeadDim + dim_base;
422 + } else {
423 + int64_t const q_row_offset =
424 + (static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) *
425 + kHeadDim +
426 + dim_base;
427 + src_ptr = q_in + q_row_offset;
428 + }
429 + v0 = *reinterpret_cast<uint4 const*>(src_ptr);
430 + v1 = *reinterpret_cast<uint4 const*>(src_ptr + 8);
431 }
385 - uint4 const v0 = *reinterpret_cast<uint4 const*>(src_ptr);
386 - uint4 const v1 = *reinterpret_cast<uint4 const*>(src_ptr + 8);
432
388 - processDeepseekV4Slot<scalar_t_in>(
389 - v0, v1, tokenIdx, slotIdx, dim_base, laneId, num_heads_q, eps, q_inout,
433 + processDeepseekV4Slot<scalar_t_in, kNumHeadsQPadded>(
434 + v0, v1, tokenIdx, slotIdx, dim_base, laneId, num_heads_q, eps, q_out,
435 k_cache, slot_mapping, position_ids, cos_sin_cache, cache_block_size,
436 kv_block_stride);
437
@@ -408,9 +453,9 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
453 // Q branch (RMSNorm + RoPE, in place) head_slot == num_heads_q
454 // KV branch (RoPE + UE8M0 quant + insert)
455 //
411 -template <typename scalar_t_in>
456 +template <typename scalar_t_in, int kNumHeadsQPadded>
457 __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid(
413 - scalar_t_in* __restrict__ q_inout, // [N, H, 512] bf16, in place
458 + scalar_t_in const* __restrict__ q_in, scalar_t_in* __restrict__ q_out,
459 scalar_t_in const* __restrict__ kv_in, uint8_t* __restrict__ k_cache,
460 int64_t const* __restrict__ slot_mapping,
461 int64_t const* __restrict__ position_ids,
@@ -435,25 +480,32 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid(
480 #endif
481
482 int const dim_base = laneId * kElemsPerLane; // in [0, 512) step 16
438 - int const slot_end =
439 - (tokenIdx >= num_tokens_insert) ? num_heads_q : (num_heads_q + 1);
440 -
441 - auto src_for_slot = [&](int s) -> scalar_t_in const* {
442 - if (s == num_heads_q) {
443 - return kv_in + static_cast<int64_t>(tokenIdx) * kHeadDim + dim_base;
483 + // Slot enumeration: live-Q + pad-Q + (KV if this token has a slot).
484 + int const slot_end = (tokenIdx >= num_tokens_insert)
485 + ? kNumHeadsQPadded
486 + : (kNumHeadsQPadded + 1);
487 +
488 + auto load_slot = [&](int s, uint4& va, uint4& vb) {
489 + // pad-Q slots skip the load — q_in beyond num_heads_q is OOB.
490 + if (s >= num_heads_q && s < kNumHeadsQPadded) return;
491 + scalar_t_in const* src;
492 + if (s == kNumHeadsQPadded) {
493 + src = kv_in + static_cast<int64_t>(tokenIdx) * kHeadDim + dim_base;
494 + } else {
495 + src = q_in +
496 + (static_cast<int64_t>(tokenIdx) * num_heads_q +
497 + static_cast<int64_t>(s)) *
498 + kHeadDim +
499 + dim_base;
500 }
445 - return q_inout +
446 - (static_cast<int64_t>(tokenIdx) * num_heads_q +
447 - static_cast<int64_t>(s)) *
448 - kHeadDim +
449 - dim_base;
501 + va = *reinterpret_cast<uint4 const*>(src);
502 + vb = *reinterpret_cast<uint4 const*>(src + 8);
503 };
504
505 if (warpId < slot_end) {
506 int curr_slot = warpId;
454 - scalar_t_in const* src_curr = src_for_slot(curr_slot);
455 - uint4 v0_curr = *reinterpret_cast<uint4 const*>(src_curr);
456 - uint4 v1_curr = *reinterpret_cast<uint4 const*>(src_curr + 8);
507 + uint4 v0_curr, v1_curr;
508 + load_slot(curr_slot, v0_curr, v1_curr);
509
510 while (curr_slot < slot_end) {
511 int const next_slot = curr_slot + warpsPerBlock;
@@ -462,14 +514,12 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid(
514 // Prefetch src for the next slot
515 uint4 v0_next, v1_next;
516 if (has_next) {
465 - scalar_t_in const* src_next = src_for_slot(next_slot);
466 - v0_next = *reinterpret_cast<uint4 const*>(src_next);
467 - v1_next = *reinterpret_cast<uint4 const*>(src_next + 8);
517 + load_slot(next_slot, v0_next, v1_next);
518 }
519
470 - processDeepseekV4Slot<scalar_t_in>(
520 + processDeepseekV4Slot<scalar_t_in, kNumHeadsQPadded>(
521 v0_curr, v1_curr, tokenIdx, curr_slot, dim_base, laneId,
472 - num_heads_q, eps, q_inout, k_cache, slot_mapping, position_ids,
522 + num_heads_q, eps, q_out, k_cache, slot_mapping, position_ids,
523 cos_sin_cache, cache_block_size, kv_block_stride);
524
525 // ── Buffer rotation: hand the prefetched LDGs to the next iter.
@@ -490,10 +540,10 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid(
540 // ────────────────────────────────────────────────────────────────────────────
541 // Launch wrapper
542 // ────────────────────────────────────────────────────────────────────────────
493 -template <typename scalar_t_in>
494 -void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
495 - scalar_t_in* q_inout, scalar_t_in const* kv_in, uint8_t* k_cache,
496 - int64_t const* slot_mapping, int64_t const* position_ids,
543 +template <typename scalar_t_in, int kNumHeadsQPadded>
544 +static void launchFusedDeepseekV4Templated(
545 + scalar_t_in const* q_in, scalar_t_in* q_out, scalar_t_in const* kv_in,
546 + uint8_t* k_cache, int64_t const* slot_mapping, int64_t const* position_ids,
547 float const* cos_sin_cache, float const eps, int const num_tokens_full,
548 int const num_tokens_insert, int const num_heads_q,
549 int const cache_block_size, int const kv_block_stride,
@@ -501,7 +551,7 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
551 constexpr int kBlockSize = 256;
552 constexpr int kWarpsPerBlock = kBlockSize / 32;
553 int64_t const total_warps =
504 - static_cast<int64_t>(num_tokens_full) * (num_heads_q + 1);
554 + static_cast<int64_t>(num_tokens_full) * (kNumHeadsQPadded + 1);
555 int const grid =
556 static_cast<int>((total_warps + kWarpsPerBlock - 1) / kWarpsPerBlock);
557
@@ -532,46 +582,85 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
582
583 if (num_tokens_full < NUM_TOKEN_CUTOFF) {
584 cudaLaunchKernelEx(
535 - &config, fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<scalar_t_in>,
536 - q_inout, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, eps,
537 - num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
585 + &config,
586 + fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<scalar_t_in,
587 + kNumHeadsQPadded>,
588 + q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache,
589 + eps, num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
590 kv_block_stride);
591 } else {
592 config.gridDim = dim3(num_tokens_full);
593 cudaLaunchKernelEx(
594 &config,
543 - fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid<scalar_t_in>,
544 - q_inout, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, eps,
545 - num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
595 + fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid<
596 + scalar_t_in, kNumHeadsQPadded>,
597 + q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache,
598 + eps, num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
599 kv_block_stride);
600 }
548 -
601 #else
602 // ROCm: use standard kernel launch syntax (no PDL/stream serialization)
603 // clang-format off
552 - fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<scalar_t_in>
604 + fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<scalar_t_in, kNumHeadsQPadded>
605 <<<grid, kBlockSize, 0, stream>>>(
554 - q_inout, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache,
555 - eps, num_tokens_full, num_tokens_insert, num_heads_q,
606 + q_in, q_out, kv_in, k_cache, slot_mapping, position_ids,
607 + cos_sin_cache, eps, num_tokens_full, num_tokens_insert, num_heads_q,
608 cache_block_size, kv_block_stride);
609 #endif
610 }
611
612 +// Runtime dispatch into one of the precompiled `kNumHeadsQPadded`
613 +// instantiations. Supported padded head counts: 8, 16, 32, 64, 128.
614 +template <typename scalar_t_in>
615 +void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
616 + scalar_t_in const* q_in, scalar_t_in* q_out, scalar_t_in const* kv_in,
617 + uint8_t* k_cache, int64_t const* slot_mapping,
618 + int64_t const* position_ids, float const* cos_sin_cache, float const eps,
619 + int const num_tokens_full, int const num_tokens_insert,
620 + int const num_heads_q, int const num_heads_q_padded,
621 + int const cache_block_size, int const kv_block_stride,
622 + cudaStream_t stream) {
623 +#define DISPATCH(N) \
624 + case N: \
625 + launchFusedDeepseekV4Templated<scalar_t_in, N>( \
626 + q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, \
627 + cos_sin_cache, eps, num_tokens_full, num_tokens_insert, num_heads_q, \
628 + cache_block_size, kv_block_stride, stream); \
629 + return;
630 +
631 + switch (num_heads_q_padded) {
632 + DISPATCH(8)
633 + DISPATCH(16)
634 + DISPATCH(32)
635 + DISPATCH(64)
636 + DISPATCH(128)
637 + default:
638 + TORCH_CHECK(false,
639 + "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert: "
640 + "unsupported num_heads_q_padded=",
641 + num_heads_q_padded,
642 + " (compiled instantiations: 8, 16, 32, 64, 128).");
643 + }
644 +#undef DISPATCH
645 +}
646 +
647 } // namespace deepseek_v4_fused_ops
648 } // namespace vllm
649
650 // ────────────────────────────────────────────────────────────────────────────
651 // Torch op wrapper
652 // ────────────────────────────────────────────────────────────────────────────
566 -void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
567 - torch::Tensor& q, // [N, H, 512] bf16, in place
653 +torch::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
654 + torch::Tensor const& q_in, // [N, num_heads_q, 512] bf16
655 torch::Tensor const& kv, // [N, 512] bf16 (read-only)
656 torch::Tensor& k_cache, // [num_blocks, block_bytes] uint8
657 torch::Tensor const& slot_mapping, // [N] int64
658 torch::Tensor const& position_ids, // [N] int64
659 torch::Tensor const& cos_sin_cache, // [max_pos, rope_dim] bf16
660 + int64_t q_head_padded, // padded Q head count for output
661 double eps, int64_t cache_block_size) {
574 - TORCH_CHECK(q.is_cuda() && q.is_contiguous(), "q must be contiguous CUDA");
662 + TORCH_CHECK(q_in.is_cuda() && q_in.is_contiguous(),
663 + "q_in must be contiguous CUDA");
664 TORCH_CHECK(kv.is_cuda() && kv.is_contiguous(), "kv must be contiguous CUDA");
665 TORCH_CHECK(k_cache.is_cuda(), "k_cache must be CUDA");
666 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(
668 TORCH_CHECK(position_ids.is_cuda() && position_ids.dtype() == torch::kInt64,
669 "position_ids must be int64 CUDA");
670 TORCH_CHECK(cos_sin_cache.is_cuda(), "cos_sin_cache must be CUDA");
582 - TORCH_CHECK(q.dim() == 3 && q.size(2) == 512, "q shape [N, H, 512]");
671 + TORCH_CHECK(q_in.dim() == 3 && q_in.size(2) == 512,
672 + "q_in shape [N, num_heads_q, 512]");
673 TORCH_CHECK(kv.dim() == 2 && kv.size(1) == 512, "kv shape [N, 512]");
584 - TORCH_CHECK(q.dtype() == kv.dtype(), "q and kv dtype must match");
674 + TORCH_CHECK(q_in.dtype() == kv.dtype(), "q_in and kv dtype must match");
675 + TORCH_CHECK(q_head_padded >= q_in.size(1),
676 + "q_head_padded must be >= q_in.size(1) (num_heads_q)");
677 TORCH_CHECK(k_cache.dtype() == torch::kUInt8, "k_cache must be uint8");
678 TORCH_CHECK(cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == 64,
679 "cos_sin_cache shape [max_pos, 64]");
@@ -591,32 +683,41 @@ void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
683 // With DP padding, slot_mapping can be shorter than q/kv/positions.
684 // Q-norm+RoPE runs on all q.size(0) rows (downstream attention uses them);
685 // KV quant+insert runs only on the first slot_mapping.size(0) rows.
594 - int const num_tokens_full = static_cast<int>(q.size(0));
686 + int const num_tokens_full = static_cast<int>(q_in.size(0));
687 int const num_tokens_insert = static_cast<int>(slot_mapping.size(0));
688 TORCH_CHECK(static_cast<int>(kv.size(0)) == num_tokens_full &&
689 static_cast<int>(position_ids.size(0)) == num_tokens_full,
690 "q/kv/position_ids row counts must match");
691 TORCH_CHECK(num_tokens_insert <= num_tokens_full,
692 "slot_mapping must not exceed q row count");
601 - int const num_heads_q = static_cast<int>(q.size(1));
693 + int const num_heads_q = static_cast<int>(q_in.size(1));
694 + int const num_heads_q_padded = static_cast<int>(q_head_padded);
695 int const cache_block_size_i = static_cast<int>(cache_block_size);
696 int const kv_block_stride = static_cast<int>(k_cache.stride(0));
697
605 - at::cuda::OptionalCUDAGuard device_guard(device_of(q));
698 + at::cuda::OptionalCUDAGuard device_guard(device_of(q_in));
699 auto stream = at::cuda::getCurrentCUDAStream();
700
701 + // Allocate the padded q output. The kernel writes every element (live
702 + // region gets RMSNorm+RoPE; pad region gets zeros), so `empty` is safe.
703 + torch::Tensor q_out = torch::empty(
704 + {q_in.size(0), q_head_padded, q_in.size(2)}, q_in.options());
705 +
706 VLLM_DISPATCH_HALF_TYPES(
609 - q.scalar_type(), "fused_deepseek_v4_qnorm_rope_kv_insert", [&] {
707 + q_in.scalar_type(), "fused_deepseek_v4_qnorm_rope_kv_insert", [&] {
708 using qkv_scalar_t = scalar_t;
709 vllm::deepseek_v4_fused_ops::
710 launchFusedDeepseekV4QNormRopeKVRopeQuantInsert<qkv_scalar_t>(
613 - reinterpret_cast<qkv_scalar_t*>(q.data_ptr()),
711 + reinterpret_cast<qkv_scalar_t const*>(q_in.data_ptr()),
712 + reinterpret_cast<qkv_scalar_t*>(q_out.data_ptr()),
713 reinterpret_cast<qkv_scalar_t const*>(kv.data_ptr()),
714 reinterpret_cast<uint8_t*>(k_cache.data_ptr()),
715 reinterpret_cast<int64_t const*>(slot_mapping.data_ptr()),
716 reinterpret_cast<int64_t const*>(position_ids.data_ptr()),
717 cos_sin_cache.data_ptr<float>(), static_cast<float>(eps),
718 num_tokens_full, num_tokens_insert, num_heads_q,
620 - cache_block_size_i, kv_block_stride, stream);
719 + num_heads_q_padded, cache_block_size_i, kv_block_stride,
720 + stream);
721 });
622 -}
\ No newline at end of file
722 + return q_out;
723 +}
csrc/ops.h
+4 -3
@@ -70,10 +70,11 @@ void rms_norm(torch::Tensor& out, torch::Tensor& input, torch::Tensor& weight,
70 void fused_add_rms_norm(torch::Tensor& input, torch::Tensor& residual,
71 torch::Tensor& weight, double epsilon);
72
73 -void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
74 - torch::Tensor& q, torch::Tensor const& kv, torch::Tensor& k_cache,
73 +torch::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
74 + torch::Tensor const& q_in, torch::Tensor const& kv, torch::Tensor& k_cache,
75 torch::Tensor const& slot_mapping, torch::Tensor const& position_ids,
76 - torch::Tensor const& cos_sin_cache, double eps, int64_t cache_block_size);
76 + torch::Tensor const& cos_sin_cache, int64_t q_head_padded, double eps,
77 + int64_t cache_block_size);
78
79 void apply_repetition_penalties_(torch::Tensor& logits,
80 const torch::Tensor& prompt_mask,
csrc/torch_bindings.cpp
+2 -2
@@ -99,9 +99,9 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
99 // kernel launch.
100 ops.def(
101 "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert("
102 - "Tensor! q, Tensor kv, Tensor! k_cache, "
102 + "Tensor q_in, Tensor kv, Tensor! k_cache, "
103 "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, "
104 - "float eps, int cache_block_size) -> ()");
104 + "int q_head_padded, float eps, int cache_block_size) -> Tensor");
105 ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert", torch::kCUDA,
106 &fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert);
107
tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
+72 -17
@@ -120,9 +120,19 @@ pytestmark = pytest.mark.skipif(
120 )
121
122
123 -def _call_fused(q, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs):
124 - torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
125 - q, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs
123 +def _call_fused(
124 + q_in, q_head_padded, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs
125 +):
126 + return torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
127 + q_in,
128 + kv,
129 + k_cache,
130 + slot_mapping,
131 + positions,
132 + cos_sin_cache,
133 + q_head_padded,
134 + eps,
135 + bs,
136 )
137
138
@@ -130,8 +140,23 @@ def _call_fused(q, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs)
140
141
142 @pytest.mark.parametrize("num_tokens", [1, 4, 17, 64, 2048])
133 -@pytest.mark.parametrize("n_heads", [8, 64])
134 -def test_q_path_matches_reference(num_tokens: int, n_heads: int):
143 +@pytest.mark.parametrize(
144 + "n_heads,padded_heads",
145 + [
146 + # Each supported padded_heads instantiation: padded (n_heads <
147 + # padded_heads) and unpadded (n_heads == padded_heads).
148 + (1, 8),
149 + (8, 8),
150 + (8, 16),
151 + (16, 16),
152 + (16, 32),
153 + (32, 32),
154 + (8, 64),
155 + (64, 64),
156 + (64, 128),
157 + ],
158 +)
159 +def test_q_path_matches_reference(num_tokens: int, n_heads: int, padded_heads: int):
160 torch.manual_seed(0)
161 device = "cuda"
162 dtype = torch.bfloat16
@@ -156,10 +181,16 @@ def test_q_path_matches_reference(num_tokens: int, n_heads: int):
181 num_blocks, bs, HEAD_BYTES, dtype=torch.uint8, device=device
182 ).view(num_blocks, -1)
183 slot_mapping = torch.full((num_tokens,), -1, dtype=torch.int64, device=device)
159 - q_fused = q.clone()
160 - _call_fused(q_fused, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs)
184 + q_out = _call_fused(
185 + q, padded_heads, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs
186 + )
187
162 - torch.testing.assert_close(q_fused, q_ref, rtol=1e-2, atol=1e-2)
188 + torch.testing.assert_close(q_out[:, :n_heads], q_ref, rtol=1e-2, atol=1e-2)
189 + if n_heads < padded_heads:
190 + pad_region = q_out[:, n_heads:padded_heads]
191 + assert pad_region.abs().max().item() == 0.0, (
192 + "padded head slots must be exact zero"
193 + )
194
195
196 # ── 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):
232 kv_ref, k_cache_ref, slot_mapping, block_size=block_size
233 )
234
204 - # ── Fused path (dummy q, single head) ──────────────────────────────────
235 + # ── Fused path (dummy q, padded to FlashMLA's min head count 64) ───────
236 k_cache_fused = torch.zeros_like(k_cache_ref)
237 q_dummy = torch.zeros(num_tokens, 1, HEAD_DIM, dtype=dtype, device=device)
207 - _call_fused(
238 + _ = _call_fused(
239 q_dummy,
240 + 64,
241 kv,
242 k_cache_fused,
243 slot_mapping,
@@ -298,8 +330,9 @@ def test_kv_path_with_dp_padding(num_tokens: int, pad: int, block_size: int):
330 # Fused: pass full-sized q/kv/positions, shorter slot_mapping.
331 q_dummy = torch.zeros(total, 1, HEAD_DIM, dtype=dtype, device=device)
332 k_cache_fused = torch.zeros_like(k_cache_ref)
301 - _call_fused(
333 + _ = _call_fused(
334 q_dummy,
335 + 64,
336 kv,
337 k_cache_fused,
338 slot_mapping,
@@ -316,9 +349,26 @@ def test_kv_path_with_dp_padding(num_tokens: int, pad: int, block_size: int):
349
350
351 @pytest.mark.parametrize("num_tokens", [1, 4, 17, 2048])
319 -@pytest.mark.parametrize("n_heads", [8, 64])
352 +@pytest.mark.parametrize(
353 + "n_heads,padded_heads",
354 + [
355 + # Each supported padded_heads instantiation: padded (n_heads <
356 + # padded_heads) and unpadded (n_heads == padded_heads).
357 + (1, 8),
358 + (8, 8),
359 + (8, 16),
360 + (16, 16),
361 + (16, 32),
362 + (32, 32),
363 + (8, 64),
364 + (64, 64),
365 + (64, 128),
366 + ],
367 +)
368 @pytest.mark.parametrize("block_size", [16, 64])
321 -def test_combined_q_and_kv(num_tokens: int, n_heads: int, block_size: int):
369 +def test_combined_q_and_kv(
370 + num_tokens: int, n_heads: int, padded_heads: int, block_size: int
371 +):
372 torch.manual_seed(2)
373 device = "cuda"
374 dtype = torch.bfloat16
@@ -345,10 +395,10 @@ def test_combined_q_and_kv(num_tokens: int, n_heads: int, block_size: int):
395 )
396
397 # Fused single call.
348 - q_fused = q.clone()
398 k_cache_fused = torch.zeros_like(k_cache_ref)
350 - _call_fused(
351 - q_fused,
399 + q_out = _call_fused(
400 + q,
401 + padded_heads,
402 kv,
403 k_cache_fused,
404 slot_mapping,
@@ -358,5 +408,10 @@ def test_combined_q_and_kv(num_tokens: int, n_heads: int, block_size: int):
408 block_size,
409 )
410
361 - torch.testing.assert_close(q_fused, q_ref, rtol=1e-2, atol=1e-2)
411 + torch.testing.assert_close(q_out[:, :n_heads], q_ref, rtol=1e-2, atol=1e-2)
412 + if n_heads < padded_heads:
413 + pad_region = q_out[:, n_heads:padded_heads]
414 + assert pad_region.abs().max().item() == 0.0, (
415 + "padded head slots must be exact zero"
416 + )
417 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
@@ -8,7 +8,7 @@ the existing separate operations (inverse RoPE via rotate_neox + FP8 quant
8 via per_token_group_quant_fp8).
9
10 The reference faithfully reproduces the exact flow in
11 -deepseek_v4/nvidia/ops/attention.py:295-310:
11 +deepseek_v4/attention.py:295-310:
12 1. Apply inverse RoPE (NeoX style, last rope_dim=64 dims of each head)
13 2. Reshape [T, H, head_dim] -> [T, G, D]
14 3. Transpose+flatten to [G*T, D], quantize, reshape back
@@ -668,7 +668,7 @@ def _unfused_inv_rope_fp8_quant(
668 nope_dim: int = NOPE_DIM,
669 rope_dim: int = ROPE_DIM,
670 ) -> tuple[torch.Tensor, torch.Tensor]:
671 - """Unfused path matching deepseek_v4/nvidia/ops/attention.py:295-310.
671 + """Unfused path matching deepseek_v4/attention.py:295-310.
672
673 Uses the production CUDA RoPE kernel + per_token_group_quant_fp8.
674 """
vllm/models/deepseek_v4/amd/model.py
+1 -1
@@ -53,7 +53,7 @@ from vllm.model_executor.models.utils import (
53 maybe_prefix,
54 )
55 from vllm.model_executor.utils import set_weight_attrs
56 -from vllm.models.deepseek_v4.nvidia.ops.attention import (
56 +from vllm.models.deepseek_v4.attention import (
57 DeepseekV4Indexer,
58 DeepseekV4MLAModules,
59 DeepseekV4MultiHeadLatentAttentionWrapper,
vllm/models/deepseek_v4/amd/rocm.py
+5 -1
@@ -32,7 +32,7 @@ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
32 from vllm.v1.worker.workspace import current_workspace_manager
33
34 if TYPE_CHECKING:
35 - from vllm.models.deepseek_v4.nvidia.ops.attention import (
35 + from vllm.models.deepseek_v4.attention import (
36 DeepseekV4MLAAttention,
37 )
38
@@ -592,6 +592,10 @@ class DeepseekV4ROCMAiterMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
592
593 backend_cls = DeepseekV4ROCMAiterMLASparseBackend
594
595 + @classmethod
596 + def get_padded_num_q_heads(cls, num_heads: int) -> int:
597 + return num_heads
598 +
599 @classmethod
600 def forward_mqa( # type: ignore[override]
601 cls,
vllm/models/deepseek_v4/attention.py renamed
+23 -40
@@ -156,18 +156,6 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
156 self.head_dim = head_dim
157 self.scale = scale
158
159 - # FlashMLA sparse kernel only supports 64 or 128 heads; pad up to the
160 - # next supported size. Must match DeepseekV4MLAAttention.padded_heads.
161 - if num_heads <= 64:
162 - self.padded_heads = 64
163 - elif num_heads <= 128:
164 - self.padded_heads = 128
165 - else:
166 - raise ValueError(
167 - f"DeepseekV4 attention does not support {num_heads} heads "
168 - "(must be <= 128)."
169 - )
170 -
159 self.q_lora_rank = q_lora_rank
160 self.kv_lora_rank = kv_lora_rank
161 self.window_size = window_size
@@ -263,6 +251,9 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
251 indexer=self.indexer,
252 topk_indices_buffer=self.topk_indices_buffer,
253 )
254 + # Mirror the inner layer's padded head count (single source of truth).
255 + self.padded_heads = self.mla_attn.padded_heads
256 +
257 # Register this layer in the compilation config's static forward context
258 # This allows the custom op to retrieve the layer during execution
259 compilation_config = mla_modules.vllm_config.compilation_config
@@ -450,7 +441,7 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
441
442 def wq_b_kv_insert() -> torch.Tensor:
443 q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim)
453 - self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)
444 + q = self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)
445 return q
446
447 # 3-way overlap (matches TRT-LLM PR #14142 Level 1): default runs
@@ -484,7 +475,7 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
475
476 def wq_b_kv_insert() -> torch.Tensor:
477 q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim)
487 - self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)
478 + q = self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)
479 return q
480
481 q, _ = maybe_execute_in_parallel(
@@ -497,12 +488,7 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
488 else:
489 # SWA-only layer: no compressor, no overlap.
490 q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim)
500 - self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)
501 -
502 - # Pad q to FlashMLA-required head count (64 or 128)
503 - if self.n_local_heads < self.padded_heads:
504 - pad_size = self.padded_heads - self.n_local_heads
505 - q = F.pad(q, (0, 0, 0, pad_size), value=0.0)
491 + q = self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)
492
493 # MLA attention writes into the pre-allocated `out` buffer
494 # ([num_tokens, padded_heads, head_dim]).
@@ -516,9 +502,17 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
502 attn_metadata: (
503 dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None
504 ),
519 - ) -> None:
505 + ) -> torch.Tensor:
506 if not isinstance(attn_metadata, dict):
521 - return
507 + # Profile run: kernel doesn't fire; produce a padded tensor so
508 + # downstream FlashMLA gets the right shape.
509 + if self.n_local_heads < self.padded_heads:
510 + return F.pad(
511 + q,
512 + (0, 0, 0, self.padded_heads - self.n_local_heads),
513 + value=0.0,
514 + )
515 + return q
516
517 swa_metadata = cast(
518 "DeepseekSparseSWAMetadata | None",
@@ -530,16 +524,19 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
524 swa_kv_cache_2d = swa_kv_cache.view(swa_kv_cache.shape[0], -1)
525
526 # Horizontally fused:
533 - # Q side: q_head_norm (per-head RMSNorm, no weight) + GPT-J RoPE
527 + # Q side: q_head_norm (per-head RMSNorm, no weight) + GPT-J RoPE,
528 + # with zero-fill for the padding head slots. The kernel
529 + # allocates and returns the padded q tensor.
530 # KV side: GPT-J RoPE + UE8M0 FP8 quant + paged cache insert
531 # kv is unchanged; mla_attn reads kv solely via swa_kv_cache.
536 - torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
532 + return torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
533 q,
534 kv,
535 swa_kv_cache_2d,
536 swa_metadata.slot_mapping,
537 positions.to(torch.int64),
538 self.rotary_emb.cos_sin_cache,
539 + self.padded_heads,
540 self.eps,
541 swa_metadata.block_size,
542 )
@@ -607,9 +604,6 @@ direct_register_custom_op(
604
605
606 class DeepseekV4MLAAttention(nn.Module, AttentionLayerBase):
610 - # FlashMLA FP8 sparse only supports 64 or 128 heads
611 - SUPPORTED_HEAD_COUNTS = (64, 128)
612 -
607 def __init__(
608 self,
609 num_heads: int,
@@ -655,19 +649,8 @@ class DeepseekV4MLAAttention(nn.Module, AttentionLayerBase):
649 self.aux_stream = aux_stream
650 self.ln_events = [torch.cuda.Event(), torch.cuda.Event()]
651
658 - # Determine padded head count for FlashMLA
659 - if num_heads not in self.SUPPORTED_HEAD_COUNTS:
660 - if num_heads < 64:
661 - self.padded_heads = 64
662 - elif num_heads < 128:
663 - self.padded_heads = 128
664 - else:
665 - raise ValueError(
666 - f"DeepseekV4MLAAttention does not support {num_heads} heads. "
667 - f"Supported: <= 128 (will be padded to 64 or 128)"
668 - )
669 - else:
670 - self.padded_heads = num_heads
652 + # Padded Q head count is dictated by the selected impl.
653 + self.padded_heads = self.impl_cls.get_padded_num_q_heads(num_heads)
654
655 # Store attention sink
656 assert attn_sink is not None
vllm/models/deepseek_v4/nvidia/flashmla.py
+23 -1
@@ -28,7 +28,7 @@ from vllm.v1.attention.ops.flashmla import (
28 from vllm.v1.worker.workspace import current_workspace_manager
29
30 if TYPE_CHECKING:
31 - from vllm.models.deepseek_v4.nvidia.ops.attention import (
31 + from vllm.models.deepseek_v4.attention import (
32 DeepseekV4MLAAttention,
33 )
34 from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
@@ -63,6 +63,18 @@ class DeepseekV4SparseMLAAttentionImpl(SparseMLAAttentionImpl[FlashMLASparseMeta
63 ) -> None:
64 raise NotImplementedError
65
66 + @classmethod
67 + @abstractmethod
68 + def get_padded_num_q_heads(cls, num_heads: int) -> int:
69 + """Q head count the backend wants q allocated at.
70 +
71 + The MLA wrapper allocates the q/output buffers at
72 + ``[N, get_padded_num_q_heads(n_local_heads), head_dim]``. Must
73 + satisfy ``result >= num_heads``. Backends with no padding constraint
74 + return ``num_heads``.
75 + """
76 + raise NotImplementedError
77 +
78
79 class DeepseekV4FlashMLASparseBackend(FlashMLASparseBackend):
80 @staticmethod
@@ -104,6 +116,16 @@ class DeepseekV4FlashMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
116
117 backend_cls = DeepseekV4FlashMLASparseBackend
118
119 + @classmethod
120 + def get_padded_num_q_heads(cls, num_heads: int) -> int:
121 + # FP8 decode kernel only supports h_q = 64 or 128.
122 + if num_heads > 128:
123 + raise ValueError(
124 + f"DeepseekV4 FlashMLA does not support {num_heads} heads "
125 + "(FP8 decode kernel requires h_q in {64, 128})."
126 + )
127 + return 64 if num_heads <= 64 else 128
128 +
129 @classmethod
130 def forward_mqa( # type: ignore[override]
131 cls,
vllm/models/deepseek_v4/nvidia/model.py
+1 -1
@@ -54,7 +54,7 @@ from vllm.model_executor.models.utils import (
54 maybe_prefix,
55 )
56 from vllm.model_executor.utils import set_weight_attrs
57 -from vllm.models.deepseek_v4.nvidia.ops.attention import (
57 +from vllm.models.deepseek_v4.attention import (
58 DeepseekV4Indexer,
59 DeepseekV4MLAModules,
60 DeepseekV4MultiHeadLatentAttentionWrapper,