@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 5ff3bbfb08

[Perf][DSv4] Fuse inverse RoPE + FP8 quant into FlashInfer sparse MLA (#58621)

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

Yongye Zhu committed Sep 24, 2026 at 23:26 UTC 5ff3bbfb089a22430ed7649dcd625c287928b51c
5 files changed +583 -104
tests/kernels/attention/test_flashmla_sparse.py
+275
@@ -511,6 +511,86 @@ def test_flashinfer_mixed_sparse_indices_separates_window_and_padded_width():
511 assert sparse_lens.cpu().tolist() == [padded_width]
512
513
514 +def test_flashinfer_mixed_sparse_indices_noncausal_rows_have_no_active_gaps():
515 + """A DSpark non-causal row keeps -1 out of the kernel's active ranges.
516 +
517 + Each draft token's causal window holds its first visible entries, the rest
518 + of the block starts at column `window`, and the compressed entries follow
519 + directly; `sparse_topk_lens` ends right after them.
520 + """
521 + from vllm.models.deepseek_v4.common.ops.cache_utils import (
522 + build_flashinfer_mixed_sparse_indices,
523 + )
524 +
525 + device = torch.device("cuda")
526 + window, width = 4, 8
527 + # One request: 1 context token, a 2-token draft block, slots 0..2 visible.
528 + decode_swa = torch.full((2, width), -1, dtype=torch.int32, device=device)
529 + decode_swa[:, :3] = torch.arange(3, dtype=torch.int32, device=device)
530 + compressed = torch.tensor(
531 + [[100, 101, -1, -1]] * 2, dtype=torch.int32, device=device
532 + )
533 + sparse_indices, sparse_lens = build_flashinfer_mixed_sparse_indices(
534 + decode_swa_indices=decode_swa,
535 + decode_compressed_indices=compressed,
536 + decode_compressed_topk_lens=torch.tensor(
537 + [2, 2], dtype=torch.int32, device=device
538 + ),
539 + prefill_topk_indices=torch.empty((0, 4), dtype=torch.int32, device=device),
540 + query_start_loc=torch.tensor([0, 2], dtype=torch.int32, device=device),
541 + seq_lens=torch.tensor([3], dtype=torch.int32, device=device),
542 + token_to_req_indices=torch.tensor([0, 0], dtype=torch.int32, device=device),
543 + swa_block_table=torch.tensor([[0]], dtype=torch.int32, device=device),
544 + swa_block_size=64,
545 + compressed_block_table=torch.tensor([[0]], dtype=torch.int32, device=device),
546 + compressed_block_size=64,
547 + window_size=window,
548 + compress_ratio=128,
549 + topk=0,
550 + )
551 +
552 + assert sparse_indices.cpu().tolist() == [
553 + # pos 1: window [0, 1]; slot 2 spills to column 4; compressed after it.
554 + [0, 1, -1, -1, 2, 100, 101, -1, -1, -1, -1, -1],
555 + # pos 2: window [0, 1, 2]; nothing spills; compressed at column 4.
556 + [0, 1, 2, -1, 100, 101, -1, -1, -1, -1, -1, -1],
557 + ]
558 + assert sparse_lens.cpu().tolist() == [7, 6]
559 +
560 +
561 +def test_flashinfer_mixed_sparse_indices_noncausal_pad_rows_keep_min_length():
562 + """A CUDA-graph pad token of the DSpark graph (a zero-length request) gets
563 + a causal pad row's shape instead of negative spill arithmetic."""
564 + from vllm.models.deepseek_v4.common.ops.cache_utils import (
565 + build_flashinfer_mixed_sparse_indices,
566 + )
567 +
568 + device = torch.device("cuda")
569 + window, width = 128, 192
570 + decode_swa = torch.full((3, width), -1, dtype=torch.int32, device=device)
571 + decode_swa[:2, :3] = torch.arange(3, dtype=torch.int32, device=device)
572 + sparse_indices, sparse_lens = build_flashinfer_mixed_sparse_indices(
573 + decode_swa_indices=decode_swa,
574 + decode_compressed_indices=None,
575 + decode_compressed_topk_lens=None,
576 + prefill_topk_indices=torch.empty((0, 0), dtype=torch.int32, device=device),
577 + # Request 1 is the padding: no query tokens, seq_len 0.
578 + query_start_loc=torch.tensor([0, 2, 2], dtype=torch.int32, device=device),
579 + seq_lens=torch.tensor([3, 0], dtype=torch.int32, device=device),
580 + token_to_req_indices=torch.tensor([0, 0, 1], dtype=torch.int32, device=device),
581 + swa_block_table=torch.zeros((2, 1), dtype=torch.int32, device=device),
582 + swa_block_size=64,
583 + compressed_block_table=None,
584 + compressed_block_size=64,
585 + window_size=window,
586 + compress_ratio=1,
587 + topk=0,
588 + )
589 +
590 + assert sparse_lens.cpu().tolist() == [129, 128, 128]
591 + assert (sparse_indices[2] == -1).all()
592 +
593 +
594 @pytest.mark.parametrize("model", ["deepseek_v4", "deepseek_v41"])
595 @pytest.mark.parametrize("page_padding", [0, 64])
596 def test_flashinfer_sparse_forward_reads_packed_kv_rows(model, page_padding):
@@ -565,3 +645,198 @@ def test_flashinfer_sparse_forward_reads_packed_kv_rows(model, page_padding):
645 attn._forward(query, None, cache, metadata, None, True, output)
646 expected = torch.tensor([2, 3, 3], dtype=output.dtype, device=device)
647 torch.testing.assert_close(output, expected[:, None, None].expand_as(output))
648 +
649 +
650 +def _make_rope_quant_attn(mod, num_rows: int, device: str):
651 + """A bare DSv4 FlashInfer layer set up for RopeQuant: 128 heads, fp8 KV."""
652 + from types import SimpleNamespace
653 +
654 + heads, groups, rank = 128, 16, 256
655 + angles = torch.arange(2048, device=device, dtype=torch.float32)[:, None] * 0.05
656 + angles = angles + torch.arange(32, device=device)[None] * 0.1
657 + attn = object.__new__(mod.DeepseekV4FlashInferMLAAttention)
658 + attn.compress_ratio = 1
659 + attn.window_size = 128
660 + attn.scale = 512**-0.5
661 + attn.kv_cache_torch_dtype = torch.float8_e4m3fn
662 + attn._flashinfer_fp8_bmm1_scale = attn.scale
663 + attn._flashinfer_fp8_bmm2_scale = 1.0
664 + attn.attn_sink = torch.linspace(-1, 1, heads, device=device)
665 + attn.padded_heads = heads
666 + attn.topk_indices_buffer = torch.empty(
667 + (num_rows, 0), dtype=torch.int32, device=device
668 + )
669 + attn.rotary_emb = SimpleNamespace(
670 + cos_sin_cache=torch.cat((angles.cos(), angles.sin()), -1).contiguous()
671 + )
672 + attn.n_local_heads, attn.n_local_groups, attn.o_lora_rank = heads, groups, rank
673 + attn.nope_head_dim, attn.rope_head_dim = 448, 64
674 + attn._einsum_recipe, attn._tma_aligned_scales = (1, 1, 128), True
675 + attn.wo_a = SimpleNamespace(
676 + weight=(torch.randn((groups, rank, 4096), device=device) * 0.05).to(
677 + torch.float8_e4m3fn
678 + ),
679 + weight_scale=torch.ones((groups, rank, 32), device=device),
680 + )
681 + attn.wo_b = lambda z: z
682 + return attn
683 +
684 +
685 +def _rope_quant_vs_unfused(attn, q, cache, metadata, positions, num_tokens):
686 + """Run `_forward` unfused and fused; return (bf16 attention, z_unfused, z_fused)."""
687 + from vllm.models.deepseek_v4.nvidia.ops.o_proj import rope_quant_attn_out
688 +
689 + heads, groups = attn.n_local_heads, attn.n_local_groups
690 + unfused = torch.empty(q.shape, dtype=torch.bfloat16, device=q.device)
691 + attn._forward(q, None, cache, metadata, None, True, unfused)
692 + fused = rope_quant_attn_out(q.shape[0], q.device)
693 + metadata.flashinfer_sparse_index_cache.clear()
694 + attn._forward(q, None, cache, metadata, None, True, fused)
695 + z_unfused = attn._o_proj(unfused[:num_tokens, :heads], positions)
696 + z_fused = attn._o_proj(fused, positions)[:num_tokens]
697 + return unfused[:num_tokens], z_unfused, z_fused.view(num_tokens, groups, -1)
698 +
699 +
700 +def _assert_projects_alike(z_unfused, z_fused):
701 + z_unfused = z_unfused.float().view_as(z_fused)
702 + rel = (z_fused.float() - z_unfused).norm() / z_unfused.norm()
703 + assert rel < 2e-2, rel
704 +
705 +
706 +def _skip_unless_rope_quant():
707 + from vllm.platforms import current_platform
708 +
709 + if not current_platform.is_device_capability_family(100):
710 + pytest.skip("Requires FlashInfer TRTLLM sparse MLA on SM100/SM103")
711 +
712 +
713 +def test_flashinfer_rope_quant_matches_unfused_o_proj():
714 + """RopeQuant's fused output must project like the bf16 + inv-RoPE/quant path.
715 +
716 + One mixed step (two decodes, a three-token prefill) plus CUDA-graph pad rows,
717 + so both calls run over the full buffer with un-rebased cum_seq_lens_q.
718 + """
719 + from vllm.models.deepseek_v4.nvidia import flashinfer_sparse as mod
720 + from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
721 +
722 + _skip_unless_rope_quant()
723 + torch.manual_seed(0)
724 + device = "cuda"
725 + block_size = 64
726 + q_lens, seq_lens = [1, 1, 3], [5, 9, 70]
727 + num_tokens, num_rows = sum(q_lens), sum(q_lens) + 3
728 + cache = torch.randn((3, block_size, 512), device=device).to(torch.float8_e4m3fn)
729 + attn = _make_rope_quant_attn(mod, num_rows, device)
730 +
731 + decode_indices = torch.full((2, 128), -1, dtype=torch.int32, device=device)
732 + for i, seq_len in enumerate(seq_lens[:2]):
733 + decode_indices[i, :seq_len] = torch.arange(
734 + i * block_size, i * block_size + seq_len, device=device
735 + )
736 + query_start = torch.tensor([0, 1, 2, 5], dtype=torch.int32)
737 + metadata = DeepseekSparseSWAMetadata(
738 + block_table=torch.tensor([[0], [1], [2]], dtype=torch.int32, device=device),
739 + slot_mapping=torch.zeros(num_tokens, dtype=torch.int64, device=device),
740 + block_size=block_size,
741 + seq_lens=torch.tensor(seq_lens, dtype=torch.int32, device=device),
742 + query_start_loc=query_start.to(device),
743 + query_start_loc_cpu=query_start,
744 + token_to_req_indices=torch.tensor(
745 + [0, 1, 2, 2, 2], dtype=torch.int32, device=device
746 + ),
747 + decode_swa_indices=decode_indices,
748 + decode_swa_lens=torch.tensor(seq_lens[:2], dtype=torch.int32, device=device),
749 + decode_swa_width=128,
750 + is_valid_token=torch.ones(num_tokens, dtype=torch.bool, device=device),
751 + num_decodes=2,
752 + num_prefills=1,
753 + num_decode_tokens=2,
754 + num_prefill_tokens=3,
755 + max_decode_query_len=1,
756 + )
757 + positions = torch.tensor([4, 8, 67, 68, 69], device=device)
758 + q = torch.randn((num_rows, 128, 512), device=device).to(torch.float8_e4m3fn)
759 + _, z_unfused, z_fused = _rope_quant_vs_unfused(
760 + attn, q, cache, metadata, positions, num_tokens
761 + )
762 + _assert_projects_alike(z_unfused, z_fused)
763 +
764 +
765 +@pytest.mark.parametrize("context_len", [0, 20, 127, 900])
766 +def test_flashinfer_dspark_noncausal_block_sees_future_tokens(context_len):
767 + """Every DSpark draft token attends to the whole block, fused or not.
768 +
769 + The kernel counts every entry of its active index ranges as a key, so the
770 + builder must place the block around each query's causal window without -1
771 + gaps, down to an empty context.
772 + """
773 + from vllm.models.deepseek_v4.nvidia import flashinfer_sparse as mod
774 + from vllm.v1.attention.backends.mla.compressor_utils import (
775 + get_dspark_swa_index_width,
776 + )
777 + from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
778 +
779 + _skip_unless_rope_quant()
780 + torch.manual_seed(0)
781 + device = "cuda"
782 + block_size, window, draft = 64, 128, 4
783 + width = get_dspark_swa_index_width(window, draft)
784 + contexts = [context_len, context_len + 37]
785 + num_reqs, num_tokens = len(contexts), len(contexts) * draft
786 + num_rows = num_tokens + 2
787 + pages = -(-(max(contexts) + draft) // block_size)
788 + cache = torch.randn((num_reqs * pages, block_size, 512), device=device)
789 + cache = cache.clamp_(-1, 1).to(torch.float8_e4m3fn)
790 + attn = _make_rope_quant_attn(mod, num_rows, device)
791 + attn.attn_sink = torch.full((128,), -float("inf"), device=device)
792 +
793 + # As the DSpark SWA builder lays them out: the trailing window of context
794 + # and the whole block, contiguous from column 0.
795 + decode_indices = torch.full((num_tokens, width), -1, dtype=torch.int32)
796 + visible, positions = [], []
797 + for r, ctx in enumerate(contexts):
798 + slots = torch.arange(max(ctx - window, 0), ctx + draft) + r * pages * block_size
799 + for k in range(draft):
800 + decode_indices[r * draft + k, : slots.numel()] = slots.to(torch.int32)
801 + visible.append(slots)
802 + positions.append(ctx + k)
803 + query_start = torch.arange(0, num_tokens + 1, draft, dtype=torch.int32)
804 + seq_lens = torch.tensor([c + draft for c in contexts], dtype=torch.int32)
805 + metadata = DeepseekSparseSWAMetadata(
806 + block_table=torch.arange(
807 + num_reqs * pages, dtype=torch.int32, device=device
808 + ).view(num_reqs, pages),
809 + slot_mapping=torch.zeros(num_tokens, dtype=torch.int64, device=device),
810 + block_size=block_size,
811 + seq_lens=seq_lens.to(device),
812 + query_start_loc=query_start.to(device),
813 + query_start_loc_cpu=query_start,
814 + token_to_req_indices=torch.arange(
815 + num_reqs, dtype=torch.int32, device=device
816 + ).repeat_interleave(draft),
817 + decode_swa_indices=decode_indices.to(device),
818 + decode_swa_lens=torch.tensor(
819 + [v.numel() for v in visible], dtype=torch.int32, device=device
820 + ),
821 + decode_swa_width=width,
822 + is_valid_token=torch.ones(num_tokens, dtype=torch.bool, device=device),
823 + num_decodes=num_reqs,
824 + num_prefills=0,
825 + num_decode_tokens=num_tokens,
826 + num_prefill_tokens=0,
827 + max_decode_query_len=draft,
828 + )
829 + q = torch.randn((num_rows, 128, 512), device=device).clamp_(-1, 1)
830 + q = q.to(torch.float8_e4m3fn)
831 + attention, z_unfused, z_fused = _rope_quant_vs_unfused(
832 + attn, q, cache, metadata, torch.tensor(positions, device=device), num_tokens
833 + )
834 +
835 + keys_all = cache.flatten(0, 1).float()
836 + for token, slots in enumerate(visible):
837 + keys = keys_all[slots.to(device)]
838 + weights = torch.softmax(q[token].float() @ keys.T * attn.scale, -1)
839 + torch.testing.assert_close(
840 + attention[token].float(), weights @ keys, atol=0.05, rtol=0.05
841 + )
842 + _assert_projects_alike(z_unfused, z_fused)
vllm/models/deepseek_v4/attention.py
+37 -20
@@ -43,6 +43,7 @@ from vllm.distributed import get_tensor_model_parallel_world_size
43 from vllm.forward_context import get_forward_context
44 from vllm.logger import init_logger
45 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
46 +from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation
47 from vllm.model_executor.layers.layernorm import RMSNorm
48 from vllm.model_executor.layers.quantization import QuantizationConfig
49 from vllm.model_executor.models.utils import extract_layer_index
@@ -171,14 +172,20 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
172 q: torch.Tensor,
173 kv: torch.Tensor,
174 positions: torch.Tensor,
174 - output: torch.Tensor,
175 + output: "torch.Tensor | QuantizedActivation",
176 ) -> None:
177 """Platform-specific sparse MLA forward; writes attention into ``output``."""
178 raise NotImplementedError
179
180 @abstractmethod
180 - def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
181 - """Inverse-RoPE + wo_a + wo_b output projection (platform-specific)."""
181 + def _o_proj(
182 + self, o: "torch.Tensor | QuantizedActivation", positions: torch.Tensor
183 + ) -> torch.Tensor:
184 + """Inverse-RoPE + wo_a + wo_b output projection (platform-specific).
185 +
186 + ``o`` is the live heads of the bf16 attention output, or the
187 + QuantizedActivation of a layer whose ``_alloc_attn_out`` returns one.
188 + """
189 raise NotImplementedError
190
191 def _uses_fp8_ds_mla_layout(self) -> bool:
@@ -467,14 +474,9 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
474 hidden_states: torch.Tensor,
475 llama_4_scaling: torch.Tensor | None = None,
476 ) -> torch.Tensor:
470 - # Pre-allocate attention output with FlashMLA-padded head count.
471 - # The op writes into `o_padded`; we slice to n_local_heads after.
472 - num_tokens = hidden_states.shape[0]
473 - o_padded = torch.empty(
474 - (num_tokens, self.padded_heads, self.head_dim),
475 - dtype=hidden_states.dtype,
476 - device=hidden_states.device,
477 - )
477 + # The eager attention region writes into a caller-owned buffer
478 + # (breakable_cudagraph needs in-place outputs).
479 + attn_out = self._alloc_attn_out(hidden_states.shape[0], hidden_states)
480
481 # Keep the attention input preparation in the captured graph. Only the
482 # sparse indexer and MLA attention run in the eager break below.
@@ -492,12 +494,27 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
494 indexer_kv_score,
495 indexer_weights,
496 positions,
495 - o_padded,
497 + attn_out,
498 + )
499 + if isinstance(attn_out, torch.Tensor):
500 + attn_out = attn_out[:, : self.n_local_heads, :]
501 + return self._o_proj(attn_out, positions)
502 +
503 + def _alloc_attn_out(
504 + self, num_tokens: int, hidden_states: torch.Tensor
505 + ) -> "torch.Tensor | QuantizedActivation":
506 + """The buffer ``forward_mqa`` fills.
507 +
508 + A bf16 ``[num_tokens, padded_heads, head_dim]`` buffer by default, whose
509 + padding heads are sliced off before ``_o_proj``. A layer whose kernel
510 + also does the inverse RoPE and the FP8 cast returns a
511 + QuantizedActivation instead and gets it back in ``_o_proj`` whole.
512 + """
513 + return torch.empty(
514 + (num_tokens, self.padded_heads, self.head_dim),
515 + dtype=hidden_states.dtype,
516 + device=hidden_states.device,
517 )
497 - o = o_padded[:, : self.n_local_heads, :]
498 -
499 - # Inverse-RoPE + wo_a + wo_b output projection (platform-specific).
500 - return self._o_proj(o, positions)
518
519 def _split_qkv_and_norm(
520 self, qr_kv: torch.Tensor
@@ -529,7 +546,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
546 indexer_kv_score: torch.Tensor,
547 indexer_weights: torch.Tensor,
548 positions: torch.Tensor,
532 - o_padded: torch.Tensor,
549 + o_padded: "torch.Tensor | QuantizedActivation",
550 ) -> None:
551 """Wide eager region: the whole of ``_prepare_and_attn`` runs eagerly.
552
@@ -558,7 +575,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
575 indexer_kv_score: torch.Tensor,
576 indexer_weights: torch.Tensor,
577 positions: torch.Tensor,
561 - o_padded: torch.Tensor,
578 + o_padded: "torch.Tensor | QuantizedActivation",
579 ) -> None:
580 """Attention input preparation followed by the sparse indexer and MLA.
581
@@ -723,7 +740,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
740 q: torch.Tensor,
741 kv: torch.Tensor,
742 positions: torch.Tensor,
726 - out: torch.Tensor,
743 + out: "torch.Tensor | QuantizedActivation",
744 ) -> None:
745 if self.indexer is not None and index_q is not None:
746 assert index_weights is not None
@@ -736,7 +753,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
753 )
754
755 # MLA attention writes into the pre-allocated `out` buffer
739 - # ([num_tokens, padded_heads, head_dim]).
756 + # (see _alloc_attn_out).
757 self.forward_mqa(q, kv, positions, out)
758
759 def _fused_qnorm_rope_kv_insert(
vllm/models/deepseek_v4/common/ops/cache_utils.py
+72 -8
@@ -1099,6 +1099,7 @@ def build_flashinfer_mixed_sparse_indices(
1099 prefill_left_visible: torch.Tensor | None = None,
1100 prefill_right_visible: torch.Tensor | None = None,
1101 max_image_tokens: int = 0,
1102 + num_rows: int | None = None,
1103 ) -> tuple[torch.Tensor, torch.Tensor]:
1104 """Build the FlashInfer DSV4 sparse-index matrix for decode-first batches.
1105
@@ -1109,6 +1110,9 @@ def build_flashinfer_mixed_sparse_indices(
1110 tokens derive their SWA window from the position and translate local
1111 compressed indices to global slots via the block tables.
1112
1113 + ``num_rows`` (>= ``num_tokens``) sizes both outputs for a kernel call that
1114 + spans a CUDA-graph-padded batch; the rows past ``num_tokens`` stay unset.
1115 +
1116 When ``prefill_left_visible``/``prefill_right_visible`` are given (vision
1117 variant), the SWA column region widens by ``max_image_tokens`` and prefill
1118 tokens inside an image span get a bidirectionally widened window; decode
@@ -1170,13 +1174,15 @@ def build_flashinfer_mixed_sparse_indices(
1174 # by ``sparse_topk_lens``, so padding never changes the attention result.
1175 padded_topk = max(topk, decode_compressed_topk)
1176 padded_topk = (padded_topk + 3) // 4 * 4
1177 + num_rows = num_tokens if num_rows is None else num_rows
1178 + assert num_rows >= num_tokens
1179 sparse_indices = torch.empty(
1174 - (num_tokens, swa_total_width + padded_topk),
1180 + (num_rows, swa_total_width + padded_topk),
1181 dtype=torch.int32,
1182 device=decode_swa_indices.device,
1183 )
1184 sparse_topk_lens = torch.empty(
1179 - num_tokens, dtype=torch.int32, device=decode_swa_indices.device
1185 + num_rows, dtype=torch.int32, device=decode_swa_indices.device
1186 )
1187 if num_tokens == 0:
1188 return sparse_indices, sparse_topk_lens
@@ -1193,9 +1199,10 @@ def build_flashinfer_mixed_sparse_indices(
1199 if compressed_block_span is None
1200 else compressed_block_span
1201 )
1202 + # The launch grid follows the output rows, so hand over only the real ones.
1203 _BUILD_FLASHINFER_MIXED_SPARSE_INDICES_KERNEL(
1197 - sparse_indices,
1198 - sparse_topk_lens,
1204 + sparse_indices[:num_tokens],
1205 + sparse_topk_lens[:num_tokens],
1206 decode_swa_indices,
1207 decode_compressed_indices,
1208 decode_compressed_topk_lens,
@@ -1322,12 +1329,52 @@ class BuildFlashinferMixedSparseIndicesKernel(
1329 SWA_TOTAL_WIDTH: tl.constexpr = SWA_INDEX_WIDTH + IMAGE_WIDTH
1330
1331 if token_idx < NUM_DECODE_TOKENS:
1332 + # A decode row wider than the window is DSpark's non-causal block:
1333 + # the block-anchored window of context followed by the whole block,
1334 + # contiguous from column 0. The kernel treats every entry of its
1335 + # active ranges as a key -- the first min(WINDOW_SIZE, pos + 1)
1336 + # columns and [WINDOW_SIZE, sparse_topk_lens) -- so -1 must not
1337 + # appear in either. Fill the causal window with the first visible
1338 + # entries and put the rest from column WINDOW_SIZE on, followed
1339 + # directly by the compressed entries.
1340 + NONCAUSAL: tl.constexpr = SWA_INDEX_WIDTH > WINDOW_SIZE
1341 + compressed_start = SWA_TOTAL_WIDTH
1342 + if NONCAUSAL:
1343 + req_idx = tl.load(token_to_req_indices_ptr + token_idx)
1344 + query_start = tl.load(query_start_loc_ptr + req_idx)
1345 + query_len = tl.load(query_start_loc_ptr + req_idx + 1) - query_start
1346 + context_len = tl.load(seq_lens_ptr + req_idx) - query_len
1347 + token_in_query = token_idx - query_start
1348 + window_len = tl.minimum(context_len + token_in_query + 1, WINDOW_SIZE)
1349 + spill_len = (
1350 + tl.minimum(context_len, WINDOW_SIZE) + query_len - window_len
1351 + )
1352 + # A CUDA-graph pad token sits outside its empty request, so
1353 + # the kernel never reads its row; keep it the shape of a
1354 + # causal pad row (all -1, WINDOW_SIZE long) rather than let
1355 + # the formulas above go negative.
1356 + is_block_token = (token_in_query >= 0) & (token_in_query < query_len)
1357 + window_len = tl.where(is_block_token, window_len, 0)
1358 + spill_len = tl.where(is_block_token, spill_len, 0)
1359 + compressed_start = WINDOW_SIZE + spill_len
1360 for i in range(0, SWA_TOTAL_WIDTH, WINDOW_BLOCK_SIZE):
1361 offset = i + tl.arange(0, WINDOW_BLOCK_SIZE)
1362 mask = offset < SWA_TOTAL_WIDTH
1363 + src = offset
1364 + src_mask = offset < SWA_INDEX_WIDTH
1365 + if NONCAUSAL:
1366 + # Columns from compressed_start on belong to the loops below.
1367 + mask = mask & (offset < compressed_start)
1368 + spill = offset >= WINDOW_SIZE
1369 + src = tl.where(spill, offset - WINDOW_SIZE + window_len, offset)
1370 + src_mask = tl.where(
1371 + spill,
1372 + offset < WINDOW_SIZE + spill_len,
1373 + offset < window_len,
1374 + )
1375 values = tl.load(
1329 - decode_swa_indices_ptr + token_idx * decode_swa_stride + offset,
1330 - mask=offset < SWA_INDEX_WIDTH,
1376 + decode_swa_indices_ptr + token_idx * decode_swa_stride + src,
1377 + mask=src_mask,
1378 other=-1,
1379 )
1380 values = _remap_flashinfer_index(values, swa_block_size, swa_block_span)
@@ -1372,12 +1419,27 @@ class BuildFlashinferMixedSparseIndicesKernel(
1419 tl.store(
1420 sparse_indices_ptr
1421 + token_idx * sparse_indices_stride
1375 - + SWA_TOTAL_WIDTH
1422 + + compressed_start
1423 + offset,
1424 values,
1425 mask=mask,
1426 )
1427
1428 + if NONCAUSAL:
1429 + # The row's remaining columns, left over by the shifted
1430 + # compressed block.
1431 + for i in range(0, SWA_TOTAL_WIDTH - WINDOW_SIZE, WINDOW_BLOCK_SIZE):
1432 + offset = i + tl.arange(0, WINDOW_BLOCK_SIZE)
1433 + tl.store(
1434 + sparse_indices_ptr
1435 + + token_idx * sparse_indices_stride
1436 + + compressed_start
1437 + + PADDED_TOP_K
1438 + + offset,
1439 + tl.full((WINDOW_BLOCK_SIZE,), -1, dtype=tl.int32),
1440 + mask=offset < SWA_TOTAL_WIDTH - compressed_start,
1441 + )
1442 +
1443 if DECODE_COMPRESSED_TOPK == 0:
1444 compressed_len = tl.zeros((), dtype=tl.int32)
1445 elif not DECODE_COMPRESSED_INDICES_ARE_LOCAL:
@@ -1388,7 +1450,9 @@ class BuildFlashinferMixedSparseIndicesKernel(
1450 else:
1451 compressed_len = tl.full((), DECODE_COMPRESSED_TOPK, dtype=tl.int32)
1452
1391 - tl.store(sparse_topk_lens_ptr + token_idx, SWA_TOTAL_WIDTH + compressed_len)
1453 + tl.store(
1454 + sparse_topk_lens_ptr + token_idx, compressed_start + compressed_len
1455 + )
1456 return
1457
1458 prefill_idx = token_idx - NUM_DECODE_TOKENS
vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py
+116 -60
@@ -9,6 +9,8 @@ import torch
9 from vllm.config import VllmConfig
10 from vllm.config.cache import CacheDType
11 from vllm.forward_context import get_forward_context
12 +from vllm.logger import init_logger
13 +from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation
14 from vllm.models.deepseek_v4.attention import DeepseekV4Attention
15 from vllm.models.deepseek_v4.common.ops import (
16 build_flashinfer_mixed_sparse_indices,
@@ -17,6 +19,8 @@ from vllm.models.deepseek_v4.common.ops import (
19 from vllm.models.deepseek_v4.nvidia.ops.o_proj import (
20 compute_fp8_einsum_recipe,
21 deep_gemm_fp8_o_proj,
22 + rope_quant_attn_out,
23 + rope_quant_unsupported_reason,
24 )
25 from vllm.models.deepseek_v4.sparse_mla import (
26 DeepseekV4FlashMLAMetadata,
@@ -37,6 +41,8 @@ from vllm.v1.attention.backends.mla.sparse_swa import (
41 if TYPE_CHECKING:
42 from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
43
44 +logger = init_logger(__name__)
45 +
46 _FLASHINFER_DSV4_WORKSPACE_BUFFER_SIZE = 128 * 1024 * 1024
47 _flashinfer_dsv4_workspace_by_device: dict[torch.device, torch.Tensor] = {}
48
@@ -203,7 +209,9 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
209 def get_padded_num_q_heads(cls, num_heads: int) -> int:
210 return _pad_to_supported_q_heads(num_heads)
211
206 - def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
212 + def _o_proj(
213 + self, o: torch.Tensor | QuantizedActivation, positions: torch.Tensor
214 + ) -> torch.Tensor:
215 return deep_gemm_fp8_o_proj(
216 o,
217 positions,
@@ -219,9 +227,30 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
227 tma_aligned_scales=self._tma_aligned_scales,
228 )
229
222 - def __init__(self, *args, **kwargs) -> None:
223 - super().__init__(*args, **kwargs)
230 + def _alloc_attn_out(
231 + self, num_tokens: int, hidden_states: torch.Tensor
232 + ) -> torch.Tensor | QuantizedActivation:
233 + if self._fuse_rope_quant:
234 + return rope_quant_attn_out(num_tokens, hidden_states.device)
235 + return super()._alloc_attn_out(num_tokens, hidden_states)
236 +
237 + def __init__(self, vllm_config: VllmConfig, *args, **kwargs) -> None:
238 + super().__init__(vllm_config, *args, **kwargs)
239 self._einsum_recipe, self._tma_aligned_scales = compute_fp8_einsum_recipe()
240 + # RopeQuant derives each query's RoPE position as seq_len - q_len + i,
241 + # which matches `positions` for DSpark's non-causal draft block too,
242 + # since every call passes the real per-request seq_lens. The one
243 + # exception is a draft block running past max_model_len, whose clamped
244 + # positions it does not reproduce; that only costs those drafts.
245 + reason = rope_quant_unsupported_reason(self)
246 + self._fuse_rope_quant = reason is None
247 + if reason is None:
248 + logger.info_once(
249 + "FLASHINFER_MLA_SPARSE_DSV4 fuses the inverse RoPE and FP8 "
250 + "quant of the attention output."
251 + )
252 + else:
253 + logger.debug_once("FLASHINFER_MLA_SPARSE_DSV4 RopeQuant off: %s", reason)
254 # Per-tensor FP8 scale buffers + precomputed scalar BMM scales. Only the
255 # per-tensor FP8 cache path consumes these; bf16 reads ``self.scale``.
256 if self.kv_cache_torch_dtype != torch.float8_e4m3fn:
@@ -253,32 +282,39 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
282 q: torch.Tensor,
283 kv: torch.Tensor,
284 positions: torch.Tensor,
256 - output: torch.Tensor,
285 + output: torch.Tensor | QuantizedActivation,
286 ) -> None:
258 - # The TRTLLM-gen kernel requires h_q in {64, 128}, so the output buffer
259 - # is allocated at the padded head count while q arrives at the local
260 - # head count; _forward pads q to match before the launcher.
261 - assert output.shape[0] == q.shape[0] and output.shape[-1] == q.shape[-1], (
262 - f"output buffer shape {output.shape} incompatible with q shape {q.shape}"
263 - )
264 - assert output.shape[1] >= q.shape[1], (
265 - f"output heads {output.shape[1]} must be >= q heads {q.shape[1]}"
266 - )
267 - # Per-tensor FP8 q produces a bf16 attention output.
268 - expected_output_dtype = (
269 - torch.bfloat16 if q.dtype == torch.float8_e4m3fn else q.dtype
270 - )
271 - assert output.dtype == expected_output_dtype, (
272 - f"output dtype {output.dtype} must match expected {expected_output_dtype} "
273 - f"for q dtype {q.dtype}"
274 - )
287 + if isinstance(output, QuantizedActivation):
288 + assert output.data.shape[0] == q.shape[0]
289 + else:
290 + # The TRTLLM-gen kernel requires h_q in {64, 128}, so the output
291 + # buffer is allocated at the padded head count while q arrives at
292 + # the local head count; _forward pads q to match before the launcher.
293 + assert output.shape[0] == q.shape[0] and output.shape[-1] == q.shape[-1], (
294 + f"output buffer shape {output.shape} incompatible with q {q.shape}"
295 + )
296 + assert output.shape[1] >= q.shape[1], (
297 + f"output heads {output.shape[1]} must be >= q heads {q.shape[1]}"
298 + )
299 + # Per-tensor FP8 q produces a bf16 attention output.
300 + expected_output_dtype = (
301 + torch.bfloat16 if q.dtype == torch.float8_e4m3fn else q.dtype
302 + )
303 + assert output.dtype == expected_output_dtype, (
304 + f"output dtype {output.dtype} must match expected "
305 + f"{expected_output_dtype} for q dtype {q.dtype}"
306 + )
307
308 forward_context = get_forward_context()
309 attn_metadata = forward_context.attn_metadata
310 if attn_metadata is None:
311 # Warmup dummy run: FlashInfer reads the cache directly and lazily
312 # allocates its workspace, so nothing to reserve here.
281 - output.zero_()
313 + if isinstance(output, QuantizedActivation):
314 + output.data.zero_()
315 + output.scale.zero_()
316 + else:
317 + output.zero_()
318 return
319
320 assert isinstance(attn_metadata, dict)
@@ -313,11 +349,12 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
349 swa_metadata: "DeepseekSparseSWAMetadata",
350 attn_metadata: DeepseekV4FlashMLAMetadata | None,
351 swa_only: bool,
352 + num_rows: int | None = None,
353 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
354 """Build the combined sparse-index tensors for the mixed batch.
355
356 Returns ``(compressed_kv_cache, seq_lens, sparse_indices,
320 - sparse_topk_lens)``.
357 + sparse_topk_lens)``, sized ``num_rows`` (default: the real tokens).
358 """
359 num_decodes = swa_metadata.num_decodes
360 num_prefills = swa_metadata.num_prefills
@@ -412,6 +449,8 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
449 if swa_only
450 else ("c128a" if self.compress_ratio == 128 else "c4a")
451 )
452 + if num_rows is not None:
453 + cache_key = f"{cache_key}:{num_rows}"
454 cached_sparse = swa_metadata.flashinfer_sparse_index_cache.get(cache_key, None)
455 if cached_sparse is None:
456 swa_block_span = _packed_block_span(swa_k_cache)
@@ -439,8 +478,9 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
478 prefill_right_visible=swa_metadata.prefill_right_visible,
479 # getattr for tests that bypass __init__ via object.__new__.
480 max_image_tokens=getattr(self, "max_image_tokens", 0),
481 + num_rows=num_rows,
482 )
443 - if cache_key != "c4a":
483 + if not cache_key.startswith("c4a"):
484 swa_metadata.flashinfer_sparse_index_cache[cache_key] = (
485 sparse_indices,
486 sparse_topk_lens,
@@ -457,7 +497,7 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
497 swa_metadata: "DeepseekSparseSWAMetadata",
498 attn_metadata: DeepseekV4FlashMLAMetadata | None,
499 swa_only: bool,
460 - output: torch.Tensor,
500 + output: torch.Tensor | QuantizedActivation,
501 ) -> None:
502 assert self.kv_cache_torch_dtype in (torch.bfloat16, torch.float8_e4m3fn)
503 num_decodes = swa_metadata.num_decodes
@@ -468,6 +508,11 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
508 num_tokens = num_decode_tokens + num_prefill_tokens
509 if num_tokens == 0:
510 return
511 + # RopeQuant's output is group-major with a group stride of exactly
512 + # sum_q tokens, so both calls span every row of the (CUDA-graph
513 + # padded) buffers and select their tokens through cum_seq_lens_q
514 + # alone; the kernel never touches rows outside it.
515 + rope_quant = isinstance(output, QuantizedActivation)
516
517 (
518 compressed_kv_cache,
@@ -480,12 +525,12 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
525 swa_metadata=swa_metadata,
526 attn_metadata=attn_metadata,
527 swa_only=swa_only,
528 + num_rows=q.shape[0] if rope_quant else None,
529 )
530
531 # CUDA graph execution can pad q/output past the scheduled token count;
532 # restrict to the real tokens (the launcher validates sparse indices).
487 - query = q[:num_tokens]
488 - output = output[:num_tokens]
533 + query = q if rope_quant else q[:num_tokens]
534 bmm1_scale: float | torch.Tensor = self.scale
535 bmm2_scale: float | torch.Tensor = 1.0
536 if self.kv_cache_torch_dtype == torch.float8_e4m3fn:
@@ -499,7 +544,9 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
544 # The TRTLLM-gen sparse-MLA kernel requires h_q in {64, 128}; zero-pad
545 # the query heads to the allocated output head count. Padded heads attend
546 # to the shared KV and are sliced off downstream (output is padded too).
502 - padded_heads = output.shape[1]
547 + padded_heads = (
548 + output.shape[1] if isinstance(output, torch.Tensor) else self.padded_heads
549 + )
550 if query.shape[1] < padded_heads:
551 padded_query = query.new_zeros(
552 (query.shape[0], padded_heads, query.shape[2])
@@ -512,51 +559,60 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
559 query_start_loc_cpu = swa_metadata.query_start_loc_cpu
560 assert query_start_loc is not None and query_start_loc_cpu is not None
561
515 - # Keep the TRTLLM-gen decode/prefill split: the launcher is tuned for
516 - # uniform-q batches, and this avoids flattening mixed batches into one call.
517 - if num_decode_tokens > 0:
518 - decode_cu = query_start_loc[: num_decodes + 1]
562 + def run(
563 + rows: slice, reqs: slice, cum_seq_lens_q: torch.Tensor, max_q_len: int
564 + ) -> None:
565 + if rope_quant:
566 + assert isinstance(output, QuantizedActivation)
567 + rows = slice(None)
568 + out_kwargs = dict(
569 + out=output.data,
570 + dsv4_output_scale=output.scale,
571 + dsv4_inv_rope_cos_sin_cache=self.rotary_emb.cos_sin_cache,
572 + )
573 + else:
574 + assert isinstance(output, torch.Tensor)
575 + out_kwargs = dict(out=output[rows])
576 flashinfer_trtllm_batch_decode_sparse_mla_dsv4(
520 - query=query[:num_decode_tokens],
577 + query=query[rows],
578 swa_kv_cache=swa_k_cache,
579 workspace_buffer=workspace,
523 - sparse_indices=sparse_indices[:num_decode_tokens],
580 + sparse_indices=sparse_indices[rows],
581 sparse_indices_are_storage_offsets=True,
582 compressed_kv_cache=compressed_kv_cache,
526 - sparse_topk_lens=sparse_topk_lens[:num_decode_tokens],
527 - seq_lens=seq_lens[:num_decodes],
528 - out=output[:num_decode_tokens],
583 + sparse_topk_lens=sparse_topk_lens[rows],
584 + seq_lens=seq_lens[reqs],
585 bmm1_scale=bmm1_scale,
586 bmm2_scale=bmm2_scale,
587 sinks=self.attn_sink,
532 - cum_seq_lens_q=decode_cu,
533 - max_q_len=swa_metadata.max_decode_query_len,
588 + cum_seq_lens_q=cum_seq_lens_q,
589 + max_q_len=max_q_len,
590 + **out_kwargs,
591 )
592
536 - if num_prefill_tokens > 0:
537 - # The prefill query view re-anchors at offset 0, so rebase the
538 - # cumulative query offsets to start at 0.
539 - prefill_cu = (
540 - query_start_loc[num_decodes : num_reqs + 1]
541 - - query_start_loc[num_decodes]
593 + # Keep the TRTLLM-gen decode/prefill split: the launcher is tuned for
594 + # uniform-q batches, and this avoids flattening mixed batches into one call.
595 + if num_decode_tokens > 0:
596 + run(
597 + slice(0, num_decode_tokens),
598 + slice(0, num_decodes),
599 + query_start_loc[: num_decodes + 1],
600 + swa_metadata.max_decode_query_len,
601 )
602 +
603 + if num_prefill_tokens > 0:
604 + prefill_cu = query_start_loc[num_decodes : num_reqs + 1]
605 + if not rope_quant:
606 + # The prefill query view re-anchors at offset 0, so rebase the
607 + # cumulative query offsets to start at 0.
608 + prefill_cu = prefill_cu - query_start_loc[num_decodes]
609 prefill_cu_cpu = query_start_loc_cpu[num_decodes : num_reqs + 1]
610 prefill_lens_cpu = prefill_cu_cpu[1:] - prefill_cu_cpu[:-1]
545 - flashinfer_trtllm_batch_decode_sparse_mla_dsv4(
546 - query=query[num_decode_tokens:num_tokens],
547 - swa_kv_cache=swa_k_cache,
548 - workspace_buffer=workspace,
549 - sparse_indices=sparse_indices[num_decode_tokens:num_tokens],
550 - sparse_indices_are_storage_offsets=True,
551 - compressed_kv_cache=compressed_kv_cache,
552 - sparse_topk_lens=sparse_topk_lens[num_decode_tokens:num_tokens],
553 - seq_lens=seq_lens[num_decodes:num_reqs],
554 - out=output[num_decode_tokens:num_tokens],
555 - bmm1_scale=bmm1_scale,
556 - bmm2_scale=bmm2_scale,
557 - sinks=self.attn_sink,
558 - cum_seq_lens_q=prefill_cu,
559 - max_q_len=int(prefill_lens_cpu.max().item()),
611 + run(
612 + slice(num_decode_tokens, num_tokens),
613 + slice(num_decodes, num_reqs),
614 + prefill_cu,
615 + int(prefill_lens_cpu.max().item()),
616 )
617
618
vllm/models/deepseek_v4/nvidia/ops/o_proj.py
+83 -16
@@ -5,11 +5,15 @@ from collections.abc import Callable
5 import torch
6 import torch.nn as nn
7
8 +from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation
9 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
10 + kFp8Dynamic128Sym,
11 +)
12 from vllm.models.deepseek_v4.common.ops.fused_inv_rope_fp8_quant import (
13 fused_inv_rope_fp8_quant,
14 )
15 from vllm.platforms import current_platform
12 -from vllm.utils.deep_gemm import fp8_einsum
16 +from vllm.utils.deep_gemm import fp8_einsum, get_tma_aligned_size
17
18
19 def compute_fp8_einsum_recipe(
@@ -29,7 +33,7 @@ def compute_fp8_einsum_recipe(
33
34
35 def deep_gemm_fp8_o_proj(
32 - o: torch.Tensor,
36 + o: torch.Tensor | QuantizedActivation,
37 positions: torch.Tensor,
38 cos_sin_cache: torch.Tensor,
39 wo_a: nn.Module,
@@ -49,23 +53,31 @@ def deep_gemm_fp8_o_proj(
53 layer selects the recipe at initialization. ``wo_b`` is any callable over
54 the flattened ``z``: the projection module itself, or a wrapper that also
55 reduce-scatters its output (DeepSeek-V4.1 GEMM-RS).
56 +
57 + A QuantizedActivation ``o`` was already rotated and cast by the attention
58 + kernel, one ``wo_a`` group per slot with the live groups first, so only
59 + the padding groups are sliced off.
60 """
61 use_fp8 = wo_a.weight.dtype == torch.float8_e4m3fn
54 - o_proj_input, o_scale = fused_inv_rope_fp8_quant(
55 - o,
56 - positions,
57 - cos_sin_cache,
58 - n_groups=n_groups,
59 - heads_per_group=heads_per_group,
60 - nope_dim=nope_dim,
61 - rope_dim=rope_dim,
62 - quant_group_size=einsum_recipe[2],
63 - tma_aligned_scales=tma_aligned_scales,
64 - quantize=use_fp8,
65 - )
62 + if isinstance(o, QuantizedActivation):
63 + assert use_fp8
64 + o_proj_input, o_scale = o.data[:, :n_groups], o.scale[:, :n_groups]
65 + else:
66 + o_proj_input, o_scale = fused_inv_rope_fp8_quant(
67 + o,
68 + positions,
69 + cos_sin_cache,
70 + n_groups=n_groups,
71 + heads_per_group=heads_per_group,
72 + nope_dim=nope_dim,
73 + rope_dim=rope_dim,
74 + quant_group_size=einsum_recipe[2],
75 + tma_aligned_scales=tma_aligned_scales,
76 + quantize=use_fp8,
77 + )
78 z = torch.empty(
67 - (o.shape[0], n_groups, o_lora_rank),
68 - device=o.device,
79 + (o_proj_input.shape[0], n_groups, o_lora_rank),
80 + device=o_proj_input.device,
81 dtype=torch.bfloat16,
82 )
83 if use_fp8:
@@ -89,3 +101,58 @@ def deep_gemm_fp8_o_proj(
101 out=z.transpose(0, 1),
102 )
103 return wo_b(z.flatten(1))
104 +
105 +
106 +# TRTLLM-gen's DSv4 RopeQuant epilogue (flashinfer-ai/flashinfer#4918) has one
107 +# fixed schedule: 128 query heads in groups of 8 512-wide heads, and FP8 with
108 +# one packed UE8M0 scale per 128 elements.
109 +_ROPE_QUANT_NUM_HEADS = 128
110 +_ROPE_QUANT_HEADS_PER_GROUP = 8
111 +_ROPE_QUANT_HEAD_DIM = 512
112 +_ROPE_QUANT_BLOCK = 128
113 +
114 +
115 +def rope_quant_unsupported_reason(layer: torch.nn.Module) -> str | None:
116 + """Why ``layer`` cannot take FlashInfer's fused inverse RoPE + FP8 output.
117 +
118 + DSv4 always has 8 heads per ``wo_a`` group at any TP size, and the rest of
119 + the kernel's fixed shape (head/RoPE dims, per-128 einsum scales, FP32 RoPE
120 + cache) is what the SM100 layer uses anyway. What varies is the local head
121 + count and the configs the layer also serves unquantized.
122 + """
123 + if layer.n_local_heads != _ROPE_QUANT_NUM_HEADS:
124 + # Padding fewer heads up to the kernel's fixed 128 costs more
125 + # attention than the fused epilogue saves.
126 + return f"RopeQuant needs {_ROPE_QUANT_NUM_HEADS} local query heads"
127 + if layer.kv_cache_torch_dtype != torch.float8_e4m3fn:
128 + return "RopeQuant needs the per-tensor FP8 KV cache"
129 + if layer.wo_a.weight.dtype != torch.float8_e4m3fn:
130 + return "RopeQuant needs an FP8 wo_a"
131 + return None
132 +
133 +
134 +def rope_quant_attn_out(num_tokens: int, device: torch.device) -> QuantizedActivation:
135 + """Allocate the output pair the RopeQuant kernel writes.
136 +
137 + Values are group-major, ``[num_tokens, 16, 8 * 512]`` over physical
138 + ``[group, token, K]``. Scales are MN-major packed UE8M0, one int32 per
139 + head over its four 128-wide blocks. Both are what ``fp8_einsum`` consumes.
140 + """
141 + n_groups = _ROPE_QUANT_NUM_HEADS // _ROPE_QUANT_HEADS_PER_GROUP
142 + group_width = _ROPE_QUANT_HEADS_PER_GROUP * _ROPE_QUANT_HEAD_DIM
143 + data = torch.empty(
144 + (n_groups, num_tokens, group_width), dtype=torch.float8_e4m3fn, device=device
145 + ).transpose(0, 1)
146 + aligned = get_tma_aligned_size(num_tokens, torch.int32.itemsize)
147 + scale = torch.empty(
148 + (n_groups, group_width // (_ROPE_QUANT_BLOCK * 4), aligned),
149 + dtype=torch.int32,
150 + device=device,
151 + ).permute(2, 0, 1)[:num_tokens]
152 + return QuantizedActivation(
153 + data=data,
154 + scale=scale,
155 + orig_dtype=torch.bfloat16,
156 + orig_shape=data.shape,
157 + quant_key=kFp8Dynamic128Sym,
158 + )