[GLM5] Support FlashMLA FP8 KV cache (Hopper & Blackwell) (#47090)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Jun 30, 2026 at 18:13 UTC
f098ee70c730a20dbcec6ed0dbe6e4a041ac2848
3 files changed
+307
-44
tests/kernels/test_fused_deepseek_v32_norm_rope.py
+139
@@ -292,6 +292,76 @@ def test_fused_norm_rope_no_indexer(num_tokens: int):
292
assert (topk == 7).all(), "topk buffer should be untouched on shared layer"
293
294
295
+@pytest.mark.parametrize("num_tokens", [1, 4, 17, 512])
296
+def test_fused_norm_rope_ds_mla(num_tokens: int):
297
+ """fp8_ds_mla MLA cache layout (FlashMLA sparse, bf16-query path; SM90/SM100).
298
+
299
+ Per-token 656-byte entry: 512 fp8 NoPE (4 per-128 tiles, dynamic float32
300
+ scale) | 4 float32 scales | 64 bf16 (unquantized) RoPE.
301
+ """
302
+ torch.manual_seed(5)
303
+ dev = "cuda"
304
+ max_pos = 8192
305
+ pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) % max_pos
306
+
307
+ q_c = torch.randn(num_tokens, Q_LORA, device=dev, dtype=torch.bfloat16)
308
+ kv_c = torch.randn(num_tokens, KV_LORA, device=dev, dtype=torch.bfloat16)
309
+ k_pe = torch.randn(num_tokens, ROPE_DIM, device=dev, dtype=torch.bfloat16)
310
+ qw = torch.randn(Q_LORA, device=dev, dtype=torch.bfloat16)
311
+ kvw = torch.randn(KV_LORA, device=dev, dtype=torch.bfloat16)
312
+ mla_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
313
+
314
+ bs = max_pos
315
+ mla_cache = torch.zeros(1, bs, 656, device=dev, dtype=torch.uint8)
316
+ slot = torch.arange(num_tokens, device=dev, dtype=torch.int64)
317
+ topk = torch.full((num_tokens, 2048), 7, device=dev, dtype=torch.int32)
318
+
319
+ q_out = K.fused_norm_rope(
320
+ pos,
321
+ q_c,
322
+ qw,
323
+ EPS,
324
+ kv_c,
325
+ kvw,
326
+ EPS,
327
+ k_pe,
328
+ mla_cos_sin,
329
+ None,
330
+ None,
331
+ None,
332
+ EPS,
333
+ None,
334
+ topk,
335
+ slot_mapping=slot,
336
+ indexer_k_cache=None,
337
+ mla_kv_cache=mla_cache,
338
+ mla_kv_cache_dtype="fp8_ds_mla",
339
+ mla_k_scale=None,
340
+ has_indexer=False,
341
+ index_rope_interleave=False,
342
+ )
343
+
344
+ assert_bf16(q_out, rms_norm(q_c, qw), "q_c rmsnorm (ds_mla)")
345
+
346
+ kv_ref = rms_norm(kv_c, kvw) # [N, 512] fp32
347
+ kpe_ref = rope(k_pe.float(), pos, mla_cos_sin, interleave=True) # [N, 64]
348
+ tiles = kv_ref.view(num_tokens, 4, 128)
349
+ ref_scale = torch.clamp(tiles.abs().amax(dim=-1) / FP8_MAX, min=1.1754944e-38)
350
+ ref_nope = (tiles / ref_scale[..., None]).reshape(num_tokens, KV_LORA).to(FP8)
351
+
352
+ cache = mla_cache[0, :num_tokens] # [N, 656] uint8
353
+ nope = cache[:, :KV_LORA].view(FP8)
354
+ scales = cache.view(torch.float32)[:, KV_LORA // 4 : KV_LORA // 4 + 4]
355
+ rope_off = KV_LORA // 2 + 8
356
+ rope_vals = cache.view(torch.bfloat16)[:, rope_off : rope_off + ROPE_DIM]
357
+
358
+ torch.testing.assert_close(scales, ref_scale, rtol=1e-2, atol=1e-6)
359
+ assert_fp8(nope, ref_nope, "ds_mla NoPE fp8")
360
+ assert_bf16(rope_vals, kpe_ref, "ds_mla RoPE bf16")
361
+ # No indexer on this call: top-k buffer must be untouched.
362
+ assert (topk == 7).all(), "topk buffer should be untouched (no indexer)"
363
+
364
+
365
# ── fused_q ──────────────────────────────────────────────────────────────────
366
367
@@ -400,6 +470,75 @@ def test_fused_q_no_indexer(num_tokens: int):
470
assert_fp8(mqa[:, :, KV_LORA:], (qpe_ref / s).to(FP8), "mqa q_pe")
471
472
473
+@pytest.mark.parametrize("num_tokens", [1, 17, 512])
474
+@pytest.mark.parametrize("has_indexer", [True, False])
475
+def test_fused_q_bf16_query(num_tokens: int, has_indexer: bool):
476
+ """bf16-query path (FlashMLA sparse, SM90/SM100): only the RoPE'd q_pe is
477
+ produced (bf16, unquantized); ql_nope is consumed directly by the caller."""
478
+ torch.manual_seed(6)
479
+ dev = "cuda"
480
+ max_pos = 8192
481
+ pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) % max_pos
482
+
483
+ q_pe = torch.randn(
484
+ num_tokens, NUM_HEADS, ROPE_DIM, device=dev, dtype=torch.bfloat16
485
+ )
486
+ ql_nope = torch.randn(
487
+ num_tokens, NUM_HEADS, KV_LORA, device=dev, dtype=torch.bfloat16
488
+ )
489
+ q_scale = torch.tensor([0.37], device=dev, dtype=torch.float32)
490
+ q_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
491
+
492
+ index_q = index_w = idx_cos_sin = None
493
+ if has_indexer:
494
+ index_q = torch.randn(
495
+ num_tokens, INDEX_HEADS, INDEX_HEAD_DIM, device=dev, dtype=torch.bfloat16
496
+ )
497
+ index_w = torch.randn(num_tokens, INDEX_HEADS, device=dev, dtype=torch.float32)
498
+ idx_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
499
+
500
+ iq_fp8, iw_out, q_pe_out = K.fused_q(
501
+ pos,
502
+ q_pe,
503
+ q_cos_sin,
504
+ index_q,
505
+ idx_cos_sin,
506
+ ql_nope,
507
+ q_scale,
508
+ index_w,
509
+ INDEX_HEAD_DIM**-0.5,
510
+ INDEX_HEADS**-0.5,
511
+ has_indexer=has_indexer,
512
+ index_rope_interleave=False,
513
+ quantize_mqa=False,
514
+ )
515
+
516
+ # MQA query: only the RoPE'd q_pe, bf16, unquantized.
517
+ assert q_pe_out.dtype == torch.bfloat16
518
+ assert q_pe_out.shape == (num_tokens, NUM_HEADS, ROPE_DIM)
519
+ qpe_ref = rope(
520
+ q_pe.float(),
521
+ pos.unsqueeze(-1).expand(num_tokens, NUM_HEADS),
522
+ q_cos_sin,
523
+ interleave=True,
524
+ )
525
+ assert_bf16(q_pe_out, qpe_ref, "bf16 q_pe RoPE")
526
+
527
+ # Indexer-Q is unchanged on this path (still UE8M0 fp8 + folded weights).
528
+ if has_indexer:
529
+ assert index_q is not None
530
+ iq_ref = rope(
531
+ index_q.float(),
532
+ pos.unsqueeze(-1).expand(num_tokens, INDEX_HEADS),
533
+ idx_cos_sin,
534
+ interleave=False,
535
+ )
536
+ q_ref, scale_ref = ue8m0_quant(iq_ref)
537
+ assert_fp8(iq_fp8, q_ref, "indexer-Q fp8 (bf16-query path)")
538
+ iw_ref = index_w * scale_ref * (INDEX_HEAD_DIM**-0.5) * (INDEX_HEADS**-0.5)
539
+ torch.testing.assert_close(iw_out, iw_ref, rtol=1e-3, atol=1e-3)
540
+
541
+
542
# ── fused_eh_norm (MTP) ──────────────────────────────────────────────────────
543
544
vllm/models/deepseek_v32/nvidia/attention.py
+33
-13
@@ -269,20 +269,30 @@ class DeepseekV32Attention(MLAAttention):
269
# Runtime toggle for index_share_for_mtp_iteration: MTP draft step 0
270
# computes the top-k, steps 1+ set this True to reuse it.
271
self.skip_topk = False
272
- # Single fused fp8 path: Triton fused norm/rope/cache + fused-q write a
273
- # single fp8 MQA query and the contiguous [kv_c; k_pe] MLA cache layout.
274
- # This requires an fp8 KV cache and a sparse MLA backend that accepts a
275
- # quantized query (FlashInfer sparse on SM100).
276
- assert (
277
- is_quantized_kv_cache(self.kv_cache_dtype)
278
- and self.impl.supports_quant_query_input
279
- ), (
280
- "deepseek_v32 (nvidia) requires an fp8 KV cache served by the "
281
- "FlashInfer sparse MLA backend (which accepts a quantized query). "
282
- "Launch with --kv-cache-dtype fp8."
272
+ # Fused fp8 paths: Triton fused norm/rope/cache + fused-q. Two layouts,
273
+ # picked by the sparse MLA backend's query support:
274
+ # * supports_quant_query_input (FlashInfer sparse, SM100): per-tensor
275
+ # fp8 cache + a single packed fp8 MQA query.
276
+ # * not supported (FlashMLA sparse, SM90/SM100): fp8_ds_mla cache
277
+ # (per-128 block-scaled fp8 NoPE + unquantized bf16 RoPE) + a bf16
278
+ # (ql_nope, q_pe) query tuple. FA3 cannot mix a bf16 query with an
279
+ # fp8 KV cache, so FlashMLA (which dequantizes internally) is used.
280
+ # FlashMLA sparse runs on both Hopper and Blackwell, so this is the
281
+ # only DSA path on SM90 and an opt-in alternative on SM100.
282
+ assert is_quantized_kv_cache(self.kv_cache_dtype), (
283
+ "deepseek_v32 (nvidia) requires an fp8 KV cache served by a sparse "
284
+ "MLA backend. Launch with --kv-cache-dtype fp8 (FlashInfer sparse) "
285
+ "or --kv-cache-dtype fp8_ds_mla (FlashMLA sparse)."
286
)
287
+ self._fp8_query = self.impl.supports_quant_query_input
288
+ if not self._fp8_query:
289
+ assert self.kv_cache_dtype == "fp8_ds_mla", (
290
+ "deepseek_v32 (nvidia) on a bf16-query sparse MLA backend "
291
+ "(FlashMLA sparse) requires the fp8_ds_mla KV cache layout. "
292
+ "Launch with --kv-cache-dtype fp8_ds_mla."
293
+ )
294
# The paged KV cache is stored as uint8 and viewed as fp8 for the decode
285
- # (per-tensor fp8; never the fp8_ds_mla layout on this path).
295
+ # (per-tensor fp8). The fp8_ds_mla layout is consumed as raw bytes.
296
self._fp8_kv_needs_view = self.kv_cache_dtype != "fp8_ds_mla"
297
# GLM-5.2 uses interleaved indexer RoPE; DeepSeek-V3.2 uses NeoX.
298
self._index_rope_interleave = getattr(config, "indexer_rope_interleave", False)
@@ -465,6 +475,7 @@ class DeepseekV32Attention(MLAAttention):
475
indexer_n_head_scale,
476
has_indexer=has_indexer,
477
index_rope_interleave=self._index_rope_interleave,
478
+ quantize_mqa=self._fp8_query,
479
)
480
481
if self.indexer is not None:
@@ -496,8 +507,17 @@ class DeepseekV32Attention(MLAAttention):
507
kv_cache = self.kv_cache
508
if self._fp8_kv_needs_view:
509
kv_cache = kv_cache.view(torch.float8_e4m3fn)
510
+ if self._fp8_query:
511
+ # FlashInfer sparse: single packed fp8 query.
512
+ mqa_q_arg: torch.Tensor | tuple[torch.Tensor, torch.Tensor] = mqa_q[
513
+ :num_actual
514
+ ]
515
+ else:
516
+ # FlashMLA sparse: bf16 (ql_nope, q_pe) tuple. mqa_q is the RoPE'd
517
+ # q_pe; ql_nope is consumed directly.
518
+ mqa_q_arg = (ql_nope[:num_actual], mqa_q[:num_actual])
519
attn_out, _ = self.impl.forward_mqa( # type: ignore[attr-defined]
500
- mqa_q[:num_actual], kv_cache, attn_metadata, self
520
+ mqa_q_arg, kv_cache, attn_metadata, self
521
)
522
x = attn_out.view(
523
num_actual, self.num_local_heads, self.kv_lora_rank
vllm/models/deepseek_v32/nvidia/kernels.py
+135
-31
@@ -132,6 +132,14 @@ def _fused_norm_rope_kernel(
132
mla_cache_entry_stride,
133
MLA_CACHE_FP8: tl.constexpr,
134
mla_cache_scale_ptr,
135
+ # fp8_ds_mla cache views (block-scaled fp8 NoPE + unquantized bf16 RoPE).
136
+ # mla_cache_ptr is the fp8 (1-byte) view, so the block/entry strides above
137
+ # are byte offsets; these two share the same buffer as fp32 / bf16 views.
138
+ mla_cache_ds_scale_ptr,
139
+ mla_cache_ds_rope_ptr,
140
+ MLA_CACHE_DS_MLA: tl.constexpr,
141
+ MLA_NUM_TILES: tl.constexpr,
142
+ MLA_TILE_DIM: tl.constexpr,
143
# Top k indices
144
topk_indices_ptr,
145
topk_indices_stride,
@@ -209,6 +217,37 @@ def _fused_norm_rope_kernel(
217
mla_block_size = mla_cache_block_stride // mla_cache_entry_stride
218
mla_block_idx = slot_idx // mla_block_size
219
mla_block_off = slot_idx % mla_block_size
220
+
221
+ if MLA_CACHE_DS_MLA:
222
+ # fp8_ds_mla layout (DeepSeek-V3.2, KV_DIM == 512): per-128-element
223
+ # tile of the NoPE is dynamically quantized to fp8 with its own
224
+ # float32 scale; the RoPE tail is stored unquantized in bf16.
225
+ # bytes [0, KV_DIM) : KV_DIM fp8 NoPE values
226
+ # bytes [KV_DIM, KV_DIM + 16) : MLA_NUM_TILES float32 scales
227
+ # bytes [KV_DIM + 16, ...) : 2 * KPE_HALF_ROT_DIM bf16 RoPE
228
+ # mla_cache_block_stride / mla_cache_entry_stride are byte strides
229
+ # (mla_cache_ptr is the 1-byte fp8 view of the uint8 cache).
230
+ byte_base = (
231
+ mla_block_idx * mla_cache_block_stride
232
+ + mla_block_off * mla_cache_entry_stride
233
+ )
234
+ kv_2d = tl.reshape(kv_c, (MLA_NUM_TILES, MLA_TILE_DIM))
235
+ tile_amax = tl.max(tl.abs(kv_2d), axis=1, keep_dims=True)
236
+ # scale = amax / 448 (fp8 e4m3 max), matching the reference
237
+ # concat_and_cache_ds_mla kernel; floored to FLT_MIN.
238
+ tile_scale = tl.maximum(tile_amax * (1.0 / 448.0), 1.1754944e-38)
239
+ kv_c_fp8 = tl.reshape((kv_2d / tile_scale).to(tl.float8e4nv), (KV_DIM,))
240
+ tl.store(mla_cache_ptr + byte_base + kv_block, kv_c_fp8)
241
+ tile_off = tl.arange(0, MLA_NUM_TILES)
242
+ tl.store(
243
+ mla_cache_ds_scale_ptr + byte_base // 4 + KV_DIM // 4 + tile_off,
244
+ tl.reshape(tile_scale, (MLA_NUM_TILES,)),
245
+ )
246
+ rope_dst = mla_cache_ds_rope_ptr + byte_base // 2 + (KV_DIM // 2 + 8)
247
+ tl.store(rope_dst + dim_off * 2, r1.to(tl.bfloat16))
248
+ tl.store(rope_dst + dim_off * 2 + 1, r2.to(tl.bfloat16))
249
+ return
250
+
251
dst = (
252
mla_cache_ptr
253
+ mla_block_idx * mla_cache_block_stride
@@ -397,12 +436,29 @@ def fused_norm_rope(
436
)
437
438
# --- MLA KV cache setup ---
400
- mla_cache_fp8 = mla_kv_cache_dtype != "auto"
439
+ mla_cache_ds_mla = mla_kv_cache_dtype == "fp8_ds_mla"
440
+ mla_cache_fp8 = mla_kv_cache_dtype not in ("auto", "fp8_ds_mla")
441
+ mla_num_tiles = 1
442
+ mla_ds_scale_view = torch.empty(0, dtype=torch.float32, device=device)
443
+ mla_ds_rope_view = torch.empty(0, dtype=torch.bfloat16, device=device)
444
if mla_kv_cache is not None:
402
- mla_block_stride = mla_kv_cache.stride(0)
403
- mla_entry_stride = mla_kv_cache.stride(1)
404
- if mla_cache_fp8 and mla_kv_cache.dtype == torch.uint8:
405
- mla_kv_cache = mla_kv_cache.view(torch.float8_e4m3fn)
445
+ if mla_cache_ds_mla:
446
+ # 656-byte custom layout addressed in bytes; mla_cache_ptr is the
447
+ # 1-byte fp8 view, so block/entry strides are byte offsets and the
448
+ # fp32/bf16 views share the same buffer.
449
+ assert kv_dim == 512, "fp8_ds_mla requires kv_lora_rank == 512"
450
+ mla_num_tiles = kv_dim // 128
451
+ u8_cache = mla_kv_cache.view(torch.uint8)
452
+ mla_block_stride = u8_cache.stride(0)
453
+ mla_entry_stride = u8_cache.stride(1)
454
+ mla_ds_scale_view = u8_cache.view(torch.float32)
455
+ mla_ds_rope_view = u8_cache.view(torch.bfloat16)
456
+ mla_kv_cache = u8_cache.view(torch.float8_e4m3fn)
457
+ else:
458
+ mla_block_stride = mla_kv_cache.stride(0)
459
+ mla_entry_stride = mla_kv_cache.stride(1)
460
+ if mla_cache_fp8 and mla_kv_cache.dtype == torch.uint8:
461
+ mla_kv_cache = mla_kv_cache.view(torch.float8_e4m3fn)
462
if mla_k_scale is None:
463
mla_k_scale = torch.ones(1, dtype=torch.float32, device=device)
464
else:
@@ -461,6 +517,11 @@ def fused_norm_rope(
517
mla_entry_stride,
518
mla_cache_fp8,
519
mla_k_scale,
520
+ mla_ds_scale_view,
521
+ mla_ds_rope_view,
522
+ mla_cache_ds_mla,
523
+ mla_num_tiles,
524
+ kv_dim // mla_num_tiles if mla_cache_ds_mla else 1,
525
# Top k indices buffer
526
topk_indices_buffer,
527
topk_indices_buffer.stride(0),
@@ -506,6 +567,11 @@ def _fused_q_kernel(
567
q_scale_ptr,
568
QL_NOPE_DIM: tl.constexpr,
569
QL_NOPE_BLOCK: tl.constexpr,
570
+ # bf16 MQA query RoPE output (when QUANTIZE_MQA is False); the NoPE part is
571
+ # consumed directly from ql_nope, so only the RoPE'd q_pe is written here.
572
+ q_pe_out_ptr,
573
+ q_pe_out_stride0,
574
+ q_pe_out_stride1,
575
# Index weights
576
index_weights_ptr,
577
index_weights_stride,
@@ -515,13 +581,17 @@ def _fused_q_kernel(
581
index_weights_out_stride,
582
HAS_INDEXER: tl.constexpr,
583
INDEX_ROPE_INTERLEAVE: tl.constexpr,
584
+ QUANTIZE_MQA: tl.constexpr,
585
):
586
pid = tl.program_id(0)
587
tok_idx = tl.program_id(1)
588
head_idx = tl.program_id(2)
589
590
if pid == 2:
524
- # ql_nope quantize + pack into the front of mqa_q_fp8.
591
+ # ql_nope quantize + pack into the front of mqa_q_fp8. On the bf16
592
+ # query path ql_nope is consumed as-is (no pack), so skip entirely.
593
+ if not QUANTIZE_MQA:
594
+ return
595
if 2 * head_idx >= NUM_Q_HEADS:
596
return
597
@@ -581,23 +651,34 @@ def _fused_q_kernel(
651
).to(tl.float32)
652
r1 = x1 * cos - x2 * sin
653
r2 = x2 * cos + x1 * sin
584
- tl.store(
585
- mqa_q_fp8_ptr
586
- + tok_idx * mqa_q_fp8_stride0
587
- + q_head_idx * mqa_q_fp8_stride1
588
- + QL_NOPE_DIM
589
- + rot_off * 2,
590
- (r1 / scale).to(tl.float8e4nv),
591
- )
592
- tl.store(
593
- mqa_q_fp8_ptr
594
- + tok_idx * mqa_q_fp8_stride0
595
- + q_head_idx * mqa_q_fp8_stride1
596
- + QL_NOPE_DIM
597
- + rot_off * 2
598
- + 1,
599
- (r2 / scale).to(tl.float8e4nv),
600
- )
654
+ if QUANTIZE_MQA:
655
+ tl.store(
656
+ mqa_q_fp8_ptr
657
+ + tok_idx * mqa_q_fp8_stride0
658
+ + q_head_idx * mqa_q_fp8_stride1
659
+ + QL_NOPE_DIM
660
+ + rot_off * 2,
661
+ (r1 / scale).to(tl.float8e4nv),
662
+ )
663
+ tl.store(
664
+ mqa_q_fp8_ptr
665
+ + tok_idx * mqa_q_fp8_stride0
666
+ + q_head_idx * mqa_q_fp8_stride1
667
+ + QL_NOPE_DIM
668
+ + rot_off * 2
669
+ + 1,
670
+ (r2 / scale).to(tl.float8e4nv),
671
+ )
672
+ else:
673
+ # bf16 query: write the RoPE'd q_pe unquantized.
674
+ out_ty = q_pe_out_ptr.dtype.element_ty
675
+ q_pe_dst = (
676
+ q_pe_out_ptr
677
+ + tok_idx * q_pe_out_stride0
678
+ + q_head_idx * q_pe_out_stride1
679
+ )
680
+ tl.store(q_pe_dst + rot_off * 2, r1.to(out_ty))
681
+ tl.store(q_pe_dst + rot_off * 2 + 1, r2.to(out_ty))
682
return
683
elif pid == 1:
684
# Index Q RoPE + fp8 quant, all in registers. The roped bf16 index_q is
@@ -681,7 +762,16 @@ def fused_q(
762
index_weights_head_scale: float,
763
has_indexer: bool = True,
764
index_rope_interleave: bool = False,
765
+ quantize_mqa: bool = True,
766
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
767
+ """Fuse the MQA-query and indexer-query RoPE/quantization.
768
+
769
+ Returns ``(index_q_fp8, index_weights_out, mqa_q)``. When ``quantize_mqa``
770
+ is True (FlashInfer sparse, fp8 query) ``mqa_q`` is a single fp8 tensor
771
+ packing ``[ql_nope; q_pe]``. When False (FlashMLA sparse, bf16 query) it is
772
+ the RoPE'd ``q_pe`` in bf16; the caller pairs it with ``ql_nope`` as the
773
+ ``(ql_nope, q_pe)`` tuple the backend expects.
774
+ """
775
assert positions.ndim == 1
776
assert q_pe.ndim == 3
777
assert q_pe_cos_sin_cache.ndim == 2
@@ -705,13 +795,23 @@ def fused_q(
795
num_index_q_heads = index_q.shape[1]
796
index_q_head_dim = index_q.shape[2]
797
grid_heads = max(mqa_grid_heads, num_index_q_heads)
708
- mqa_q_fp8 = torch.empty(
709
- q_pe.shape[0],
710
- q_pe.shape[1],
711
- ql_nope.shape[2] + q_pe.shape[2],
712
- dtype=torch.float8_e4m3fn,
713
- device=q_pe.device,
714
- )
798
+ if quantize_mqa:
799
+ # fp8 path: pack [ql_nope; q_pe] into a single fp8 tensor.
800
+ mqa_q_fp8 = torch.empty(
801
+ q_pe.shape[0],
802
+ q_pe.shape[1],
803
+ ql_nope.shape[2] + q_pe.shape[2],
804
+ dtype=torch.float8_e4m3fn,
805
+ device=q_pe.device,
806
+ )
807
+ # Placeholder; pid 0 packs q_pe into mqa_q_fp8 instead.
808
+ q_pe_out = mqa_q_fp8
809
+ mqa_q = mqa_q_fp8
810
+ else:
811
+ # bf16 path: only the RoPE'd q_pe is produced; ql_nope used directly.
812
+ q_pe_out = torch.empty_like(q_pe)
813
+ mqa_q_fp8 = q_pe_out # unused placeholder for the fp8 pack pointer
814
+ mqa_q = q_pe_out
815
816
index_q_fp8 = torch.empty_like(index_q, dtype=torch.float8_e4m3fn)
817
index_weights_out = torch.empty_like(index_weights, dtype=torch.float32)
@@ -744,6 +844,9 @@ def fused_q(
844
q_scale,
845
ql_nope.shape[2],
846
triton.next_power_of_2(ql_nope.shape[2]),
847
+ q_pe_out,
848
+ q_pe_out.stride(0),
849
+ q_pe_out.stride(1),
850
index_weights,
851
index_weights.stride(0),
852
index_weights_softmax_scale,
@@ -752,12 +855,13 @@ def fused_q(
855
index_weights_out.stride(0),
856
HAS_INDEXER=has_indexer,
857
INDEX_ROPE_INTERLEAVE=index_rope_interleave,
858
+ QUANTIZE_MQA=quantize_mqa,
859
# num_warps=1 is optimal here: each program is a single 128-element
860
# rope+quant, so the kernel is program-count/occupancy bound, not
861
# per-program compute bound (swept 1/2/4/8 — 1 wins or ties everywhere).
862
num_warps=1,
863
)
760
- return index_q_fp8, index_weights_out, mqa_q_fp8
864
+ return index_q_fp8, index_weights_out, mqa_q
865
866
867
@triton.jit