[Model][DSv4.1] FlashMLA mega attention and the NVFP4 compressed KV cache (#56935)
Signed-off-by: Yongye Zhu <yongye@inferact.ai> Signed-off-by: 云挚 <ningyunxiao.nyx@antgroup.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Yongye Zhu <yongye@inferact.ai> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: Yunxiao Ning <73625538+foraxe@users.noreply.github.com> Co-authored-by: AI Assistant <noreply@openai.com>
Yongye Zhu committed
Sep 16, 2026 at 03:07 UTC
d6a1677d5504244c566eb900ca605cb5511f4ab4
20 files changed
+1811
-156
benchmarks/kernels/benchmark_dsv41_mega_attn.py
new
+272
@@ -0,0 +1,272 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""FlashMLA mega attention vs the split-KV decode pipeline (DeepSeek V4.1).
4
+
5
+Times the full set of per-layer ops each path runs in a decode step, captured
6
+in a CUDA graph -- eager timing is dominated by launch overhead and flatters
7
+whichever path launches fewer kernels.
8
+
9
+ mega: fused_qnorm_rope_kv_rope_quant_insert (Q pad + KV insert)
10
+ + mega kernel
11
+ split-KV: fused_qnorm_rope_kv_rope_quant_insert (Q RoPE + pad + KV insert)
12
+ + flash_mla_with_kvcache + fused_inv_rope_fp8_quant
13
+
14
+Both arms fuse their Q preparation into the KV insert, so that op is charged
15
+to both; the mega arm's Q side is only a zero-pad (nothing at all when the
16
+shard is already at the kernel head count) because its kernel does the Q RoPE,
17
+the inverse RoPE and the FP8 cast.
18
+
19
+``--local-heads`` is the variable that decides the outcome: the mega kernel
20
+absorbs work proportional to the live head count while paying for the padded
21
+one, so it wins at 64 live heads and is a wash at 16 (TP4 on a 64-head model).
22
+
23
+Both arms use the same KV records so the comparison is kernel-for-kernel:
24
+a V4.1 fp8 (528 B) sliding-window cache and a compressed cache that is either
25
+V4.1 fp8 or V4.1 NVFP4 (288 B) -- FlashMLA's SM100 sparse decode reads the
26
+NVFP4 compressed record in both the fused and unfused kernels.
27
+
28
+Run: .venv/bin/python benchmarks/kernels/benchmark_dsv41_mega_attn.py
29
+"""
30
+
31
+import argparse
32
+
33
+import torch
34
+
35
+import vllm.v1.attention.ops.flashmla as fm
36
+from vllm.models.deepseek_v4.common.ops.fused_inv_rope_fp8_quant import (
37
+ fused_inv_rope_fp8_quant,
38
+)
39
+from vllm.models.deepseek_v41.common.ops import quantize_and_insert_k_cache
40
+from vllm.models.deepseek_v41.common.ops.fused_compress_quant_cache import (
41
+ rope_quant_insert,
42
+)
43
+from vllm.models.deepseek_v41.common.ops.fused_layout import permute_q_to_fused
44
+from vllm.models.deepseek_v41.nvidia.flash_mla_mega_attn import (
45
+ alloc_mega_attn_output,
46
+ is_flashmla_mega_attn_supported,
47
+)
48
+from vllm.utils.math_utils import round_up
49
+
50
+HEAD_DIM, ROPE_DIM = 512, 64
51
+V41_BYTES, V41_FP4_BYTES = 528, 288
52
+SWA_BLOCK, COMPRESSED_BLOCK = 32, 128
53
+
54
+
55
+def make_cos_sin_cache(max_pos, device):
56
+ inv_freq = 1.0 / (
57
+ 10000 ** (torch.arange(0, ROPE_DIM, 2, device=device).float() / ROPE_DIM)
58
+ )
59
+ freqs = torch.outer(torch.arange(max_pos, device=device).float(), inv_freq)
60
+ return torch.cat([freqs.cos(), freqs.sin()], -1)
61
+
62
+
63
+def empty_paged_cache(num_rows, block_size, bytes_per_token, device):
64
+ """A zeroed paged cache view [num_blocks, block_size, bytes_per_token]."""
65
+ num_blocks = (num_rows + block_size - 1) // block_size + 1
66
+ page = round_up(block_size * bytes_per_token, 512)
67
+ backing = torch.zeros(num_blocks, page, dtype=torch.uint8, device=device)
68
+ return backing, backing.as_strided(
69
+ (num_blocks, block_size, bytes_per_token), (page, bytes_per_token, 1)
70
+ )
71
+
72
+
73
+def fill_cache(view, bytes_per_token, cos_sin, device):
74
+ n = view.shape[0] * view.shape[1]
75
+ rows = torch.randn(n, HEAD_DIM, device=device, dtype=torch.bfloat16)
76
+ slots = torch.arange(n, dtype=torch.int64, device=device)
77
+ if bytes_per_token == V41_FP4_BYTES:
78
+ # A page-aligned cache has more slots than the cos_sin table has rows;
79
+ # wrap so the insert never indexes past it.
80
+ positions = slots % cos_sin.shape[0]
81
+ rope_quant_insert(rows, positions, cos_sin, view, slots, 1)
82
+ else:
83
+ quantize_and_insert_k_cache(
84
+ rows,
85
+ view.reshape(view.shape[0], -1),
86
+ slots,
87
+ block_size=view.shape[1],
88
+ bytes_per_token=bytes_per_token,
89
+ )
90
+
91
+
92
+def time_graph(fn, iters=100, warmup=5):
93
+ """Capture ``fn`` in a CUDA graph and time its replay."""
94
+ s = torch.cuda.Stream()
95
+ s.wait_stream(torch.cuda.current_stream())
96
+ with torch.cuda.stream(s):
97
+ for _ in range(3):
98
+ fn()
99
+ torch.cuda.current_stream().wait_stream(s)
100
+ g = torch.cuda.CUDAGraph()
101
+ with torch.cuda.graph(g):
102
+ fn()
103
+ for _ in range(warmup):
104
+ g.replay()
105
+ torch.accelerator.synchronize()
106
+ start = torch.cuda.Event(enable_timing=True)
107
+ end = torch.cuda.Event(enable_timing=True)
108
+ start.record()
109
+ for _ in range(iters):
110
+ g.replay()
111
+ end.record()
112
+ torch.accelerator.synchronize()
113
+ return start.elapsed_time(end) * 1000 / iters
114
+
115
+
116
+def bench_decode(s_q, device, topk_extra, extra_bytes, local_heads):
117
+ padded_heads, scale = 64, HEAD_DIM**-0.5
118
+ n_wv_group = padded_heads // 8
119
+ cos_sin = make_cos_sin_cache(65536, device)
120
+ positions = torch.randint(0, 65536, (s_q,), device=device, dtype=torch.int64)
121
+ pos32 = positions.to(torch.int32)
122
+
123
+ # Per-step inputs: the wq_b output (local heads) and this step's KV row.
124
+ q_local = torch.randn(
125
+ s_q, local_heads, HEAD_DIM, device=device, dtype=torch.bfloat16
126
+ )
127
+ q_fused = permute_q_to_fused(q_local)
128
+ kv = torch.randn(s_q, HEAD_DIM, device=device, dtype=torch.bfloat16)
129
+
130
+ n_swa = s_q * 128 + 128
131
+ swa_backing, swa = empty_paged_cache(n_swa, SWA_BLOCK, V41_BYTES, device)
132
+ fill_cache(swa, V41_BYTES, cos_sin, device)
133
+ _, extra = empty_paged_cache(65536, COMPRESSED_BLOCK, extra_bytes, device)
134
+ fill_cache(extra, extra_bytes, cos_sin, device)
135
+
136
+ slot_mapping = torch.arange(s_q, dtype=torch.int64, device=device)
137
+ swa_idx = torch.randint(0, n_swa, (s_q, 128), device=device, dtype=torch.int32)
138
+ swa_len = torch.full((s_q,), 128, device=device, dtype=torch.int32)
139
+ ex_idx = torch.randint(
140
+ 0, 65536, (s_q, topk_extra), device=device, dtype=torch.int32
141
+ )
142
+ ex_len = torch.full((s_q,), topk_extra, device=device, dtype=torch.int32)
143
+ sink = torch.zeros(padded_heads, device=device)
144
+ sched = fm.FlashMLASchedMeta()
145
+ out = alloc_mega_attn_output(s_q, n_wv_group, device)
146
+ swa4 = swa.unsqueeze(-2)
147
+ extra4 = extra.unsqueeze(-2)
148
+ swa_2d = swa_backing
149
+
150
+ def mega():
151
+ # One launch zero-pads the fused-layout Q and inserts the SWA KV.
152
+ q_pad = torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
153
+ q_fused,
154
+ kv,
155
+ swa_2d,
156
+ slot_mapping,
157
+ positions,
158
+ cos_sin,
159
+ 0 if padded_heads == local_heads else padded_heads,
160
+ 1e-6,
161
+ SWA_BLOCK,
162
+ False, # apply_q_norm
163
+ True, # kv_mxfp8
164
+ False, # apply_q_rope: the mega kernel rotates Q itself
165
+ True, # is_q_interleaved
166
+ )
167
+ torch.ops._flashmla_C.fused_norm_rope_attn_rope_cast_decode(
168
+ q_fused if padded_heads == local_heads else q_pad,
169
+ swa4,
170
+ swa_idx,
171
+ scale,
172
+ HEAD_DIM,
173
+ sink,
174
+ swa_len,
175
+ extra4,
176
+ ex_idx,
177
+ ex_len,
178
+ False,
179
+ 0.0,
180
+ pos32,
181
+ False,
182
+ 64,
183
+ cos_sin,
184
+ n_wv_group,
185
+ 32,
186
+ True,
187
+ True,
188
+ True,
189
+ out.data,
190
+ out.scale,
191
+ )
192
+
193
+ def split_kv():
194
+ q_pad = torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
195
+ q_local,
196
+ kv,
197
+ swa_2d,
198
+ slot_mapping,
199
+ positions,
200
+ cos_sin,
201
+ padded_heads,
202
+ 1e-6,
203
+ SWA_BLOCK,
204
+ False,
205
+ True,
206
+ )
207
+ o = fm.flash_mla_with_kvcache(
208
+ q=q_pad.unsqueeze(1),
209
+ k_cache=swa4,
210
+ block_table=None,
211
+ head_dim_v=HEAD_DIM,
212
+ tile_scheduler_metadata=sched,
213
+ cache_seqlens=None,
214
+ is_fp8_kvcache=True,
215
+ indices=swa_idx.view(s_q, 1, -1),
216
+ topk_length=swa_len,
217
+ softmax_scale=scale,
218
+ attn_sink=sink,
219
+ extra_k_cache=extra4,
220
+ extra_indices_in_kvcache=ex_idx.view(s_q, 1, -1),
221
+ extra_topk_length=ex_len,
222
+ )[0]
223
+ fused_inv_rope_fp8_quant(
224
+ o.squeeze(1)[:, :local_heads],
225
+ positions,
226
+ cos_sin,
227
+ n_groups=local_heads // 8,
228
+ heads_per_group=8,
229
+ quant_group_size=32,
230
+ tma_aligned_scales=True,
231
+ )
232
+
233
+ return time_graph(mega), time_graph(split_kv)
234
+
235
+
236
+def main():
237
+ parser = argparse.ArgumentParser()
238
+ parser.add_argument(
239
+ "--decode-s-q",
240
+ nargs="+",
241
+ type=int,
242
+ default=[1, 2, 4, 8, 16, 32, 64, 128, 256, 512],
243
+ )
244
+ parser.add_argument("--topk-extra", type=int, default=512)
245
+ parser.add_argument("--local-heads", type=int, default=16)
246
+ parser.add_argument(
247
+ "--extra-bytes",
248
+ type=int,
249
+ choices=[V41_BYTES, V41_FP4_BYTES],
250
+ default=V41_BYTES,
251
+ help="compressed-cache record: 528 (V4.1 fp8) or 288 (V4.1 NVFP4)",
252
+ )
253
+ args = parser.parse_args()
254
+ ok, reason = is_flashmla_mega_attn_supported()
255
+ if not ok:
256
+ raise SystemExit(reason)
257
+ device = torch.device("cuda")
258
+ torch.manual_seed(0)
259
+ print(
260
+ f"decode, cudagraph (topk_swa=128, topk_extra={args.topk_extra}, "
261
+ f"local_heads={args.local_heads}, extra={args.extra_bytes}B)"
262
+ )
263
+ print(f"{'s_q':>6} {'mega_us':>9} {'split_kv_us':>12} {'speedup':>8}")
264
+ for s_q in args.decode_s_q:
265
+ m, s = bench_decode(
266
+ s_q, device, args.topk_extra, args.extra_bytes, args.local_heads
267
+ )
268
+ print(f"{s_q:>6d} {m:9.1f} {s:12.1f} {s / m:7.2f}x")
269
+
270
+
271
+if __name__ == "__main__":
272
+ main()
csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
+134
-72
@@ -3,7 +3,10 @@
3
* SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4
*
5
* Horizontally-fused DeepseekV4-MLA kernel:
6
- * - Q side: optional per-head RMSNorm + GPT-J RoPE on last ROPE_DIM
6
+ * - Q side: optional per-head RMSNorm (apply_q_norm) and GPT-J RoPE on the
7
+ * last ROPE_DIM (apply_q_rope), plus the zero-pad to
8
+ * q_head_padded, in the head-major layout or (is_q_interleaved)
9
+ * FlashMLA's mega-attention chunk-interleaved one
10
* - KV side: GPT-J RoPE on last ROPE_DIM + UE8M0 FP8 quant on NoPE + paged
11
* cache insert
12
*
@@ -174,6 +177,30 @@ __device__ __forceinline__ float warpSum(float val) {
177
return val;
178
}
179
180
+// Offset of the 16 elements a lane owns inside a Q tensor of `num_heads`
181
+// heads. The default layout is head-major: a head's 512 dims are contiguous.
182
+// `Q_INTERLEAVED` is FlashMLA's mega-attention layout, 32 chunks of
183
+// [num_heads, 16] per token, which changes the address only -- lane L still
184
+// owns dims [16L, 16L+16) of its head either way, so `dim_base` stays the
185
+// true dim offset and the norm and RoPE below are equally valid in both.
186
+// What it does change is where the padding heads land: at the tail of every
187
+// chunk rather than after the live heads.
188
+template <bool Q_INTERLEAVED>
189
+__device__ __forceinline__ int64_t qElemOffset(int const tokenIdx,
190
+ int const slotIdx,
191
+ int const laneId,
192
+ int const dim_base,
193
+ int const num_heads) {
194
+ if constexpr (Q_INTERLEAVED) {
195
+ return static_cast<int64_t>(tokenIdx) * num_heads * kHeadDim +
196
+ static_cast<int64_t>(laneId) * num_heads * kElemsPerLane +
197
+ static_cast<int64_t>(slotIdx) * kElemsPerLane;
198
+ } else {
199
+ return (static_cast<int64_t>(tokenIdx) * num_heads + slotIdx) * kHeadDim +
200
+ dim_base;
201
+ }
202
+}
203
+
204
// ────────────────────────────────────────────────────────────────────────────
205
// Per-slot inner pipeline
206
// ────────────────────────────────────────────────────────────────────────────
@@ -192,10 +219,16 @@ __device__ __forceinline__ float warpSum(float val) {
219
// + paged-cache
220
// insert)
221
//
195
-// `kv_mxfp8` is grid-uniform and only steers the KV branch, so it stays a
196
-// runtime argument rather than doubling every (head count, q-norm)
197
-// instantiation of a kernel whose hot path is the Q branch.
198
-template <typename scalar_t_in, int kNumHeadsQPadded, bool APPLY_Q_NORM>
222
+// `kv_mxfp8` and `apply_q_rope` are grid-uniform, so they stay runtime
223
+// arguments rather than multiplying every (head count, q-norm, q-layout)
224
+// instantiation of a kernel whose hot path is the Q branch. `Q_INTERLEAVED`
225
+// cannot: it picks the address arithmetic, which has to fold.
226
+//
227
+// The three Q knobs are independent. The mega-attention layer uses
228
+// Q_INTERLEAVED with both norm and RoPE off, because its attention kernel
229
+// applies those itself -- not because the layout forbids them.
230
+template <typename scalar_t_in, int kNumHeadsQPadded, bool APPLY_Q_NORM,
231
+ bool Q_INTERLEAVED = false>
232
__device__ __forceinline__ void processDeepseekV4Slot(
233
uint4 v0, uint4 v1, int const tokenIdx, int const slotIdx,
234
int const dim_base, int const laneId, int const num_heads_q,
@@ -203,7 +236,7 @@ __device__ __forceinline__ void processDeepseekV4Slot(
236
uint8_t* __restrict__ k_cache, int64_t const* __restrict__ slot_mapping,
237
int64_t const* __restrict__ position_ids,
238
float const* __restrict__ cos_sin_cache, int const cache_block_size,
206
- int const kv_block_stride, bool const kv_mxfp8) {
239
+ int const kv_block_stride, bool const kv_mxfp8, bool const apply_q_rope) {
240
using Converter = vllm::_typeConvert<scalar_t_in>;
241
bool const isKV = (slotIdx == kNumHeadsQPadded);
242
bool const isPadQ = !isKV && (slotIdx >= num_heads_q);
@@ -213,10 +246,8 @@ __device__ __forceinline__ void processDeepseekV4Slot(
246
// zero literal is correct. Matches the live-Q branch's vectorized store.
247
if (isPadQ) {
248
scalar_t_in* dst =
216
- q_out +
217
- (static_cast<int64_t>(tokenIdx) * kNumHeadsQPadded + slotIdx) *
218
- kHeadDim +
219
- dim_base;
249
+ q_out + qElemOffset<Q_INTERLEAVED>(tokenIdx, slotIdx, laneId, dim_base,
250
+ kNumHeadsQPadded);
251
uint4 const zero4 = {0u, 0u, 0u, 0u};
252
*reinterpret_cast<uint4*>(dst) = zero4;
253
*reinterpret_cast<uint4*>(dst + 8) = zero4;
@@ -264,7 +295,9 @@ __device__ __forceinline__ void processDeepseekV4Slot(
295
296
// ── GPT-J RoPE on dims [NOPE_DIM, HEAD_DIM) ─────────────────────────────
297
// All math in fp32. cos_sin_cache is loaded as fp32 (its native storage).
267
- bool const is_rope_lane = dim_base >= kNopeDim;
298
+ // KV is always rotated; Q only when the caller has not already had it
299
+ // rotated elsewhere (the mega-attention kernel does its own).
300
+ bool const is_rope_lane = dim_base >= kNopeDim && (isKV || apply_q_rope);
301
if (is_rope_lane) {
302
int64_t const pos = position_ids[tokenIdx];
303
constexpr int kHalfRope = kRopeDim / 2;
@@ -314,10 +347,8 @@ __device__ __forceinline__ void processDeepseekV4Slot(
347
make_float2(elements[8 + 2 * i], elements[8 + 2 * i + 1]));
348
}
349
scalar_t_in* dst =
317
- q_out +
318
- (static_cast<int64_t>(tokenIdx) * kNumHeadsQPadded + slotIdx) *
319
- kHeadDim +
320
- dim_base;
350
+ q_out + qElemOffset<Q_INTERLEAVED>(tokenIdx, slotIdx, laneId, dim_base,
351
+ kNumHeadsQPadded);
352
*reinterpret_cast<uint4*>(dst) = out0;
353
*reinterpret_cast<uint4*>(dst + 8) = out1;
354
} else {
@@ -432,14 +463,17 @@ __device__ __forceinline__ void processDeepseekV4Slot(
463
// `kNumHeadsQPadded` is a template parameter (compile-time constant) so the
464
// divisions in the grid math and the KV-sentinel comparison fold to fast
465
// constant operations. The launch wrapper dispatches the runtime value to
435
-// the matching instantiation.
466
+// the matching instantiation. `kNumHeadsQPadded == 0` leaves one slot per
467
+// token and is the KV-insert-only instantiation, for a Q that is already in
468
+// the shape its attention kernel wants.
469
//
470
// With DP padding, q/kv/position_ids can have more rows than slot_mapping.
471
// The live-Q and pad-Q branches cover all `num_tokens_full` rows (downstream
472
// attention uses them). The KV branch only inserts the first
473
// `num_tokens_insert` tokens (= slot_mapping length) into the paged cache.
474
//
442
-template <typename scalar_t_in, int kNumHeadsQPadded, bool APPLY_Q_NORM>
475
+template <typename scalar_t_in, int kNumHeadsQPadded, bool APPLY_Q_NORM,
476
+ bool Q_INTERLEAVED>
477
__global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
478
scalar_t_in const* __restrict__ q_in, // [N, num_heads_q, 512]
479
scalar_t_in* __restrict__ q_out, // [N, kNumHeadsQPadded, 512]
@@ -454,7 +488,8 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
488
int const num_heads_q, // live Q heads (input layout)
489
int const cache_block_size, // tokens per paged-cache block
490
int const kv_block_stride, // bytes per paged-cache block
457
- bool const kv_mxfp8) { // V4.1 all-dims MXFP8 record
491
+ bool const kv_mxfp8, // V4.1 all-dims MXFP8 record
492
+ bool const apply_q_rope) { // rotate Q too, not just KV
493
#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
494
// BF16 _typeConvert specialization is unavailable on pre-Ampere. The
495
// DeepseekV4 kernel only runs with bf16 inputs in practice, so compile a
@@ -497,20 +532,18 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
532
if (isKV) {
533
src_ptr = kv_in + static_cast<int64_t>(tokenIdx) * kHeadDim + dim_base;
534
} else {
500
- int64_t const q_row_offset =
501
- (static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) *
502
- kHeadDim +
503
- dim_base;
504
- src_ptr = q_in + q_row_offset;
535
+ src_ptr = q_in + qElemOffset<Q_INTERLEAVED>(tokenIdx, slotIdx, laneId,
536
+ dim_base, num_heads_q);
537
}
538
v0 = *reinterpret_cast<uint4 const*>(src_ptr);
539
v1 = *reinterpret_cast<uint4 const*>(src_ptr + 8);
540
}
541
510
- processDeepseekV4Slot<scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM>(
542
+ processDeepseekV4Slot<scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM,
543
+ Q_INTERLEAVED>(
544
v0, v1, tokenIdx, slotIdx, dim_base, laneId, num_heads_q, eps, q_out,
545
k_cache, slot_mapping, position_ids, cos_sin_cache, cache_block_size,
513
- kv_block_stride, kv_mxfp8);
546
+ kv_block_stride, kv_mxfp8, apply_q_rope);
547
548
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
549
cudaTriggerProgrammaticLaunchCompletion();
@@ -530,7 +563,8 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
563
// Q branch (optional RMSNorm + RoPE, in place) head_slot == num_heads_q
564
// KV branch (RoPE + UE8M0 quant + insert)
565
//
533
-template <typename scalar_t_in, int kNumHeadsQPadded, bool APPLY_Q_NORM>
566
+template <typename scalar_t_in, int kNumHeadsQPadded, bool APPLY_Q_NORM,
567
+ bool Q_INTERLEAVED>
568
__global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid(
569
scalar_t_in const* __restrict__ q_in, scalar_t_in* __restrict__ q_out,
570
scalar_t_in const* __restrict__ kv_in, uint8_t* __restrict__ k_cache,
@@ -539,7 +573,7 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid(
573
float const* __restrict__ cos_sin_cache, float const eps,
574
int const num_tokens_full, int const num_tokens_insert,
575
int const num_heads_q, int const cache_block_size,
542
- int const kv_block_stride, bool const kv_mxfp8) {
576
+ int const kv_block_stride, bool const kv_mxfp8, bool const apply_q_rope) {
577
#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
578
if constexpr (std::is_same_v<scalar_t_in, c10::BFloat16>) {
579
return;
@@ -569,11 +603,8 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid(
603
if (s == kNumHeadsQPadded) {
604
src = kv_in + static_cast<int64_t>(tokenIdx) * kHeadDim + dim_base;
605
} else {
572
- src = q_in +
573
- (static_cast<int64_t>(tokenIdx) * num_heads_q +
574
- static_cast<int64_t>(s)) *
575
- kHeadDim +
576
- dim_base;
606
+ src = q_in + qElemOffset<Q_INTERLEAVED>(tokenIdx, s, laneId, dim_base,
607
+ num_heads_q);
608
}
609
va = *reinterpret_cast<uint4 const*>(src);
610
vb = *reinterpret_cast<uint4 const*>(src + 8);
@@ -594,10 +625,12 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid(
625
load_slot(next_slot, v0_next, v1_next);
626
}
627
597
- processDeepseekV4Slot<scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM>(
628
+ processDeepseekV4Slot<scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM,
629
+ Q_INTERLEAVED>(
630
v0_curr, v1_curr, tokenIdx, curr_slot, dim_base, laneId,
631
num_heads_q, eps, q_out, k_cache, slot_mapping, position_ids,
600
- cos_sin_cache, cache_block_size, kv_block_stride, kv_mxfp8);
632
+ cos_sin_cache, cache_block_size, kv_block_stride, kv_mxfp8,
633
+ apply_q_rope);
634
635
// ── Buffer rotation: hand the prefetched LDGs to the next iter.
636
v0_curr = v0_next;
@@ -617,14 +650,15 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid(
650
// ────────────────────────────────────────────────────────────────────────────
651
// Launch wrapper
652
// ────────────────────────────────────────────────────────────────────────────
620
-template <typename scalar_t_in, int kNumHeadsQPadded, bool APPLY_Q_NORM>
653
+template <typename scalar_t_in, int kNumHeadsQPadded, bool APPLY_Q_NORM,
654
+ bool Q_INTERLEAVED>
655
static void launchFusedDeepseekV4Templated(
656
scalar_t_in const* q_in, scalar_t_in* q_out, scalar_t_in const* kv_in,
657
uint8_t* k_cache, int64_t const* slot_mapping, int64_t const* position_ids,
658
float const* cos_sin_cache, float const eps, int const num_tokens_full,
659
int const num_tokens_insert, int const num_heads_q,
660
int const cache_block_size, int const kv_block_stride, bool const kv_mxfp8,
627
- cudaStream_t stream) {
661
+ bool const apply_q_rope, cudaStream_t stream) {
662
constexpr int kBlockSize = 256;
663
constexpr int kWarpsPerBlock = kBlockSize / 32;
664
int64_t const total_warps =
@@ -663,37 +697,40 @@ static void launchFusedDeepseekV4Templated(
697
// grid instead. Only reachable above NUM_TOKEN_CUTOFF tokens in a single
698
// insert, where it is worth ~2x.
699
if (kNumHeadsQPadded == 0 || num_tokens_full < NUM_TOKEN_CUTOFF) {
666
- cudaLaunchKernelEx(&config,
667
- fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<
668
- scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM>,
669
- q_in, q_out, kv_in, k_cache, slot_mapping, position_ids,
670
- cos_sin_cache, eps, num_tokens_full, num_tokens_insert,
671
- num_heads_q, cache_block_size, kv_block_stride,
672
- kv_mxfp8);
700
+ cudaLaunchKernelEx(
701
+ &config,
702
+ fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<
703
+ scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM, Q_INTERLEAVED>,
704
+ q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache,
705
+ eps, num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
706
+ kv_block_stride, kv_mxfp8, apply_q_rope);
707
} else {
708
config.gridDim = dim3(num_tokens_full);
709
cudaLaunchKernelEx(
710
&config,
711
fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid<
678
- scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM>,
712
+ scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM, Q_INTERLEAVED>,
713
q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache,
714
eps, num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
681
- kv_block_stride, kv_mxfp8);
715
+ kv_block_stride, kv_mxfp8, apply_q_rope);
716
}
717
#else
718
// ROCm: use standard kernel launch syntax (no PDL/stream serialization)
719
// clang-format off
720
fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<
687
- scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM>
721
+ scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM, Q_INTERLEAVED>
722
<<<grid, kBlockSize, 0, stream>>>(
723
q_in, q_out, kv_in, k_cache, slot_mapping, position_ids,
724
cos_sin_cache, eps, num_tokens_full, num_tokens_insert, num_heads_q,
691
- cache_block_size, kv_block_stride, kv_mxfp8);
725
+ cache_block_size, kv_block_stride, kv_mxfp8, apply_q_rope);
726
#endif
727
}
728
729
// Runtime dispatch into one of the precompiled `kNumHeadsQPadded`
696
-// instantiations. Supported padded head counts: 8, 16, 32, 64, 128.
730
+// instantiations. Supported padded head counts: 8, 16, 32, 64, 128, plus 0
731
+// for a KV-only launch that writes no Q at all. `apply_q_norm` and
732
+// `is_q_interleaved` are template parameters, `apply_q_rope` is not (see
733
+// processDeepseekV4Slot), so each head count carries four instantiations.
734
template <typename scalar_t_in>
735
void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
736
scalar_t_in const* q_in, scalar_t_in* q_out, scalar_t_in const* kv_in,
@@ -702,22 +739,32 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
739
int const num_tokens_full, int const num_tokens_insert,
740
int const num_heads_q, int const num_heads_q_padded,
741
int const cache_block_size, int const kv_block_stride,
705
- bool const apply_q_norm, bool const kv_mxfp8, cudaStream_t stream) {
706
-#define DISPATCH(N) \
707
- case N: \
708
- if (apply_q_norm) { \
709
- launchFusedDeepseekV4Templated<scalar_t_in, N, true>( \
710
- q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, \
711
- cos_sin_cache, eps, num_tokens_full, num_tokens_insert, \
712
- num_heads_q, cache_block_size, kv_block_stride, kv_mxfp8, stream); \
713
- } else { \
714
- launchFusedDeepseekV4Templated<scalar_t_in, N, false>( \
715
- q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, \
716
- cos_sin_cache, eps, num_tokens_full, num_tokens_insert, \
717
- num_heads_q, cache_block_size, kv_block_stride, kv_mxfp8, stream); \
718
- } \
742
+ bool const apply_q_norm, bool const kv_mxfp8, bool const apply_q_rope,
743
+ bool const is_q_interleaved, cudaStream_t stream) {
744
+#define LAUNCH(N, NORM, INTERLEAVED) \
745
+ launchFusedDeepseekV4Templated<scalar_t_in, N, NORM, INTERLEAVED>( \
746
+ q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, \
747
+ eps, num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size, \
748
+ kv_block_stride, kv_mxfp8, apply_q_rope, stream)
749
+#define DISPATCH(N) \
750
+ case N: \
751
+ if (is_q_interleaved) { \
752
+ if (apply_q_norm) { \
753
+ LAUNCH(N, true, true); \
754
+ } else { \
755
+ LAUNCH(N, false, true); \
756
+ } \
757
+ } else if (apply_q_norm) { \
758
+ LAUNCH(N, true, false); \
759
+ } else { \
760
+ LAUNCH(N, false, false); \
761
+ } \
762
return;
763
764
+ if (num_heads_q_padded == 0) {
765
+ LAUNCH(0, false, false);
766
+ return;
767
+ }
768
switch (num_heads_q_padded) {
769
DISPATCH(8)
770
DISPATCH(16)
@@ -729,9 +776,10 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
776
"fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert: "
777
"unsupported num_heads_q_padded=",
778
num_heads_q_padded,
732
- " (compiled instantiations: 8, 16, 32, 64, 128).");
779
+ " (compiled instantiations: 0, 8, 16, 32, 64, 128).");
780
}
781
#undef DISPATCH
782
+#undef LAUNCH
783
}
784
785
// ────────────────────────────────────────────────────────────────────────────
@@ -1059,10 +1107,15 @@ void fused_deepseek_v4_kv_rope_insert(
1107
// Zero query heads schedules only the existing KV branch, preserving its
1108
// RoPE rounding and quantization without allocating any query tensors.
1109
if (packed) {
1062
- vllm::deepseek_v4_fused_ops::launchFusedDeepseekV4Templated<scalar_t, 0, false>(
1110
+ // With zero query heads no Q slot is scheduled, so neither Q knob is
1111
+ // reachable: Q_INTERLEAVED picks addressing nothing uses, and apply_q_rope
1112
+ // gates a branch only a Q slot takes (KV is rotated unconditionally).
1113
+ vllm::deepseek_v4_fused_ops::launchFusedDeepseekV4Templated<scalar_t, 0,
1114
+ false, false>(
1115
nullptr, nullptr, input, cache, slots, positions, cos_sin, 0.0f,
1116
num_tokens, num_tokens, 0, static_cast<int>(cache_block_size),
1065
- static_cast<int>(k_cache.stride(0)), kv_mxfp8, stream);
1117
+ static_cast<int>(k_cache.stride(0)), kv_mxfp8, /*apply_q_rope=*/true,
1118
+ stream);
1119
} else if (fp8) {
1120
vllm::deepseek_v4_fused_ops::launchFullCacheKernel<scalar_t, false, true>(
1121
nullptr, nullptr, 0, 0, input, cache, slots, positions, cos_sin,
@@ -1085,8 +1138,9 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
1138
torch::stable::Tensor const& slot_mapping, // [N] int64
1139
torch::stable::Tensor const& position_ids, // [N] int64
1140
torch::stable::Tensor const& cos_sin_cache, // [max_pos, rope_dim] bf16
1088
- int64_t q_head_padded, // padded Q head count for output
1089
- double eps, int64_t cache_block_size, bool apply_q_norm, bool kv_mxfp8) {
1141
+ int64_t q_head_padded, // padded Q head count for output, 0 for no Q
1142
+ double eps, int64_t cache_block_size, bool apply_q_norm, bool kv_mxfp8,
1143
+ bool apply_q_rope, bool is_q_interleaved) {
1144
STD_TORCH_CHECK(q_in.device().is_cuda() && q_in.is_contiguous(),
1145
"q_in must be contiguous CUDA");
1146
STD_TORCH_CHECK(kv.device().is_cuda() && kv.is_contiguous(),
@@ -1106,8 +1160,11 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
1160
STD_TORCH_CHECK(kv.dim() == 2 && kv.size(1) == 512, "kv shape [N, 512]");
1161
STD_TORCH_CHECK(q_in.scalar_type() == kv.scalar_type(),
1162
"q_in and kv dtype must match");
1109
- STD_TORCH_CHECK(q_head_padded >= q_in.size(1),
1110
- "q_head_padded must be >= q_in.size(1) (num_heads_q)");
1163
+ // q_head_padded == 0 asks for the KV insert alone: Q is left untouched and
1164
+ // an empty tensor comes back, for a caller whose Q is already in the shape
1165
+ // its attention kernel reads.
1166
+ STD_TORCH_CHECK(q_head_padded == 0 || q_head_padded >= q_in.size(1),
1167
+ "q_head_padded must be 0 or >= q_in.size(1) (num_heads_q)");
1168
STD_TORCH_CHECK(k_cache.scalar_type() == torch::headeronly::ScalarType::Byte,
1169
"k_cache must be uint8");
1170
STD_TORCH_CHECK(cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == 64,
@@ -1138,8 +1195,12 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
1195
// Allocate the padded q output. The kernel writes every element (live
1196
// region gets optional RMSNorm+RoPE; pad region gets zeros), so `empty` is
1197
// safe.
1141
- auto q_out = torch::stable::new_empty(
1142
- q_in, {q_in.size(0), q_head_padded, q_in.size(2)}, q_in.scalar_type());
1198
+ auto q_out =
1199
+ q_head_padded == 0
1200
+ ? torch::stable::new_empty(q_in, {int64_t{0}}, q_in.scalar_type())
1201
+ : torch::stable::new_empty(
1202
+ q_in, {q_in.size(0), q_head_padded, q_in.size(2)},
1203
+ q_in.scalar_type());
1204
1205
VLLM_STABLE_DISPATCH_HALF_TYPES(
1206
q_in.scalar_type(), "fused_deepseek_v4_qnorm_rope_kv_insert", [&] {
@@ -1155,7 +1216,8 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
1216
cos_sin_cache.const_data_ptr<float>(), static_cast<float>(eps),
1217
num_tokens_full, num_tokens_insert, num_heads_q,
1218
num_heads_q_padded, cache_block_size_i, kv_block_stride,
1158
- apply_q_norm, kv_mxfp8, stream);
1219
+ apply_q_norm, kv_mxfp8, apply_q_rope, is_q_interleaved,
1220
+ stream);
1221
});
1222
return q_out;
1223
}
csrc/libtorch_stable/ops.h
+2
-1
@@ -271,7 +271,8 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
271
torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping,
272
torch::stable::Tensor const& position_ids,
273
torch::stable::Tensor const& cos_sin_cache, int64_t q_head_padded,
274
- double eps, int64_t cache_block_size, bool apply_q_norm, bool kv_mxfp8);
274
+ double eps, int64_t cache_block_size, bool apply_q_norm, bool kv_mxfp8,
275
+ bool apply_q_rope, bool is_q_interleaved);
276
277
void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(
278
torch::stable::Tensor& q, torch::stable::Tensor const& kv,
csrc/libtorch_stable/torch_bindings.cpp
+8
-1
@@ -425,6 +425,12 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
425
"bool is_neox, Tensor position_ids, "
426
"int forced_token_heads_per_warp=-1) -> ()");
427
428
+ // q_head_padded is the padded Q head count of the returned tensor, or 0 to
429
+ // do the KV insert alone and return an empty tensor. The Q knobs are
430
+ // independent: apply_q_norm and apply_q_rope each drop that step for Q
431
+ // alone (KV is always rotated), and is_q_interleaved reads and writes Q in
432
+ // FlashMLA's mega-attention chunk-interleaved layout, which moves the
433
+ // padding heads to the tail of every head-dim chunk.
434
ops.def(
435
"fused_deepseek_v4_kv_rope_insert("
436
"Tensor kv, Tensor! k_cache, Tensor slot_mapping, Tensor position_ids, "
@@ -436,7 +442,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
442
"Tensor q_in, Tensor kv, Tensor! k_cache, "
443
"Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, "
444
"int q_head_padded, float eps, int cache_block_size, "
439
- "bool apply_q_norm=True, bool kv_mxfp8=False) -> Tensor");
445
+ "bool apply_q_norm=True, bool kv_mxfp8=False, bool apply_q_rope=True, "
446
+ "bool is_q_interleaved=False) -> Tensor");
447
448
// FlashInfer V4 full-cache variants: write Q in place (bf16) or to a separate
449
// FP8 tensor, and KV into a contiguous 512-wide token-strided cache.
tests/kernels/test_compressor_kv_cache.py
+133
-6
@@ -331,6 +331,19 @@ def test_v41_fused_save_compress_and_insert(
331
torch.testing.assert_close(cache_backing, expected_cache, rtol=0, atol=0)
332
333
334
+def _rotate_rope_tail(
335
+ latent_row: torch.Tensor, cos_sin_row: torch.Tensor
336
+) -> torch.Tensor:
337
+ """GPT-J RoPE of one latent row's last 64 dims, in fp32."""
338
+ row = latent_row.float()
339
+ c, s = cos_sin_row.chunk(2)
340
+ rotated = row.clone()
341
+ even, odd = row[448::2], row[449::2]
342
+ rotated[448::2] = even * c - odd * s
343
+ rotated[449::2] = odd * c + even * s
344
+ return rotated
345
+
346
+
347
def _mxfp8_record_reference(
348
latent_row: torch.Tensor, cos_sin_row: torch.Tensor
349
) -> tuple[torch.Tensor, torch.Tensor]:
@@ -339,12 +352,7 @@ def _mxfp8_record_reference(
352
Mirrors FlashMLA's ``KVCacheLayout.V41_FP8Sparse`` quantizer: rotate the
353
RoPE tail first, then scale every 32-dim tile, RoPE tiles included.
354
"""
342
- row = latent_row.float()
343
- c, s = cos_sin_row.chunk(2)
344
- rotated = row.clone()
345
- even, odd = row[448::2], row[449::2]
346
- rotated[448::2] = even * c - odd * s
347
- rotated[449::2] = odd * c + even * s
355
+ rotated = _rotate_rope_tail(latent_row, cos_sin_row)
356
quantized, scales = _ue8m0_reference(rotated, 32, 448.0)
357
return quantized.view(torch.uint8), (scales.log2() + 127).to(torch.uint8)
358
@@ -461,6 +469,125 @@ def test_v41_mxfp8_cache_round_trip():
469
assert (error <= tolerance.repeat_interleave(32, dim=-1)).all()
470
471
472
+# e2m1 magnitudes indexed by the low 3 bits of a code; bit 3 is the sign.
473
+_E2M1_MAGNITUDES = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]
474
+
475
+
476
+def _decode_nvfp4_row(packed: torch.Tensor, scales: torch.Tensor) -> torch.Tensor:
477
+ """Unpack one 256-byte e2m1 row and apply its 32 e4m3 tile scales."""
478
+ mags = torch.tensor(_E2M1_MAGNITUDES, device=packed.device)
479
+ codes = torch.empty(512, dtype=torch.uint8, device=packed.device)
480
+ codes[0::2] = packed & 0xF # even element in the low nibble
481
+ codes[1::2] = packed >> 4
482
+ vals = mags[(codes & 7).long()] * torch.where(codes >= 8, -1.0, 1.0)
483
+ return (
484
+ vals.view(32, 16) * scales.view(torch.float8_e4m3fn).float().view(32, 1)
485
+ ).flatten()
486
+
487
+
488
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only")
489
+@pytest.mark.parametrize("compress_ratio", [1, 2])
490
+def test_v41_rope_insert_nvfp4_record(compress_ratio: int):
491
+ """The 288-byte V4.1 NVFP4 record: 256 B of e2m1 pairs then 32 e4m3 scales.
492
+
493
+ Scales are byte-exact against the reference (``amax / 6`` clamped to the
494
+ e4m3 range); values are checked as round-to-nearest on the e2m1 grid, whose
495
+ coarsest step is 2 between 4 and 6, so half a step is ``1.0 * scale``.
496
+ """
497
+ from vllm.models.deepseek_v41.common.ops.fused_compress_quant_cache import (
498
+ rope_quant_insert,
499
+ )
500
+
501
+ torch.manual_seed(13)
502
+ device = "cuda"
503
+ cache_block = 64
504
+ cache_stride = math.ceil(cache_block * 288 / 512) * 512
505
+ num_tokens = 12
506
+
507
+ positions = torch.arange(num_tokens, dtype=torch.int64, device=device)
508
+ latent = torch.randn(num_tokens, 512, dtype=torch.bfloat16, device=device)
509
+ angles = torch.randn(64, 32, device=device)
510
+ cos_sin = torch.cat((angles.cos(), angles.sin()), dim=-1)
511
+
512
+ cache_backing = torch.full((2, cache_stride), 165, dtype=torch.uint8, device=device)
513
+ cache = cache_backing.as_strided((2, cache_block, 288), (cache_stride, 288, 1))
514
+ slots = torch.arange(num_tokens, dtype=torch.int64, device=device)
515
+ slots[3] = -1 # covers the negative-slot skip
516
+
517
+ rope_quant_insert(latent, positions, cos_sin, cache, slots, compress_ratio)
518
+
519
+ untouched = torch.full_like(cache_backing, 165)
520
+ for t in range(num_tokens):
521
+ slot = slots[t].item()
522
+ pos = positions[t].item()
523
+ if slot < 0 or (pos + 1) % compress_ratio:
524
+ continue
525
+ page, row = divmod(slot, cache_block)
526
+ cs = cos_sin[pos // compress_ratio * compress_ratio]
527
+ rotated = _rotate_rope_tail(latent[t], cs)
528
+ amax = rotated.view(32, 16).abs().amax(-1)
529
+ scale = (amax / 6.0).clamp(2.0**-9, 448.0).to(torch.float8_e4m3fn)
530
+
531
+ got_scales = cache_backing[page, cache_block * 256 + row * 32 :][:32]
532
+ torch.testing.assert_close(got_scales, scale.view(torch.uint8), rtol=0, atol=0)
533
+ decoded = _decode_nvfp4_row(
534
+ cache_backing[page, row * 256 : (row + 1) * 256], got_scales
535
+ )
536
+ step = scale.float().repeat_interleave(16)
537
+ assert ((decoded - rotated).abs() <= step).all()
538
+ untouched[page, row * 256 : (row + 1) * 256] = cache_backing[
539
+ page, row * 256 : (row + 1) * 256
540
+ ]
541
+ untouched[page, cache_block * 256 + row * 32 :][:32] = got_scales
542
+ # Rows the kernel must not have touched, page padding included.
543
+ torch.testing.assert_close(cache_backing, untouched, rtol=0, atol=0)
544
+
545
+
546
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only")
547
+def test_v41_nvfp4_gather_matches_insert():
548
+ """The NVFP4 gather dequantizes exactly what the insert kernel wrote."""
549
+ from vllm.models.deepseek_v41.common.ops import dequantize_and_gather_k_cache
550
+ from vllm.models.deepseek_v41.common.ops.fused_compress_quant_cache import (
551
+ rope_quant_insert,
552
+ )
553
+
554
+ torch.manual_seed(17)
555
+ device = "cuda"
556
+ block_size = 64
557
+ num_tokens = 70
558
+ num_blocks = 4
559
+ page_bytes = math.ceil(block_size * 288 / 512) * 512
560
+
561
+ positions = torch.arange(num_tokens, dtype=torch.int64, device=device)
562
+ latent = torch.randn(num_tokens, 512, dtype=torch.bfloat16, device=device)
563
+ angles = torch.randn(128, 32, device=device)
564
+ cos_sin = torch.cat((angles.cos(), angles.sin()), dim=-1)
565
+ backing = torch.zeros(num_blocks, page_bytes, dtype=torch.uint8, device=device)
566
+ cache = backing.as_strided((num_blocks, block_size, 288), (page_bytes, 288, 1))
567
+ slots = torch.arange(num_tokens, dtype=torch.int64, device=device)
568
+ rope_quant_insert(latent, positions, cos_sin, cache, slots, 1)
569
+
570
+ out = torch.zeros(1, num_tokens, 512, dtype=torch.bfloat16, device=device)
571
+ dequantize_and_gather_k_cache(
572
+ out,
573
+ cache,
574
+ seq_lens=torch.tensor([num_tokens], dtype=torch.int32, device=device),
575
+ gather_lens=None,
576
+ block_table=torch.arange(num_blocks, dtype=torch.int32, device=device).view(
577
+ 1, -1
578
+ ),
579
+ block_size=block_size,
580
+ offset=0,
581
+ )
582
+ for t in (0, 1, block_size, num_tokens - 1):
583
+ page, row = divmod(t, block_size)
584
+ expected = _decode_nvfp4_row(
585
+ backing[page, row * 256 : (row + 1) * 256],
586
+ backing[page, block_size * 256 + row * 32 :][:32],
587
+ )
588
+ torch.testing.assert_close(out[0, t].float(), expected, rtol=0, atol=0)
589
+
590
+
591
@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only")
592
@pytest.mark.parametrize("compress_ratio", [1, 2])
593
@pytest.mark.parametrize("store_fp8", [False, True])
tests/kernels/test_dsv41_mega_attn_layouts.py
new
+71
@@ -0,0 +1,71 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Layout plumbing for DeepSeek V4.1 mega attention.
4
+
5
+The kernel's Q and O layouts are produced by permuting ``wq_b`` rows and
6
+``wo_a`` columns once at load, so what has to hold is that the permuted GEMMs
7
+agree with a reference that permutes the activation instead.
8
+"""
9
+
10
+import torch
11
+
12
+from vllm.models.deepseek_v41.common.ops.fused_layout import (
13
+ WV_GROUP_SIZE,
14
+ o_fused_chunk_permutation,
15
+ o_fused_permutation,
16
+ permute_q_to_fused,
17
+ permute_wo_a_,
18
+ permute_wq_b_,
19
+ q_fused_permutation,
20
+)
21
+
22
+HEAD_DIM = 512
23
+
24
+
25
+def test_permuted_wq_b_gemm_matches_permuted_activation():
26
+ """Permuting wq_b's rows makes the Q GEMM emit the fused layout directly."""
27
+ torch.manual_seed(0)
28
+ num_heads, q_lora_rank, num_tokens = 16, 64, 5
29
+ weight = torch.randn(num_heads * HEAD_DIM, q_lora_rank)
30
+ # One scale row per weight row, as an MXFP8 wq_b shard carries.
31
+ scale = torch.randint(
32
+ 0, 256, (num_heads * HEAD_DIM, q_lora_rank // 32), dtype=torch.uint8
33
+ )
34
+ qr = torch.randn(num_tokens, q_lora_rank)
35
+ orig_weight, orig_scale = weight.clone(), scale.clone()
36
+
37
+ standard = (qr @ weight.T).view(num_tokens, num_heads, HEAD_DIM)
38
+ expected = permute_q_to_fused(standard)
39
+
40
+ permute_wq_b_(weight, scale, num_heads)
41
+ fused = (qr @ weight.T).view(num_tokens, num_heads, HEAD_DIM)
42
+ torch.testing.assert_close(fused, expected)
43
+ # The scale has to follow its row, or dequant reads another row's scale.
44
+ perm = q_fused_permutation(num_heads, HEAD_DIM)
45
+ torch.testing.assert_close(weight, orig_weight[perm])
46
+ torch.testing.assert_close(scale, orig_scale[perm])
47
+
48
+
49
+def test_permuted_wo_a_consumes_fused_output():
50
+ """Permuting wo_a's input columns lets it read the kernel's O layout."""
51
+ torch.manual_seed(1)
52
+ out_features, num_tokens = 32, 4
53
+ in_features = WV_GROUP_SIZE * HEAD_DIM
54
+ weight = torch.randn(out_features, in_features)
55
+ scale = torch.randint(0, 256, (out_features, in_features // 32), dtype=torch.uint8)
56
+ standard_o = torch.randn(num_tokens, in_features)
57
+ orig_weight, orig_scale = weight.clone(), scale.clone()
58
+
59
+ expected = standard_o @ weight.T
60
+ # The kernel emits the same values in the fused chunk order.
61
+ perm = o_fused_permutation(WV_GROUP_SIZE, HEAD_DIM)
62
+ fused_o = standard_o[:, perm]
63
+
64
+ permute_wo_a_(weight, scale, WV_GROUP_SIZE)
65
+ torch.testing.assert_close(weight, orig_weight[:, perm])
66
+ torch.testing.assert_close(
67
+ scale, orig_scale[:, o_fused_chunk_permutation(WV_GROUP_SIZE, HEAD_DIM)]
68
+ )
69
+ # Unlike wq_b, this permutation sits inside the 4096-term reduction, so the
70
+ # summation order changes and the result differs in the last few ulps.
71
+ torch.testing.assert_close(fused_o @ weight.T, expected, rtol=1e-3, atol=1e-3)
tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
+183
@@ -522,6 +522,189 @@ def test_kv_path_mxfp8_record_matches_reference(num_tokens: int, block_size: int
522
)
523
524
525
+# ── Test 2d: mega-attention Q layout (is_q_interleaved) ──────────────────────
526
+
527
+Q_CHUNK = 16
528
+NUM_Q_CHUNKS = HEAD_DIM // Q_CHUNK # 32
529
+
530
+
531
+@pytest.mark.parametrize("num_tokens", [1, 17, 2048])
532
+@pytest.mark.parametrize("n_heads,q_head_padded", [(16, 64), (64, 128), (64, 0)])
533
+def test_q_interleaved_pads_without_rope(
534
+ num_tokens: int, n_heads: int, q_head_padded: int
535
+):
536
+ """``is_q_interleaved`` with no norm or RoPE only zero-pads Q.
537
+
538
+ The mega-attention kernel norms and rotates Q itself, so every live
539
+ element must survive bit-exact while the padding heads -- which interleave
540
+ with the live ones inside each of the 32 head-dim chunks -- come back
541
+ zero. ``q_head_padded=0`` drops the Q pass entirely and returns nothing,
542
+ for a shard already at the kernel's head count.
543
+ """
544
+ torch.manual_seed(5)
545
+ device = "cuda"
546
+ dtype = torch.bfloat16
547
+ block_size = 64
548
+
549
+ q = torch.randn(num_tokens, n_heads, HEAD_DIM, dtype=dtype, device=device)
550
+ kv = torch.randn(num_tokens, HEAD_DIM, dtype=dtype, device=device)
551
+ positions = torch.arange(num_tokens, dtype=torch.int64, device=device)
552
+ cos_sin_cache = make_cos_sin_cache(4096, ROPE_DIM, torch.float32, device)
553
+ num_blocks = (num_tokens + block_size - 1) // block_size + 1
554
+ slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device)
555
+ k_cache = torch.zeros(
556
+ num_blocks, block_size * V41_HEAD_BYTES, dtype=torch.uint8, device=device
557
+ )
558
+
559
+ q_out = torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
560
+ q,
561
+ kv,
562
+ k_cache,
563
+ slot_mapping,
564
+ positions,
565
+ cos_sin_cache,
566
+ q_head_padded,
567
+ 1e-6,
568
+ block_size,
569
+ False, # apply_q_norm
570
+ True, # kv_mxfp8
571
+ False, # apply_q_rope
572
+ True, # is_q_interleaved
573
+ )
574
+
575
+ # The KV half is unaffected by the Q mode: it must still be inserted.
576
+ assert k_cache.any()
577
+
578
+ if q_head_padded == 0:
579
+ assert q_out.numel() == 0
580
+ return
581
+ assert q_out.shape == (num_tokens, q_head_padded, HEAD_DIM)
582
+ got = q_out.view(num_tokens, NUM_Q_CHUNKS, q_head_padded, Q_CHUNK)
583
+ want = q.view(num_tokens, NUM_Q_CHUNKS, n_heads, Q_CHUNK)
584
+ torch.testing.assert_close(got[:, :, :n_heads], want, rtol=0, atol=0)
585
+ assert (got[:, :, n_heads:] == 0).all()
586
+
587
+
588
+@pytest.mark.parametrize("n_heads", [8, 16, 32])
589
+def test_q_interleaved_survives_graph_replay(n_heads: int):
590
+ """A captured fused insert refills its padded Q buffer on every replay.
591
+
592
+ The op allocates ``q_out`` itself, so the tensor a capture hands back has
593
+ to keep receiving the new padding when the graph is replayed over changed
594
+ input -- that is what lets the layer prepare Q inside the captured region
595
+ rather than in ``forward_mqa``.
596
+ """
597
+ torch.manual_seed(6)
598
+ device = "cuda"
599
+ dtype = torch.bfloat16
600
+ block_size, num_tokens, q_head_padded = 64, 7, 64
601
+
602
+ q = torch.randn(num_tokens, n_heads, HEAD_DIM, dtype=dtype, device=device)
603
+ kv = torch.randn(num_tokens, HEAD_DIM, dtype=dtype, device=device)
604
+ positions = torch.arange(num_tokens, dtype=torch.int64, device=device)
605
+ cos_sin_cache = make_cos_sin_cache(4096, ROPE_DIM, torch.float32, device)
606
+ num_blocks = (num_tokens + block_size - 1) // block_size + 1
607
+ slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device)
608
+ k_cache = torch.zeros(
609
+ num_blocks, block_size * V41_HEAD_BYTES, dtype=torch.uint8, device=device
610
+ )
611
+
612
+ def prepare():
613
+ return torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
614
+ q,
615
+ kv,
616
+ k_cache,
617
+ slot_mapping,
618
+ positions,
619
+ cos_sin_cache,
620
+ q_head_padded,
621
+ 1e-6,
622
+ block_size,
623
+ False, # apply_q_norm
624
+ True, # kv_mxfp8
625
+ False, # apply_q_rope
626
+ True, # is_q_interleaved
627
+ )
628
+
629
+ # Capture takes its warmup on a side stream.
630
+ side = torch.cuda.Stream()
631
+ side.wait_stream(torch.cuda.current_stream())
632
+ with torch.cuda.stream(side):
633
+ for _ in range(3):
634
+ prepare()
635
+ torch.cuda.current_stream().wait_stream(side)
636
+ graph = torch.cuda.CUDAGraph()
637
+ with torch.cuda.graph(graph):
638
+ q_out = prepare()
639
+
640
+ for factor in (0.25, -0.75):
641
+ q.copy_(torch.randn_like(q) * factor)
642
+ graph.replay()
643
+ torch.accelerator.synchronize()
644
+ got = q_out.view(num_tokens, NUM_Q_CHUNKS, q_head_padded, Q_CHUNK)
645
+ want = q.view(num_tokens, NUM_Q_CHUNKS, n_heads, Q_CHUNK)
646
+ torch.testing.assert_close(got[:, :, :n_heads], want, rtol=0, atol=0)
647
+ assert (got[:, :, n_heads:] == 0).all()
648
+
649
+
650
+@pytest.mark.parametrize("n_heads", [16, 64])
651
+@pytest.mark.parametrize("apply_q_norm", [False, True])
652
+@pytest.mark.parametrize("apply_q_rope", [False, True])
653
+def test_q_interleaved_is_orthogonal_to_norm_and_rope(
654
+ n_heads: int, apply_q_norm: bool, apply_q_rope: bool
655
+):
656
+ """The Q layout composes with whatever norm and RoPE it is given.
657
+
658
+ A lane owns the same 512-dim slice of its head in either layout -- only
659
+ the address differs -- so the head-major result permuted into the
660
+ interleaved layout must equal running the interleaved path on a permuted
661
+ input, bit for bit, for every norm/RoPE combination.
662
+ """
663
+ from vllm.models.deepseek_v41.common.ops.fused_layout import permute_q_to_fused
664
+
665
+ torch.manual_seed(7)
666
+ device = "cuda"
667
+ dtype = torch.bfloat16
668
+ block_size, num_tokens, q_head_padded = 64, 5, 64
669
+
670
+ q = torch.randn(num_tokens, n_heads, HEAD_DIM, dtype=dtype, device=device)
671
+ kv = torch.randn(num_tokens, HEAD_DIM, dtype=dtype, device=device)
672
+ positions = torch.arange(num_tokens, dtype=torch.int64, device=device)
673
+ cos_sin_cache = make_cos_sin_cache(4096, ROPE_DIM, torch.float32, device)
674
+ num_blocks = (num_tokens + block_size - 1) // block_size + 1
675
+ slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device)
676
+
677
+ def run(q_in, is_q_interleaved):
678
+ cache = torch.zeros(
679
+ num_blocks, block_size * V41_HEAD_BYTES, dtype=torch.uint8, device=device
680
+ )
681
+ out = torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
682
+ q_in,
683
+ kv,
684
+ cache,
685
+ slot_mapping,
686
+ positions,
687
+ cos_sin_cache,
688
+ q_head_padded,
689
+ 1e-6,
690
+ block_size,
691
+ apply_q_norm,
692
+ True, # kv_mxfp8
693
+ apply_q_rope,
694
+ is_q_interleaved,
695
+ )
696
+ return out, cache
697
+
698
+ head_major, cache_major = run(q, False)
699
+ interleaved, cache_inter = run(permute_q_to_fused(q), True)
700
+
701
+ torch.testing.assert_close(
702
+ interleaved, permute_q_to_fused(head_major), rtol=0, atol=0
703
+ )
704
+ # The Q layout must not leak into the KV half.
705
+ torch.testing.assert_close(cache_inter, cache_major, rtol=0, atol=0)
706
+
707
+
708
# ── Test 2b: DP padding (slot_mapping shorter than q/kv) ─────────────────────
709
710
vllm/models/deepseek_v41/amd/rocm.py
+2
-1
@@ -636,7 +636,8 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
636
transpose_scale=False,
637
)
638
639
- def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
639
+ def _o_proj(self, attn_out: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
640
+ o = attn_out[:, : self.n_local_heads, :]
641
# ROCm BF16 reference wo_a path (inverse RoPE + einsum) + wo_b.
642
z = rocm_inv_rope_einsum(
643
self.rotary_emb,
vllm/models/deepseek_v41/attention.py
+150
-54
@@ -44,6 +44,7 @@ from vllm.config import (
44
VllmConfig,
45
get_current_vllm_config,
46
)
47
+from vllm.config.cache import CacheDType
48
from vllm.distributed import get_tensor_model_parallel_world_size
49
from vllm.forward_context import get_forward_context
50
from vllm.logger import init_logger
@@ -127,34 +128,40 @@ def _use_v41_mxfp8_kv_record() -> bool:
128
129
def _resolve_dsv4_kv_cache_dtype(
130
use_fp8_ds_mla_layout: bool,
130
- kv_cache_dtype: str,
131
+ kv_cache_dtype: CacheDType,
132
cache_config: CacheConfig | None,
132
-) -> tuple[str, torch.dtype]:
133
+ packed_kv_cache_dtype: CacheDType = "fp8_ds_mla",
134
+) -> tuple[CacheDType, torch.dtype]:
135
"""Map ``(layout, --kv-cache-dtype)`` to ``(cache_dtype_str, torch_dtype)``.
136
137
Both layouts are paged; they differ in the per-token block format. The
136
- ``fp8_ds_mla`` format is UE8M0 block-scaled fp8 packed as ``uint8`` (the
137
- canonical ``fp8_ds_mla`` string is written back onto ``cache_config`` so the
138
- page-size specs pick the 576B per-token slot). Plain-row backends store each
139
- token's KV row in its element dtype: bf16 or per-tensor FP8 E4M3.
138
+ packed formats are ``uint8``-backed: ``fp8_ds_mla`` is UE8M0 block-scaled
139
+ fp8 throughout, ``nvfp4_ds_mla`` keeps that sliding-window record and
140
+ stores the compressed cache as NVFP4. An unspecific ``--kv-cache-dtype``
141
+ (``auto`` / ``fp8``) resolves to ``packed_kv_cache_dtype``, the record this
142
+ layer's kernel prefers, and the canonical string is written back onto
143
+ ``cache_config`` so the page-size specs pick the right per-token slot.
144
+ Plain-row backends store each token's KV row in its element dtype: bf16 or
145
+ per-tensor FP8 E4M3.
146
"""
147
if use_fp8_ds_mla_layout:
142
- # fp8_ds_mla block format: UE8M0 block-scaled fp8 packed as uint8.
143
- if kv_cache_dtype == "auto":
144
- kv_cache_dtype = "fp8"
145
- if not kv_cache_dtype.startswith("fp8"):
148
+ if kv_cache_dtype in ("auto", "fp8"):
149
+ kv_cache_dtype = packed_kv_cache_dtype
150
+ elif not kv_cache_dtype.endswith("_ds_mla"):
151
raise ValueError(
147
- "DeepseekV4 fp8_ds_mla layout only supports fp8 "
148
- f"kv-cache, got {kv_cache_dtype}. Please set "
149
- "`--kv-cache-dtype fp8` or select a backend that supports "
150
- "bfloat16 KV cache."
152
+ "DeepseekV4 packed KV layouts only support fp8 kv-cache, got "
153
+ f"{kv_cache_dtype}. Please set `--kv-cache-dtype fp8` or "
154
+ "select a backend that supports bfloat16 KV cache."
155
)
152
- if kv_cache_dtype != "fp8_ds_mla":
153
- if cache_config is not None:
154
- cache_config.cache_dtype = "fp8_ds_mla"
155
- kv_cache_dtype = "fp8_ds_mla"
156
- logger.info_once("Using DeepSeek's fp8_ds_mla KV cache format.")
156
+ if cache_config is not None and cache_config.cache_dtype != kv_cache_dtype:
157
+ cache_config.cache_dtype = kv_cache_dtype
158
+ logger.info_once("Using DeepSeek's %s KV cache format.", kv_cache_dtype)
159
return kv_cache_dtype, torch.uint8
160
+ if kv_cache_dtype.endswith("_ds_mla"):
161
+ raise ValueError(
162
+ f"{kv_cache_dtype} is a packed FlashMLA DeepSeek V4.1 KV cache "
163
+ "format; the selected backend stores plain KV rows."
164
+ )
165
166
# Plain bf16 / per-tensor fp8 KV row (FlashInfer).
167
if kv_cache_dtype.startswith("fp8"):
@@ -189,6 +196,28 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
196
# path to pre-reserve that workspace.
197
PREFILL_CHUNK_SIZE: ClassVar[int] = 4
198
199
+ # ---- attention-interface contract, declared by the platform subclass ----
200
+
201
+ @property
202
+ def accepts_unnormed_unroped_query(self) -> bool:
203
+ """Whether ``forward_mqa``'s ``q`` is the raw ``wq_b`` output.
204
+
205
+ True when the attention kernel applies the Q norm and RoPE itself, and
206
+ reads Q in its own chunk-interleaved layout, so the layer only
207
+ zero-pads Q to ``padded_heads`` and inserts KV before calling it.
208
+ """
209
+ return False
210
+
211
+ @property
212
+ def packed_kv_cache_dtype(self) -> CacheDType:
213
+ """The packed KV record this layer's kernel prefers.
214
+
215
+ What an unspecific ``--kv-cache-dtype`` (``auto`` / ``fp8``) resolves
216
+ to. Mega attention overrides it: its kernel is the one that can read
217
+ an NVFP4 compressed cache.
218
+ """
219
+ return "fp8_ds_mla"
220
+
221
@classmethod
222
@abstractmethod
223
def get_padded_num_q_heads(cls, num_heads: int) -> int:
@@ -213,8 +242,19 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
242
raise NotImplementedError
243
244
@abstractmethod
216
- def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
217
- """Inverse-RoPE + wo_a + wo_b output projection (platform-specific)."""
245
+ def _o_proj(
246
+ self,
247
+ attn_out: "torch.Tensor | QuantizedActivation",
248
+ positions: torch.Tensor,
249
+ ) -> torch.Tensor:
250
+ """Project whatever ``_alloc_attn_out`` produced through wo_a and wo_b.
251
+
252
+ Takes the buffer whole, so each layer owns the shape it allocated: the
253
+ bf16 layers slice off their padding heads and apply the inverse RoPE,
254
+ while a layer whose attention kernel already did the inverse RoPE and
255
+ the FP8 cast gets a QuantizedActivation and has only wo_a and wo_b
256
+ left.
257
+ """
258
raise NotImplementedError
259
260
def _uses_fp8_ds_mla_layout(self) -> bool:
@@ -459,11 +499,27 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
499
# Resolve the kv-cache dtype from this backend's block format. The same
500
# resolution drives the SWA cache tensor dtype below.
501
self.kv_cache_dtype, self.kv_cache_torch_dtype = _resolve_dsv4_kv_cache_dtype(
462
- self._uses_fp8_ds_mla_layout(), cache_config.cache_dtype, cache_config
502
+ self._uses_fp8_ds_mla_layout(),
503
+ cache_config.cache_dtype,
504
+ cache_config,
505
+ self.packed_kv_cache_dtype,
506
)
507
self.kv_mxfp8 = _use_v41_mxfp8_kv_record()
465
- self.kv_bytes_per_token = 528 if self.kv_mxfp8 else 584
508
+ self.swa_bytes_per_token = 528 if self.kv_mxfp8 else 584
509
+ # nvfp4_ds_mla keeps the MXFP8 sliding-window record and stores the
510
+ # compressed cache as NVFP4 (256 B of e2m1 pairs + 32 e4m3 scales).
511
+ self.compressed_bytes_per_token = (
512
+ 288 if self.kv_cache_dtype == "nvfp4_ds_mla" else self.swa_bytes_per_token
513
+ )
514
+ # One alignment for every page in the block: the block stride is their
515
+ # sum, and 512 satisfies both TMA strides in play (512 for the V4.1
516
+ # fp8 record, 256 for NVFP4).
517
self.kv_page_alignment = 512 if self.kv_mxfp8 else 576
518
+ if self.kv_cache_dtype == "nvfp4_ds_mla" and not self.kv_mxfp8:
519
+ raise ValueError(
520
+ "nvfp4_ds_mla needs the V4.1 KV records, which FlashMLA "
521
+ "decodes only on SM100."
522
+ )
523
524
self.swa_cache_layer = DeepseekV4SWACache(
525
head_dim=self.head_dim,
@@ -473,7 +529,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
529
cache_config=cache_config,
530
backend_cls=self.swa_backend_cls,
531
block_size=32,
476
- packed_bytes_per_token=self.kv_bytes_per_token,
532
+ packed_bytes_per_token=self.swa_bytes_per_token,
533
packed_page_alignment=self.kv_page_alignment,
534
)
535
@@ -553,8 +609,12 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
609
block_size=self.swa_cache_layer.block_size,
610
)
611
612
+ # Every backend that gathers a chunk's KV through
613
+ # combine_topk_swa_indices needs its Triton kernel warmed, mega
614
+ # attention included -- it calls it from _forward_prefill_mega.
615
if self.backend_cls.get_name() in (
616
"FLASHMLA_SPARSE_DSV41",
617
+ "FLASHMLA_MEGA_ATTN_DSV41",
618
"ROCM_FLASHMLA_SPARSE_DSV4",
619
):
620
from vllm.models.deepseek_v41.common.ops.cache_utils import (
@@ -569,14 +629,10 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
629
hidden_states: torch.Tensor,
630
llama_4_scaling: torch.Tensor | None = None,
631
) -> torch.Tensor:
572
- # Pre-allocate attention output with FlashMLA-padded head count.
573
- # The op writes into `o_padded`; we slice to n_local_heads after.
574
- num_tokens = hidden_states.shape[0]
575
- o_padded = torch.empty(
576
- (num_tokens, self.padded_heads, self.head_dim),
577
- dtype=hidden_states.dtype,
578
- device=hidden_states.device,
579
- )
632
+ # The eager attention region writes into a caller-owned buffer
633
+ # (breakable_cudagraph needs in-place outputs); its shape and how it is
634
+ # projected afterwards follow the interface contract above.
635
+ attn_out = self._alloc_attn_out(hidden_states.shape[0], hidden_states)
636
637
# Keep the attention input preparation in the captured graph. Only the
638
# sparse indexer and MLA attention run in the eager break below.
@@ -593,12 +649,19 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
649
kv_score,
650
indexer_weights,
651
positions,
596
- o_padded,
652
+ attn_out,
653
)
598
- o = o_padded[:, : self.n_local_heads, :]
654
+ return self._o_proj(attn_out, positions)
655
600
- # Inverse-RoPE + wo_a + wo_b output projection (platform-specific).
601
- return self._o_proj(o, positions)
656
+ def _alloc_attn_out(
657
+ self, num_tokens: int, hidden_states: torch.Tensor
658
+ ) -> "torch.Tensor | QuantizedActivation":
659
+ """The buffer ``forward_mqa`` fills, per the interface contract."""
660
+ return torch.empty(
661
+ (num_tokens, self.padded_heads, self.head_dim),
662
+ dtype=hidden_states.dtype,
663
+ device=hidden_states.device,
664
+ )
665
666
@cached_property
667
def _can_fuse_query_quant(self) -> bool:
@@ -652,7 +715,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
715
kv_score: torch.Tensor,
716
indexer_weights: torch.Tensor,
717
positions: torch.Tensor,
655
- o_padded: torch.Tensor,
718
+ attn_out: "torch.Tensor | QuantizedActivation",
719
) -> None:
720
"""Wide eager region: the whole of ``_prepare_and_attn`` runs eagerly.
721
@@ -667,7 +730,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
730
kv_score,
731
indexer_weights,
732
positions,
670
- o_padded,
733
+ attn_out,
734
)
735
736
def _prepare_and_attn(
@@ -679,7 +742,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
742
kv_score: torch.Tensor,
743
indexer_weights: torch.Tensor,
744
positions: torch.Tensor,
682
- o_padded: torch.Tensor,
745
+ attn_out: "torch.Tensor | QuantizedActivation",
746
) -> None:
747
"""Attention input preparation followed by the sparse indexer and MLA.
748
@@ -751,7 +814,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
814
q,
815
kv,
816
positions,
754
- o_padded,
817
+ attn_out,
818
)
819
820
def _fused_wqa_wkv_gemm(self, hidden_states: torch.Tensor) -> torch.Tensor:
@@ -858,16 +921,28 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
921
dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None
922
),
923
) -> torch.Tensor:
924
+ """Ready ``q`` for the attention kernel and publish this step's KV.
925
+
926
+ One launch does both. With ``accepts_unnormed_unroped_query`` the
927
+ attention kernel norms and rotates Q itself and reads it in its own
928
+ chunk-interleaved layout, so the Q half of the launch is a zero-pad to
929
+ ``padded_heads`` -- and nothing at all once the shard is that wide.
930
+ """
931
if not isinstance(attn_metadata, dict):
932
# Profile run: kernel doesn't fire; produce a padded tensor so
933
# downstream FlashMLA gets the right shape.
864
- if self.n_local_heads < self.padded_heads:
865
- return F.pad(
866
- q,
867
- (0, 0, 0, self.padded_heads - self.n_local_heads),
868
- value=0.0,
869
- )
870
- return q
934
+ if self.n_local_heads >= self.padded_heads:
935
+ return q
936
+ if self.accepts_unnormed_unroped_query:
937
+ # Padding heads sit at the tail of every head-dim chunk in
938
+ # that layout, so no head-major pad of `q` reproduces it --
939
+ # and nothing reads it on a profile run.
940
+ return q.new_zeros((q.shape[0], self.padded_heads, q.shape[2]))
941
+ return F.pad(
942
+ q,
943
+ (0, 0, 0, self.padded_heads - self.n_local_heads),
944
+ value=0.0,
945
+ )
946
947
swa_metadata = cast(
948
"DeepseekSparseSWAMetadata | None",
@@ -886,23 +961,38 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
961
if cache_dtype == torch.uint8:
962
# fp8_ds_mla UE8M0 paged path. Horizontally fused:
963
# Q side: GPT-J RoPE, zero-filling the padding head slots; the
889
- # kernel allocates and returns the padded q tensor.
964
+ # kernel allocates and returns the padded q tensor. An
965
+ # interleaved Q skips the RoPE its attention kernel owns
966
+ # and keeps only the pad, which q_head_padded=0 drops
967
+ # too once the shard is already padded_heads wide.
968
# KV side: GPT-J RoPE + UE8M0 FP8 quant + paged cache insert.
969
swa_kv_cache_2d = swa_kv_cache.view(swa_kv_cache.shape[0], -1)
892
- return torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
970
+ pad_to = (
971
+ 0
972
+ if self.accepts_unnormed_unroped_query
973
+ and self.n_local_heads == self.padded_heads
974
+ else self.padded_heads
975
+ )
976
+ q_padded = torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
977
q,
978
kv,
979
swa_kv_cache_2d,
980
swa_metadata.slot_mapping,
981
positions,
982
cos_sin_cache,
899
- self.padded_heads,
983
+ pad_to,
984
self.eps,
985
swa_metadata.block_size,
902
- False,
986
+ False, # apply_q_norm: qr is normed before wq_b
987
self.kv_mxfp8,
988
+ not self.accepts_unnormed_unroped_query, # apply_q_rope
989
+ self.accepts_unnormed_unroped_query, # is_q_interleaved
990
)
991
+ return q if pad_to == 0 else q_padded
992
993
+ assert not self.accepts_unnormed_unroped_query, (
994
+ "the chunk-interleaved Q layout only pairs with a packed KV record"
995
+ )
996
# Plain-row path: the [num_blocks, block_size, 512] cache stores the KV
997
# row in its element dtype (no Q padding). bf16 rewrites q in place;
998
# per-tensor fp8 writes a separately-allocated fp8 q and quantizes the
@@ -959,7 +1049,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
1049
# fp8_ds_mla is a UE8M0 block-scaled uint8 layout whose page rounds up
1050
# to the decode kernel's TMA stride; plain bf16 / per-tensor fp8 rows
1051
# use natural element-size pages.
962
- uses_fp8_ds_mla_layout = self.kv_cache_dtype == "fp8_ds_mla"
1052
+ uses_fp8_ds_mla_layout = self.kv_cache_dtype in ("fp8_ds_mla", "nvfp4_ds_mla")
1053
return MLAAttentionSpec(
1054
block_size=vllm_config.cache_config.block_size,
1055
num_kv_heads=1,
@@ -972,7 +1062,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
1062
kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),
1063
# Packed record width; head_size stays semantic (512).
1064
state_content_bytes=(
975
- self.kv_bytes_per_token if uses_fp8_ds_mla_layout else None
1065
+ self.compressed_bytes_per_token if uses_fp8_ds_mla_layout else None
1066
),
1067
)
1068
@@ -1016,7 +1106,13 @@ class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase):
1106
def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
1107
# head_dim already carries the fp8 scale padding
1108
# tokens_per_state=1 for V3.2, >1 for DeepseekV4; same cache layout.
1019
- uses_fp8_ds_mla_layout = vllm_config.cache_config.cache_dtype == "fp8_ds_mla"
1109
+ # nvfp4_ds_mla is packed too: its compressed record is NVFP4 but the
1110
+ # sliding-window record stays the V4.1 MXFP8 one, so the indexer page
1111
+ # takes the same alignment either way.
1112
+ uses_fp8_ds_mla_layout = vllm_config.cache_config.cache_dtype in (
1113
+ "fp8_ds_mla",
1114
+ "nvfp4_ds_mla",
1115
+ )
1116
page_alignment = (
1117
576 if uses_fp8_ds_mla_layout and not _use_v41_mxfp8_kv_record() else 512
1118
)
vllm/models/deepseek_v41/common/ops/cache_utils.py
+96
-2
@@ -43,10 +43,16 @@ from vllm.utils.math_utils import next_power_of_2
43
# dims each plus a pad byte.
44
# V4.1 (528 B): all 512 dims as fp8 e4m3 (RoPE included), then 16 UE8M0
45
# scales of 32 dims each. FlashMLA's ``ModelType::V41``.
46
+# V4.1 NVFP4 (288 B): 512 e2m1 values packed two per byte (even element in
47
+# the low nibble), then 32 e4m3 scales of 16 dims each.
48
+# FlashMLA's ``ModelType::V41_FP4``, compressed cache only.
49
V4_BYTES_PER_TOKEN = 584
50
V41_BYTES_PER_TOKEN = 528
51
V41_QUANT_BLOCK = 32
52
V41_NUM_SCALES = 512 // V41_QUANT_BLOCK # 16
53
+V41_NVFP4_BYTES_PER_TOKEN = 288
54
+V41_NVFP4_QUANT_BLOCK = 16
55
+V41_NVFP4_NUM_SCALES = 512 // V41_NVFP4_QUANT_BLOCK # 32
56
57
58
@triton.jit
@@ -496,6 +502,73 @@ def _dequantize_and_gather_k_mxfp8_kernel(
502
tl.store(output_row_ptr + d, dequant.to(tl.bfloat16))
503
504
505
+@triton.jit
506
+def _dequantize_and_gather_k_nvfp4_kernel(
507
+ out_ptr,
508
+ out_stride0,
509
+ out_stride1,
510
+ k_cache_ptr,
511
+ seq_lens_ptr,
512
+ block_table_ptr,
513
+ offset,
514
+ gather_lens_ptr,
515
+ max_blocks_per_seq: tl.constexpr,
516
+ head_dim: tl.constexpr, # 512
517
+ scale_dim: tl.constexpr, # 32
518
+ quant_block: tl.constexpr, # 16
519
+ cache_block_size: tl.constexpr,
520
+ block_stride: tl.constexpr,
521
+):
522
+ """Gather and dequantize V4.1 NVFP4 rows into a bf16 workspace."""
523
+ batch_idx = tl.program_id(0)
524
+ worker_id = tl.program_id(1)
525
+ num_workers = tl.num_programs(1)
526
+
527
+ seq_len = tl.load(seq_lens_ptr + batch_idx)
528
+ if gather_lens_ptr is not None: # noqa: SIM108
529
+ gather_len = tl.load(gather_lens_ptr + batch_idx)
530
+ else:
531
+ gather_len = seq_len
532
+ start_pos = seq_len - gather_len
533
+
534
+ packed_bytes: tl.constexpr = head_dim // 2
535
+ d = tl.arange(0, head_dim)
536
+ for i in range(worker_id, gather_len, num_workers):
537
+ pos = start_pos + i
538
+ block_in_seq = pos // cache_block_size
539
+ pos_in_block = pos % cache_block_size
540
+
541
+ physical_block_idx = tl.load(
542
+ block_table_ptr + batch_idx * max_blocks_per_seq + block_in_seq
543
+ )
544
+ page = k_cache_ptr + physical_block_idx.to(tl.int64) * block_stride
545
+
546
+ packed = tl.load(
547
+ page + pos_in_block * packed_bytes + tl.arange(0, packed_bytes)
548
+ )
549
+ # Even element in the low nibble, odd in the high nibble.
550
+ codes = tl.interleave((packed & 0xF).to(tl.int32), (packed >> 4).to(tl.int32))
551
+ # e2m1 magnitudes: 0, 0.5, 1, 1.5, 2, 3, 4, 6.
552
+ mag_code = codes & 7
553
+ e = (mag_code >> 1).to(tl.float32)
554
+ m = (mag_code & 1).to(tl.float32)
555
+ mag = tl.where(mag_code < 2, m * 0.5, (1.0 + m * 0.5) * tl.exp2(e - 1.0))
556
+ vals = tl.where(codes >= 8, -mag, mag)
557
+
558
+ sf = tl.load(
559
+ page
560
+ + cache_block_size * packed_bytes
561
+ + pos_in_block * scale_dim
562
+ + tl.arange(0, scale_dim)
563
+ )
564
+ scale = sf.to(tl.float8e4nv, bitcast=True).to(tl.float32)
565
+ tiles = tl.reshape(vals, (scale_dim, quant_block))
566
+ dequant = tl.reshape(tiles * tl.reshape(scale, (scale_dim, 1)), (head_dim,))
567
+
568
+ output_row_ptr = out_ptr + batch_idx * out_stride0 + (offset + i) * out_stride1
569
+ tl.store(output_row_ptr + d, dequant.to(tl.bfloat16))
570
+
571
+
572
def dequantize_and_gather_k_cache_triton(
573
# [num_reqs, max_num_tokens, head_size]
574
out: torch.Tensor,
@@ -513,6 +586,25 @@ def dequantize_and_gather_k_cache_triton(
586
) -> None:
587
num_reqs = seq_lens.shape[0]
588
NUM_WORKERS = 128
589
+ if k_cache.shape[-1] == V41_NVFP4_BYTES_PER_TOKEN:
590
+ _dequantize_and_gather_k_nvfp4_kernel[(num_reqs, NUM_WORKERS)](
591
+ out,
592
+ out.stride(0),
593
+ out.stride(1),
594
+ k_cache,
595
+ seq_lens,
596
+ block_table,
597
+ offset,
598
+ gather_lens,
599
+ max_blocks_per_seq=block_table.shape[-1],
600
+ head_dim=512,
601
+ scale_dim=V41_NVFP4_NUM_SCALES,
602
+ quant_block=V41_NVFP4_QUANT_BLOCK,
603
+ cache_block_size=block_size,
604
+ block_stride=k_cache.stride(0),
605
+ )
606
+ return
607
+
608
if k_cache.shape[-1] == V41_BYTES_PER_TOKEN:
609
_dequantize_and_gather_k_mxfp8_kernel[(num_reqs, NUM_WORKERS)](
610
out,
@@ -581,14 +673,16 @@ def dequantize_and_gather_k_cache(
673
) -> None:
674
"""Dequantize and gather a paged DSv4 K cache.
675
584
- The record is read off ``k_cache.shape[-1]``; see the module header.
676
+ The record is read off ``k_cache.shape[-1]``; see the module header. Only
677
+ the fp8 records have a CuteDSL gather, so NVFP4 always takes the Triton
678
+ path.
679
680
``use_fnuz`` MUST match the encoder of the specific cache being read:
681
``False`` for ``compressed_k_cache`` (Triton encoder is OCP everywhere),
682
``current_platform.is_fp8_fnuz()`` for ``swa_k_cache`` (C++ encoder
683
writes FNUZ on gfx942 and OCP on gfx950).
684
"""
591
- if has_cutedsl():
685
+ if has_cutedsl() and k_cache.shape[-1] != V41_NVFP4_BYTES_PER_TOKEN:
686
# lazily import, otherwise some tests fail due to CUDA driver init failure.
687
from vllm.models.deepseek_v4.nvidia.ops.dequant_gather_k_cutedsl import (
688
_DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL,
vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py
+71
-13
@@ -4,6 +4,7 @@
4
5
import torch
6
7
+from vllm.models.deepseek_v4.common.ops.fused_indexer_q import _fp32x2_to_fp4x2
8
from vllm.platforms import current_platform
9
from vllm.triton_utils import tl, triton
10
@@ -232,13 +233,14 @@ def rope_quant_insert(
233
234
The BF16 latent supplies both NoPE quantization and RoPE input. It is read
235
only for valid slots at group boundaries. The cache dtype selects the
235
- layout: ``uint8`` is the fp8_ds_mla paged layout, whose record the
236
- per-token byte width names -- 584 B for V4 (576 value bytes and eight
237
- segregated UE8M0 scale bytes, including one zero padding scale) or 528 B
238
- for V4.1 (512 MXFP8 value bytes covering the RoPE dims too, then 16 UE8M0
239
- scales of 32 dims each). ``bfloat16`` and ``float8_e4m3fn`` are the plain
240
- [448 NoPE | 64 RoPE] rows read by FlashInfer, the latter scaled by the
241
- per-tensor ``fp8_scale``.
236
+ layout: ``uint8`` is a paged FlashMLA layout, whose record the per-token
237
+ byte width names -- 584 B for V4 (576 value bytes and eight segregated
238
+ UE8M0 scale bytes, including one zero padding scale), 528 B for V4.1
239
+ (512 MXFP8 value bytes covering the RoPE dims too, then 16 UE8M0 scales of
240
+ 32 dims each), or 288 B for V4.1 NVFP4 (256 bytes of e2m1 pairs then 32
241
+ e4m3 scales of 16 dims each), which only the compressed cache uses.
242
+ ``bfloat16`` and ``float8_e4m3fn`` are the plain [448 NoPE | 64 RoPE] rows
243
+ read by FlashInfer, the latter scaled by the per-tensor ``fp8_scale``.
244
"""
245
assert compress_ratio in (1, 2)
246
assert latent.shape[1] == 512 and latent.dtype == torch.bfloat16
@@ -249,14 +251,14 @@ def rope_quant_insert(
251
return
252
launch_kwargs = {"launch_pdl": False} if current_platform.is_cuda() else {}
253
if kv_cache.dtype == torch.uint8:
252
- assert kv_cache.shape[-1] in (584, 528), (
254
+ kernel = {
255
+ 584: _rope_quant_insert_kernel,
256
+ 528: _rope_quant_insert_mxfp8_kernel,
257
+ 288: _rope_quant_insert_nvfp4_kernel,
258
+ }.get(kv_cache.shape[-1])
259
+ assert kernel is not None, (
260
f"unsupported paged KV record width {kv_cache.shape[-1]}"
261
)
255
- kernel = (
256
- _rope_quant_insert_mxfp8_kernel
257
- if kv_cache.shape[-1] == 528
258
- else _rope_quant_insert_kernel
259
- )
262
kernel[(num_tokens,)](
263
latent,
264
positions,
@@ -402,6 +404,62 @@ def _rope_quant_insert_mxfp8_kernel(
404
tl.store(scales + tl.arange(0, 16), encoded.to(tl.uint8))
405
406
407
+@triton.jit
408
+def _rope_quant_insert_nvfp4_kernel(
409
+ latent,
410
+ positions,
411
+ cos_sin,
412
+ cache,
413
+ cache_slots,
414
+ COS_STRIDE: tl.constexpr,
415
+ CACHE_STRIDE: tl.constexpr,
416
+ CACHE_BLOCK: tl.constexpr,
417
+ COMPRESS_RATIO: tl.constexpr,
418
+ SANITIZE_CACHE_NANS: tl.constexpr,
419
+):
420
+ """V4.1 NVFP4 record: RoPE, then e2m1 with one e4m3 scale per 16 dims.
421
+
422
+ The scale is ``amax / 6`` (6 is e2m1's largest magnitude) clamped to the
423
+ e4m3 range, with no per-tensor scale on top.
424
+ """
425
+ t = tl.program_id(0)
426
+ slot = tl.load(cache_slots + t)
427
+ if slot < 0:
428
+ return
429
+ position = tl.load(positions + t)
430
+ if (position + 1) % COMPRESS_RATIO != 0:
431
+ return
432
+ d = tl.arange(0, 512)
433
+ normed = tl.load(latent + t.to(tl.int64) * 512 + d).to(tl.float32)
434
+
435
+ # NoPE pairs load (cos, sin) = (1, 0), so the rotation is the identity there.
436
+ even, odd = tl.split(tl.reshape(normed, (256, 2)))
437
+ pair = tl.arange(0, 256) - 224
438
+ cs = cos_sin + (position // COMPRESS_RATIO * COMPRESS_RATIO) * COS_STRIDE
439
+ c = tl.load(cs + tl.maximum(pair, 0), pair >= 0, other=1.0).to(tl.float32)
440
+ s = tl.load(cs + 32 + tl.maximum(pair, 0), pair >= 0, other=0.0).to(tl.float32)
441
+ rotated = tl.interleave(even * c - odd * s, odd * c + even * s)
442
+ if SANITIZE_CACHE_NANS:
443
+ rotated = tl.where(rotated == rotated, rotated, 0.0)
444
+
445
+ tiles = tl.reshape(rotated, (32, 16))
446
+ amax = tl.max(tl.abs(tiles), 1)
447
+ # 2**-9 is the smallest normal e4m3 magnitude; 448 the largest.
448
+ scale = tl.clamp(amax * (1.0 / 6.0), 0.001953125, 448.0).to(tl.float8e4nv)
449
+ # Round-to-nearest division: Triton's default div.full misplaces values
450
+ # that land exactly on an e2m1 tie.
451
+ scaled = tl.math.div_rn(tiles, tl.reshape(scale.to(tl.float32), (32, 1)))
452
+ lo, hi = tl.split(tl.reshape(tl.reshape(scaled, (512,)), (256, 2)))
453
+
454
+ page = cache + (slot // CACHE_BLOCK).to(tl.int64) * CACHE_STRIDE
455
+ tl.store(
456
+ page + (slot % CACHE_BLOCK) * 256 + tl.arange(0, 256),
457
+ _fp32x2_to_fp4x2(lo, hi),
458
+ )
459
+ scales = page + CACHE_BLOCK * 256 + (slot % CACHE_BLOCK) * 32
460
+ tl.store(scales + tl.arange(0, 32), scale.to(tl.uint8, bitcast=True))
461
+
462
+
463
@triton.jit
464
def _rope_plain_insert_kernel(
465
latent,
vllm/models/deepseek_v41/common/ops/fused_layout.py
new
+86
@@ -0,0 +1,86 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Weight permutations for FlashMLA's mega-attention kernel.
4
+
5
+The kernel reads Q with 16-element head-dim chunks interleaved across heads and
6
+writes O with 32-element chunks interleaved across the 8 heads of a ``wo_a``
7
+group. Every permutation here satisfies ``fused = standard[perm]``, and is
8
+applied once to ``wq_b`` rows and ``wo_a`` columns at load time so the
9
+surrounding GEMMs produce and consume the kernel's layouts directly -- no
10
+per-step shuffle.
11
+"""
12
+
13
+import torch
14
+
15
+HEAD_DIM = 512
16
+Q_CHUNK = 16
17
+O_CHUNK = 32
18
+WV_GROUP_SIZE = 8
19
+
20
+
21
+def q_fused_permutation(num_heads: int, head_dim: int = HEAD_DIM) -> torch.Tensor:
22
+ """``fused[(d // 16) * (H * 16) + h * 16 + d % 16] = standard[h * D + d]``."""
23
+ h = torch.arange(num_heads).view(num_heads, 1)
24
+ d = torch.arange(head_dim).view(1, head_dim)
25
+ fused_index = (d // Q_CHUNK) * (num_heads * Q_CHUNK) + h * Q_CHUNK + d % Q_CHUNK
26
+ perm = torch.empty(num_heads * head_dim, dtype=torch.long)
27
+ perm[fused_index.reshape(-1)] = torch.arange(num_heads * head_dim)
28
+ return perm
29
+
30
+
31
+def o_fused_permutation(
32
+ heads_per_group: int = WV_GROUP_SIZE, head_dim: int = HEAD_DIM
33
+) -> torch.Tensor:
34
+ """``fused[(c * G + h) * 32 + j] = standard[h * D + c * 32 + j]`` per group."""
35
+ h = torch.arange(heads_per_group).view(-1, 1, 1)
36
+ c = torch.arange(head_dim // O_CHUNK).view(1, -1, 1)
37
+ j = torch.arange(O_CHUNK).view(1, 1, -1)
38
+ fused_index = (c * heads_per_group + h) * O_CHUNK + j
39
+ perm = torch.empty(heads_per_group * head_dim, dtype=torch.long)
40
+ perm[fused_index.reshape(-1)] = torch.arange(heads_per_group * head_dim)
41
+ return perm
42
+
43
+
44
+def o_fused_chunk_permutation(
45
+ heads_per_group: int = WV_GROUP_SIZE, head_dim: int = HEAD_DIM
46
+) -> torch.Tensor:
47
+ """Per-32-element-chunk form of :func:`o_fused_permutation`, for scales."""
48
+ return o_fused_permutation(heads_per_group, head_dim)[::O_CHUNK] // O_CHUNK
49
+
50
+
51
+def permute_q_to_fused(q: torch.Tensor) -> torch.Tensor:
52
+ """``[N, H, D]`` standard layout -> the same shape in the fused layout."""
53
+ n, h, d = q.shape
54
+ perm = q_fused_permutation(h, d).to(q.device)
55
+ return q.reshape(n, h * d)[:, perm].view(n, h, d)
56
+
57
+
58
+def _bytes_view(t: torch.Tensor) -> torch.Tensor:
59
+ """Byte view of a 1-byte-element tensor, so fp8/ue8m0 can be gathered."""
60
+ return t.view(torch.uint8) if t.element_size() == 1 else t
61
+
62
+
63
+def permute_wq_b_(
64
+ weight: torch.Tensor, weight_scale: torch.Tensor, num_local_heads: int
65
+) -> None:
66
+ """Permute an MXFP8 ``wq_b`` shard's rows and per-row scales, in place."""
67
+ head_dim = weight.shape[0] // num_local_heads
68
+ perm = q_fused_permutation(num_local_heads, head_dim).to(weight.device)
69
+ for t in (weight, weight_scale):
70
+ b = _bytes_view(t)
71
+ b.copy_(b[perm])
72
+
73
+
74
+def permute_wo_a_(
75
+ weight: torch.Tensor,
76
+ weight_scale: torch.Tensor,
77
+ heads_per_group: int = WV_GROUP_SIZE,
78
+) -> None:
79
+ """Permute an MXFP8 ``wo_a`` shard's input columns and scales, in place."""
80
+ head_dim = weight.shape[1] // heads_per_group
81
+ perm = o_fused_permutation(heads_per_group, head_dim).to(weight.device)
82
+ w = _bytes_view(weight)
83
+ w.copy_(w[:, perm])
84
+ chunk_perm = o_fused_chunk_permutation(heads_per_group, head_dim)
85
+ s = _bytes_view(weight_scale)
86
+ s.copy_(s[:, chunk_perm.to(weight.device)])
vllm/models/deepseek_v41/nvidia/dspark.py
+14
@@ -477,8 +477,22 @@ class DSparkDeepseekV4ForCausalLM(nn.Module):
477
for layer in self.model.layers:
478
layer.ffn.finalize_mega_moe_weights()
479
480
+ def _finalize_attn(self) -> None:
481
+ """Run the attention backend's post-load weight step on the draft layers.
482
+
483
+ They are ordinary DeepseekV4DecoderLayers, so they take the engine-wide
484
+ attention backend and owe it whatever the target model owes it -- mega
485
+ attention permutes wq_b / wo_a here and refuses to run without it.
486
+ Idempotent, like the target's.
487
+ """
488
+ for layer in self.model.layers:
489
+ finalize = getattr(layer.attn, "finalize_loaded_weights", None)
490
+ if finalize is not None:
491
+ finalize()
492
+
493
def process_weights_after_loading(self) -> None:
494
self._finalize_moe()
495
+ self._finalize_attn()
496
497
def _remap_dspark_name(self, name: str) -> str | None:
498
"""Map a checkpoint ``mtp.{i}.*`` name to this model's parameter path.
vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py
new
+497
@@ -0,0 +1,497 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""DeepSeek V4.1 attention on FlashMLA's mega-attention kernel (SM100).
4
+
5
+One kernel does Q RoPE, sparse attention, the inverse RoPE of the output and
6
+its FP8 cast, and writes straight into the buffer ``wo_a`` consumes. So the
7
+layer declares ``accepts_unnormed_unroped_query`` -- the kernel RoPEs Q itself,
8
+leaving the fused KV insert only the zero-pad to the kernel's head count -- and
9
+its ``_alloc_attn_out`` / ``_o_proj`` pair speaks QuantizedActivation instead
10
+of bf16, since the output needs no rotation or quantization here.
11
+
12
+``wq_b`` rows and ``wo_a`` columns are permuted once at load so the surrounding
13
+GEMMs speak the kernel's chunk-interleaved layouts directly. A step's prefill
14
+and decode segments write disjoint token ranges of one output buffer pair, so
15
+a single ``wo_a`` einsum covers the whole step.
16
+"""
17
+
18
+from dataclasses import replace
19
+from typing import TYPE_CHECKING, cast
20
+
21
+import torch
22
+
23
+from vllm.config import VllmConfig
24
+from vllm.config.cache import CacheDType
25
+from vllm.forward_context import get_forward_context
26
+from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation
27
+from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic
28
+from vllm.models.deepseek_v41.common.ops import (
29
+ combine_topk_swa_indices,
30
+ compute_global_topk_indices_and_lens,
31
+ dequantize_and_gather_k_cache,
32
+)
33
+from vllm.models.deepseek_v41.common.ops.fused_layout import (
34
+ WV_GROUP_SIZE,
35
+ permute_wo_a_,
36
+ permute_wq_b_,
37
+)
38
+from vllm.models.deepseek_v41.nvidia.flashmla import DeepseekV4FlashMLAAttention
39
+from vllm.models.deepseek_v41.sparse_mla import (
40
+ DeepseekV4FlashMLAMetadata,
41
+ FlashMLAMegaAttnBackend,
42
+)
43
+from vllm.platforms import current_platform
44
+from vllm.utils.deep_gemm import fp8_einsum, get_tma_aligned_size
45
+from vllm.utils.math_utils import round_up
46
+from vllm.v1.attention.ops.flashmla import is_flashmla_sparse_supported
47
+from vllm.v1.worker.workspace import current_workspace_manager
48
+
49
+if TYPE_CHECKING:
50
+ from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
51
+
52
+_HEAD_DIM_V = 512
53
+_ROPE_DIM = 64
54
+# The kernel always emits per-32 scales, whatever num_per_channels says.
55
+_QUANT_GROUP = 32
56
+# Fixed kernel conventions: GPT-J (non-neox) RoPE, TMA-aligned col-major
57
+# scales, round-to-nearest, packed ue8m0 -- i.e. what DeepGEMM's fp8_einsum
58
+# consumes. Q norm is off: the V4.1 layer norms qr before wq_b.
59
+_ROPE_ARGS = (False, _ROPE_DIM)
60
+_SF_ARGS = (_QUANT_GROUP, True, True, True)
61
+
62
+# Decode always goes through the mega kernel: there is no batch size at which
63
+# falling back to the split-KV decode pays, so there is no minimum-token
64
+# threshold and no fallback path to maintain. What decides whether this layer
65
+# is worth selecting at all is TP, not the batch size.
66
+#
67
+# Measured on one GB300 with benchmarks/kernels/benchmark_dsv41_mega_attn.py,
68
+# CUDA graphs on, counting every per-step op each path runs (64-head model,
69
+# topk_swa=128, topk_extra=512, NVFP4 compressed cache), as mega / split-KV
70
+# microseconds:
71
+#
72
+# s_q 1 32 128 256 512
73
+# TP1 (64 live heads) 29/29 31/31 31/40 53/72 96/140
74
+# TP2 (32 live) 29/29 31/30 35/35 61/63 110/119
75
+# TP4 (16 live) 29/29 31/29 34/34 59/58 109/109
76
+# TP8 (8 live) 29/29 31/29 33/33 59/57 107/104
77
+#
78
+# The kernel absorbs the inverse RoPE and the FP8 cast, whose cost scales with
79
+# *live* heads, while its own cost scales with *padded* heads -- always 64,
80
+# since FlashMLA's FP8 decode takes only h_q in {64, 128}. So it wins at TP1
81
+# (up to 1.45x, and TP1 is also where the fused insert skips the Q pad
82
+# entirely) and from TP2 down the two cancel to within a few percent, with a
83
+# fixed ~2 us penalty at small s_q where the padded-head cost dominates.
84
+# Measure before assuming a win at high TP -- and measure under CUDA graphs:
85
+# in eager mode launch overhead swamps both paths and overstates the fused one
86
+# several-fold.
87
+
88
+
89
+def is_flashmla_mega_attn_supported() -> tuple[bool, str | None]:
90
+ """Whether FlashMLA's fused mega-attention kernels are usable here."""
91
+ is_available, maybe_reason = is_flashmla_sparse_supported()
92
+ if not is_available:
93
+ return False, maybe_reason
94
+ if not current_platform.is_device_capability_family(100):
95
+ return False, "FlashMLA mega attention requires sm_10x GPUs."
96
+ if not hasattr(torch.ops._flashmla_C, "fused_norm_rope_attn_rope_cast_decode"):
97
+ return False, "vllm._flashmla_C was built without the mega attention ops."
98
+ if not hasattr(torch.ops._flashmla_C, "permute_q_b_proj"):
99
+ return False, "vllm._flashmla_C predates the mega attention weight permutes."
100
+ return True, None
101
+
102
+
103
+def alloc_mega_attn_output(
104
+ num_tokens: int,
105
+ n_wv_group: int,
106
+ device: torch.device,
107
+) -> QuantizedActivation:
108
+ """Allocate the output buffer pair the mega-attention kernels write into.
109
+
110
+ One pair per forward step: the prefill and decode segments fill disjoint
111
+ token ranges, so a single ``fp8_einsum`` over all ``num_tokens`` can
112
+ consume the result instead of one call per segment.
113
+
114
+ The scale buffer is MN-major -- stride 1 along tokens, head-dim stride
115
+ ``ceil4(num_tokens)`` -- which is both what the kernel requires of a
116
+ caller-provided buffer and what DeepGEMM expects for this N.
117
+ """
118
+ d = WV_GROUP_SIZE * _HEAD_DIM_V
119
+ data = torch.empty(
120
+ (num_tokens, n_wv_group, d), dtype=torch.float8_e4m3fn, device=device
121
+ )
122
+ aligned = get_tma_aligned_size(num_tokens, torch.int32.itemsize)
123
+ scale = torch.empty(
124
+ (n_wv_group, d // (_QUANT_GROUP * 4), aligned),
125
+ dtype=torch.int32,
126
+ device=device,
127
+ ).permute(2, 0, 1)[:num_tokens]
128
+ return QuantizedActivation(
129
+ data=data,
130
+ scale=scale,
131
+ orig_dtype=torch.bfloat16,
132
+ orig_shape=data.shape,
133
+ quant_key=kMxfp8Dynamic,
134
+ )
135
+
136
+
137
+def _token_slice(out: QuantizedActivation, start: int, end: int) -> QuantizedActivation:
138
+ """The ``[start, end)`` token slice of a mega-attention output buffer."""
139
+ data = out.data[start:end]
140
+ return replace(out, data=data, scale=out.scale[start:end], orig_shape=data.shape)
141
+
142
+
143
+class DeepseekV4MegaAttnAttention(DeepseekV4FlashMLAAttention):
144
+ """FlashMLA mega-attention layer for DeepSeek V4.1 (SM100)."""
145
+
146
+ backend_cls = FlashMLAMegaAttnBackend
147
+
148
+ @classmethod
149
+ def is_available_for(cls, vllm_config: VllmConfig) -> bool:
150
+ """Whether this layer can serve the configured model on this device.
151
+
152
+ The kernel is SM100-only and wants ``WV_GROUP_SIZE`` heads per ``wo_a``
153
+ group, which stops holding once TP divides the head count far enough
154
+ (TP16 on a 64-head model leaves 4). ``__init__`` raises on both, so the
155
+ default-backend selector asks here first rather than turning an
156
+ unsupported topology into a startup crash.
157
+ """
158
+ if not is_flashmla_mega_attn_supported()[0]:
159
+ return False
160
+ config = vllm_config.model_config.hf_config
161
+ config = getattr(config, "text_config", config)
162
+ n_heads = getattr(config, "num_attention_heads", 0) or 0
163
+ n_groups = getattr(config, "o_groups", 0) or 0
164
+ tp_size = vllm_config.parallel_config.tensor_parallel_size
165
+ if n_heads % tp_size or n_groups % tp_size:
166
+ return False
167
+ n_local_heads, n_local_groups = n_heads // tp_size, n_groups // tp_size
168
+ if not n_local_groups or n_local_heads % WV_GROUP_SIZE:
169
+ return False
170
+ return n_local_heads // n_local_groups == WV_GROUP_SIZE
171
+
172
+ def __init__(self, vllm_config: VllmConfig, *args, **kwargs) -> None:
173
+ super().__init__(vllm_config, *args, **kwargs)
174
+ if self.n_local_heads % WV_GROUP_SIZE:
175
+ raise ValueError(
176
+ f"{self.prefix}: mega attention needs the local head count "
177
+ f"({self.n_local_heads}) to be a multiple of {WV_GROUP_SIZE}."
178
+ )
179
+ if self.n_local_heads // self.n_local_groups != WV_GROUP_SIZE:
180
+ raise ValueError(
181
+ f"{self.prefix}: mega attention needs {WV_GROUP_SIZE} heads per "
182
+ "wo_a group."
183
+ )
184
+ self.n_wv_group = self.padded_heads // WV_GROUP_SIZE
185
+ self._fused_layouts_ready = False
186
+
187
+ # ---- interface contract ------------------------------------------------
188
+
189
+ @property
190
+ def accepts_unnormed_unroped_query(self) -> bool:
191
+ return True
192
+
193
+ @property
194
+ def packed_kv_cache_dtype(self) -> CacheDType:
195
+ # This kernel is the only one that reads an NVFP4 compressed cache, so
196
+ # it is what an unspecific --kv-cache-dtype resolves to here.
197
+ return "nvfp4_ds_mla"
198
+
199
+ def _alloc_attn_out(
200
+ self, num_tokens: int, hidden_states: torch.Tensor
201
+ ) -> QuantizedActivation:
202
+ return alloc_mega_attn_output(num_tokens, self.n_wv_group, hidden_states.device)
203
+
204
+ def _o_proj(
205
+ self, attn_out: QuantizedActivation, positions: torch.Tensor
206
+ ) -> torch.Tensor:
207
+ """Grouped ``wo_a`` then ``wo_b`` over the kernel's quantized output.
208
+
209
+ The inverse RoPE and the FP8 cast happened inside the attention
210
+ kernel, and ``wo_a`` is permuted for its output layout, so there is
211
+ nothing to rotate or quantize here and ``positions`` is unused. The
212
+ scale already is DeepGEMM's packed-ue8m0 MN-major layout, so the
213
+ einsum consumes the kernel's output with no repacking.
214
+ """
215
+ groups = self.n_local_groups
216
+ z = torch.empty(
217
+ (attn_out.data.shape[0], groups, self.o_lora_rank),
218
+ dtype=torch.bfloat16,
219
+ device=attn_out.data.device,
220
+ )
221
+ fp8_einsum(
222
+ "bhr,hdr->bhd",
223
+ (attn_out.data[:, :groups], attn_out.scale[:, :groups]),
224
+ (self.wo_a.weight, self.wo_a.weight_scale),
225
+ z,
226
+ recipe=self._einsum_recipe,
227
+ )
228
+ return self.wo_b(z.flatten(1))
229
+
230
+ # ---- weights -----------------------------------------------------------
231
+
232
+ def finalize_loaded_weights(self) -> None:
233
+ """Permute wq_b rows / wo_a columns into the kernel's layouts.
234
+
235
+ Idempotent: a second post-load pass must not permute twice.
236
+ """
237
+ if self._fused_layouts_ready:
238
+ return
239
+ permute_wq_b_(
240
+ self.wq_b.weight.data, self.wq_b.weight_scale.data, self.n_local_heads
241
+ )
242
+ permute_wo_a_(
243
+ self.wo_a.weight.data,
244
+ self.wo_a.weight_scale.data,
245
+ self.n_local_heads // self.n_local_groups,
246
+ )
247
+ self._fused_layouts_ready = True
248
+
249
+ # ---- forward -----------------------------------------------------------
250
+
251
+ def forward_mqa(
252
+ self,
253
+ q: torch.Tensor,
254
+ kv: torch.Tensor,
255
+ positions: torch.Tensor,
256
+ output: QuantizedActivation,
257
+ ) -> None:
258
+ if not self._fused_layouts_ready:
259
+ raise RuntimeError(
260
+ f"{self.prefix}: wq_b / wo_a were never permuted for the mega "
261
+ "attention kernel; refusing to run with mismatched layouts."
262
+ )
263
+ attn_metadata = get_forward_context().attn_metadata
264
+ if attn_metadata is None:
265
+ # Warmup dummy run: reserve the prefill workspace, produce zeros.
266
+ self._reserve_prefill_workspace(q)
267
+ output.data.zero_()
268
+ output.scale.zero_()
269
+ return
270
+
271
+ assert isinstance(attn_metadata, dict)
272
+ # q was zero-padded to the kernel's head count by the fused KV insert.
273
+ assert q.shape[1] == self.padded_heads
274
+ flashmla_metadata = cast(
275
+ DeepseekV4FlashMLAMetadata | None,
276
+ attn_metadata.get(self.compressed_cache_prefix)
277
+ if self.compressed_cache_prefix is not None
278
+ else None,
279
+ )
280
+ swa_metadata = cast(
281
+ "DeepseekSparseSWAMetadata", attn_metadata[self.swa_cache_layer.prefix]
282
+ )
283
+ # The kernel takes int32 RoPE positions.
284
+ positions_int32 = positions.to(torch.int32)
285
+ num_decode_tokens = swa_metadata.num_decode_tokens
286
+
287
+ if swa_metadata.num_prefills > 0:
288
+ self._forward_prefill_mega(
289
+ q[num_decode_tokens:],
290
+ positions_int32[num_decode_tokens:],
291
+ flashmla_metadata,
292
+ swa_metadata,
293
+ output,
294
+ num_decode_tokens,
295
+ )
296
+ if swa_metadata.num_decodes > 0:
297
+ self._forward_decode_mega(
298
+ q[:num_decode_tokens],
299
+ positions_int32[:num_decode_tokens],
300
+ flashmla_metadata,
301
+ swa_metadata,
302
+ _token_slice(output, 0, num_decode_tokens),
303
+ )
304
+
305
+ def _reserve_prefill_workspace(self, q: torch.Tensor) -> None:
306
+ swa_only = self.compress_ratio == 0
307
+ n = 0 if swa_only else -(-self.max_model_len // self.compress_ratio)
308
+ m = n + self.window_size + self.max_num_batched_tokens
309
+ if swa_only:
310
+ top_k = 0
311
+ else:
312
+ assert self.topk_indices_buffer is not None
313
+ top_k = self.topk_indices_buffer.shape[-1]
314
+ combined_topk = round_up(top_k + self.window_size + self.max_image_tokens, 128)
315
+ current_workspace_manager().get_simultaneous(
316
+ ((self.PREFILL_CHUNK_SIZE, m, q.shape[-1]), torch.bfloat16),
317
+ ((self.max_num_batched_tokens, combined_topk), torch.int32),
318
+ ((self.max_num_batched_tokens,), torch.int32),
319
+ )
320
+
321
+ def _decode_compressed_kv_and_topk(
322
+ self,
323
+ flashmla_metadata: DeepseekV4FlashMLAMetadata | None,
324
+ swa_metadata: "DeepseekSparseSWAMetadata",
325
+ ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]:
326
+ """The compressed cache and its ``[n, topk]`` slot indices / lengths."""
327
+ if self.compress_ratio == 0:
328
+ return None, None, None
329
+ assert flashmla_metadata is not None
330
+ assert swa_metadata.is_valid_token is not None
331
+ assert self.topk_indices_buffer is not None
332
+ num_decode_tokens = swa_metadata.num_decode_tokens
333
+ indices, lens = compute_global_topk_indices_and_lens(
334
+ self.topk_indices_buffer[:num_decode_tokens],
335
+ swa_metadata.token_to_req_indices,
336
+ flashmla_metadata.block_table[: swa_metadata.num_decodes],
337
+ flashmla_metadata.block_size // self.compress_ratio,
338
+ swa_metadata.is_valid_token[:num_decode_tokens],
339
+ )
340
+ return self._compressed_kv_cache().unsqueeze(-2), indices, lens
341
+
342
+ def _forward_decode_mega(
343
+ self,
344
+ q: torch.Tensor,
345
+ positions_int32: torch.Tensor,
346
+ flashmla_metadata: DeepseekV4FlashMLAMetadata | None,
347
+ swa_metadata: "DeepseekSparseSWAMetadata",
348
+ out: QuantizedActivation,
349
+ ) -> None:
350
+ """Mega attention over the paged quantized caches, in place into ``out``.
351
+
352
+ The caches are ``[num_blocks, page, 1, bytes]``, the record taken from
353
+ ``bytes`` (528 V4.1 fp8, 288 V4.1 NVFP4; NVFP4 only as the compressed
354
+ cache beside an fp8 SWA one), and the indices ``[s_q, topk]`` int32
355
+ slot ids (``block * page + offset``, ``-1`` invalid).
356
+ """
357
+ extra_cache, extra_idx, extra_len = self._decode_compressed_kv_and_topk(
358
+ flashmla_metadata, swa_metadata
359
+ )
360
+ assert swa_metadata.decode_swa_indices is not None
361
+ torch.ops._flashmla_C.fused_norm_rope_attn_rope_cast_decode(
362
+ q,
363
+ self.swa_cache_layer.kv_cache.unsqueeze(-2),
364
+ swa_metadata.decode_swa_indices.view(q.shape[0], -1),
365
+ self.scale,
366
+ _HEAD_DIM_V,
367
+ self.attn_sink,
368
+ swa_metadata.decode_swa_lens,
369
+ extra_cache,
370
+ extra_idx,
371
+ extra_len,
372
+ False, # enable_q_norm
373
+ 0.0, # rms_norm_eps
374
+ positions_int32,
375
+ *_ROPE_ARGS,
376
+ self.rotary_emb.cos_sin_cache,
377
+ self.n_wv_group,
378
+ *_SF_ARGS,
379
+ out.data,
380
+ out.scale,
381
+ )
382
+
383
+ def _forward_prefill_mega(
384
+ self,
385
+ q: torch.Tensor,
386
+ positions_int32: torch.Tensor,
387
+ flashmla_metadata: DeepseekV4FlashMLAMetadata | None,
388
+ swa_metadata: "DeepseekSparseSWAMetadata",
389
+ out: QuantizedActivation,
390
+ token_base: int,
391
+ ) -> None:
392
+ swa_only = self.compress_ratio == 0
393
+ num_decodes = swa_metadata.num_decodes
394
+ num_decode_tokens = swa_metadata.num_decode_tokens
395
+ seq_lens = swa_metadata.prefill_seq_lens
396
+ gather_lens = swa_metadata.prefill_gather_lens
397
+ query_start_loc_cpu = swa_metadata.query_start_loc_cpu
398
+ query_start_loc = swa_metadata.query_start_loc
399
+ assert seq_lens is not None and gather_lens is not None
400
+ assert query_start_loc_cpu is not None and query_start_loc is not None
401
+ prefill_token_base = query_start_loc_cpu[num_decodes]
402
+ assert self.topk_indices_buffer is not None
403
+ topk_indices = self.topk_indices_buffer[num_decode_tokens:][
404
+ : swa_metadata.num_prefill_tokens
405
+ ]
406
+ top_k = 0 if swa_only else topk_indices.shape[-1]
407
+ chunk_plan = swa_metadata.get_prefill_chunk_plan(
408
+ compress_ratio=self.compress_ratio,
409
+ prefill_chunk_size=self.PREFILL_CHUNK_SIZE,
410
+ has_compressed=not swa_only,
411
+ )
412
+ assert chunk_plan, "prefill chunk plan must be non-empty when num_prefills > 0"
413
+ workspace_manager = current_workspace_manager()
414
+ combined_topk = round_up(top_k + self.window_size + self.max_image_tokens, 128)
415
+ for chunk_start, chunk_end, chunk_n, chunk_m in chunk_plan:
416
+ chunk_size = chunk_end - chunk_start
417
+ kv_ws, idx_ws, lens_ws = workspace_manager.get_simultaneous(
418
+ ((chunk_size, chunk_m, q.shape[-1]), torch.bfloat16),
419
+ ((self.max_num_batched_tokens, combined_topk), torch.int32),
420
+ ((self.max_num_batched_tokens,), torch.int32),
421
+ )
422
+ if not swa_only:
423
+ assert flashmla_metadata is not None
424
+ dequantize_and_gather_k_cache(
425
+ kv_ws[:chunk_size],
426
+ self._compressed_kv_cache(),
427
+ seq_lens=seq_lens[chunk_start:chunk_end] // self.compress_ratio,
428
+ gather_lens=None,
429
+ block_table=flashmla_metadata.block_table[num_decodes:][
430
+ chunk_start:chunk_end
431
+ ],
432
+ block_size=flashmla_metadata.block_size // self.compress_ratio,
433
+ offset=0,
434
+ )
435
+ dequantize_and_gather_k_cache(
436
+ kv_ws[:chunk_size],
437
+ self.swa_cache_layer.kv_cache,
438
+ seq_lens=seq_lens[chunk_start:chunk_end],
439
+ gather_lens=gather_lens[chunk_start:chunk_end],
440
+ block_table=swa_metadata.block_table[num_decodes:][
441
+ chunk_start:chunk_end
442
+ ],
443
+ block_size=swa_metadata.block_size,
444
+ offset=chunk_n,
445
+ )
446
+ first, last = num_decodes + chunk_start, num_decodes + chunk_end
447
+ qs = int(query_start_loc_cpu[first] - prefill_token_base)
448
+ qe = int(query_start_loc_cpu[last] - prefill_token_base)
449
+ combined_indices, combined_lens = combine_topk_swa_indices(
450
+ topk_indices[qs:qe],
451
+ query_start_loc[first : last + 1],
452
+ seq_lens[chunk_start:chunk_end],
453
+ gather_lens[chunk_start:chunk_end],
454
+ self.window_size,
455
+ self.compress_ratio,
456
+ top_k,
457
+ chunk_m,
458
+ chunk_n,
459
+ out=(idx_ws[: qe - qs], lens_ws[: qe - qs]),
460
+ left_visible=(
461
+ swa_metadata.prefill_left_visible[
462
+ num_decode_tokens + qs : num_decode_tokens + qe
463
+ ]
464
+ if swa_metadata.prefill_left_visible is not None
465
+ else None
466
+ ),
467
+ right_visible=(
468
+ swa_metadata.prefill_right_visible[
469
+ num_decode_tokens + qs : num_decode_tokens + qe
470
+ ]
471
+ if swa_metadata.prefill_right_visible is not None
472
+ else None
473
+ ),
474
+ max_image_tokens=self.max_image_tokens,
475
+ )
476
+ chunk_out = _token_slice(out, token_base + qs, token_base + qe)
477
+ # Mega attention over the gathered non-paged bf16 KV (RoPE already
478
+ # applied), in place into this chunk's token range. `indices`
479
+ # entries outside [0, s_kv) skip.
480
+ torch.ops._flashmla_C.fused_norm_rope_attn_rope_cast_fwd(
481
+ q[qs:qe],
482
+ kv_ws.view(-1, 1, q.shape[-1]),
483
+ combined_indices.unsqueeze(1),
484
+ self.scale,
485
+ _HEAD_DIM_V,
486
+ self.attn_sink,
487
+ combined_lens,
488
+ False, # enable_q_norm
489
+ 0.0, # rms_norm_eps
490
+ positions_int32[qs:qe],
491
+ *_ROPE_ARGS,
492
+ self.rotary_emb.cos_sin_cache,
493
+ self.n_wv_group,
494
+ *_SF_ARGS,
495
+ chunk_out.data,
496
+ chunk_out.scale,
497
+ )
vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py
+4
-2
@@ -201,7 +201,8 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
201
def get_padded_num_q_heads(cls, num_heads: int) -> int:
202
return _pad_to_supported_q_heads(num_heads)
203
204
- def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
204
+ def _o_proj(self, attn_out: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
205
+ o = attn_out[:, : self.n_local_heads, :]
206
return deep_gemm_fp8_o_proj(
207
o,
208
positions,
@@ -562,7 +563,8 @@ class DeepseekV4FlashInferSM120Attention(DeepseekV4Attention):
563
def get_padded_num_q_heads(cls, num_heads: int) -> int:
564
return _pad_to_supported_q_heads(num_heads)
565
565
- def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
566
+ def _o_proj(self, attn_out: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
567
+ o = attn_out[:, : self.n_local_heads, :]
568
return deep_gemm_fp8_o_proj(
569
o,
570
positions,
vllm/models/deepseek_v41/nvidia/flashmla.py
+2
-1
@@ -58,7 +58,8 @@ class DeepseekV4FlashMLAAttention(DeepseekV4Attention):
58
self._o_proj_block_size
59
)
60
61
- def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
61
+ def _o_proj(self, attn_out: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
62
+ o = attn_out[:, : self.n_local_heads, :]
63
return deep_gemm_fp8_o_proj(
64
o,
65
positions,
vllm/models/deepseek_v41/nvidia/model.py
+28
-2
@@ -64,6 +64,9 @@ from vllm.models.deepseek_v4.nvidia.model import (
64
make_deepseek_v4_expert_params_mapping,
65
)
66
from vllm.models.deepseek_v41.attention import DeepseekV4Attention
67
+from vllm.models.deepseek_v41.nvidia.flash_mla_mega_attn import (
68
+ DeepseekV4MegaAttnAttention,
69
+)
70
from vllm.models.deepseek_v41.nvidia.flashinfer_sparse import (
71
DeepseekV4FlashInferMLAAttention,
72
DeepseekV4FlashInferSM120Attention,
@@ -119,8 +122,9 @@ def _select_dsv4_attn_cls(vllm_config: VllmConfig) -> type[DeepseekV4Attention]:
122
123
The generic CUDA backend selector does not instantiate DSv4 layers directly,
124
so map generic sparse-MLA choices to the DSv4-specialized attention class.
122
- Without an explicit backend, SM12 defaults to FlashInfer while the other
123
- CUDA arches keep the FlashMLA path.
125
+ Without an explicit backend: SM12 takes FlashInfer, SM100 takes mega
126
+ attention where the topology allows it, and everything else keeps the
127
+ FlashMLA path.
128
"""
129
backend = vllm_config.attention_config.backend
130
device_capability = current_platform.get_device_capability()
@@ -140,6 +144,8 @@ def _select_dsv4_attn_cls(vllm_config: VllmConfig) -> type[DeepseekV4Attention]:
144
if device_capability is not None and device_capability.major == 12:
145
return DeepseekV4FlashInferSM120Attention
146
return DeepseekV4FlashInferMLAAttention
147
+ if backend is AttentionBackendEnum.FLASHMLA_MEGA_ATTN_DSV41:
148
+ return DeepseekV4MegaAttnAttention
149
if backend in (
150
AttentionBackendEnum.FLASHMLA_SPARSE,
151
AttentionBackendEnum.FLASHMLA_SPARSE_DSV4,
@@ -149,6 +155,14 @@ def _select_dsv4_attn_cls(vllm_config: VllmConfig) -> type[DeepseekV4Attention]:
155
156
if device_capability is not None and device_capability.major == 12:
157
return DeepseekV4FlashInferSM120Attention
158
+ # Mega attention is the SM100 default: it fuses Q RoPE, sparse attention,
159
+ # the output's inverse RoPE and its FP8 cast into one launch, and brings
160
+ # the 288 B NVFP4 compressed record -- the format the reference
161
+ # implementation itself stores. It declines topologies it cannot serve
162
+ # (non-SM100, TP that leaves fewer than WV_GROUP_SIZE heads per wo_a
163
+ # group, a build without the kernel), which then fall through to FlashMLA.
164
+ if DeepseekV4MegaAttnAttention.is_available_for(vllm_config):
165
+ return DeepseekV4MegaAttnAttention
166
return DeepseekV4FlashMLAAttention
167
168
@@ -936,6 +950,17 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
950
for layer in islice(self.layers, self.start_layer, self.end_layer):
951
layer.ffn.finalize_mega_moe_weights()
952
953
+ def finalize_mega_attn_weights(self) -> None:
954
+ """Permute wq_b / wo_a into FlashMLA's mega-attention layouts.
955
+
956
+ A no-op for every other attention layer, and idempotent, so a second
957
+ post-load pass cannot permute twice.
958
+ """
959
+ for layer in islice(self.layers, self.start_layer, self.end_layer):
960
+ finalize = getattr(layer.attn, "finalize_loaded_weights", None)
961
+ if finalize is not None:
962
+ finalize()
963
+
964
def finalize_mhc_broadcast_weights(self) -> None:
965
if not get_pp_group().is_first_rank or self.start_layer >= self.end_layer:
966
return
@@ -1190,6 +1215,7 @@ class DeepseekV41LLMForCausalLM(
1215
def process_weights_after_loading(self) -> None:
1216
self.model.finalize_mega_moe_weights()
1217
self.model.finalize_mhc_broadcast_weights()
1218
+ self.model.finalize_mega_attn_weights()
1219
1220
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
1221
return self.model.get_expert_mapping()
vllm/models/deepseek_v41/sparse_mla.py
+51
@@ -229,3 +229,54 @@ class DeepseekV4FlashMLABackend(DeepseekV4SparseMLABackend):
229
@staticmethod
230
def get_builder_cls() -> type[DeepseekV4FlashMLAMetadataBuilder]:
231
return DeepseekV4FlashMLAMetadataBuilder
232
+
233
+
234
+class FlashMLAMegaAttnBackend(DeepseekV4FlashMLABackend):
235
+ """FlashMLA's mega-attention kernel: Q RoPE + sparse attention + inverse
236
+ RoPE + FP8 cast of the output, in one launch.
237
+
238
+ Same metadata and KV-cache geometry as ``FLASHMLA_SPARSE_DSV41`` -- what
239
+ differs is the attention layer's interface contract (it takes an unnormed,
240
+ unroped Q and returns an already-inverse-RoPE'd, quantized output) and the
241
+ extra ``nvfp4_ds_mla`` compressed-cache record only this kernel can read.
242
+ SM100 only; the kernel has no SM90 instantiation.
243
+ """
244
+
245
+ supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
246
+ "auto",
247
+ "fp8_ds_mla",
248
+ "fp8", # alias for fp8_ds_mla
249
+ "nvfp4_ds_mla", # V4.1 fp8 SWA cache + NVFP4 compressed cache
250
+ ]
251
+
252
+ @staticmethod
253
+ def get_name() -> str:
254
+ return "FLASHMLA_MEGA_ATTN_DSV41"
255
+
256
+ @staticmethod
257
+ def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
258
+ return [128]
259
+
260
+ @classmethod
261
+ def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
262
+ return capability.major == 10
263
+
264
+ @classmethod
265
+ def supports_combination(
266
+ cls,
267
+ head_size: int,
268
+ dtype: torch.dtype,
269
+ kv_cache_dtype: CacheDType | None,
270
+ block_size: int | None,
271
+ use_mla: bool,
272
+ has_sink: bool,
273
+ use_sparse: bool,
274
+ use_mm_prefix: bool,
275
+ device_capability: DeviceCapability,
276
+ ) -> str | None:
277
+ # Imported here: the layer module imports this backend class.
278
+ from vllm.models.deepseek_v41.nvidia.flash_mla_mega_attn import (
279
+ is_flashmla_mega_attn_supported,
280
+ )
281
+
282
+ return is_flashmla_mega_attn_supported()[1]
vllm/v1/attention/backends/mla/sparse_swa.py
+4
-1
@@ -112,7 +112,10 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase):
112
# fp8_ds_mla's UE8M0 paged layout rounds its page up to the decode
113
# kernel's TMA stride; contiguous bf16/fp8 cache uses the natural
114
# element-size page.
115
- uses_fp8_ds_mla_layout = self.cache_config.cache_dtype == "fp8_ds_mla"
115
+ uses_fp8_ds_mla_layout = self.cache_config.cache_dtype in (
116
+ "fp8_ds_mla",
117
+ "nvfp4_ds_mla",
118
+ )
119
return SlidingWindowMLASpec(
120
block_size=self.block_size,
121
num_kv_heads=1,
vllm/v1/attention/backends/registry.py
+3
@@ -114,6 +114,9 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
114
"vllm.models.deepseek_v41.nvidia.flashinfer_sparse."
115
"DeepseekV4FlashInferMLASparseBackend"
116
)
117
+ FLASHMLA_MEGA_ATTN_DSV41 = (
118
+ "vllm.models.deepseek_v41.sparse_mla.FlashMLAMegaAttnBackend"
119
+ )
120
B12X = "vllm.v1.attention.backends.b12x.B12xPagedAttentionBackend"
121
FLASH_ATTN_MLA = "vllm.v1.attention.backends.mla.flashattn_mla.FlashAttnMLABackend"
122
FLASH_ATTN_MLA_SPARSE = (