From 5ff3bbfb089a22430ed7649dcd625c287928b51c Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Thu, 24 Sep 2026 23:26:50 -0700 Subject: [PATCH] [Perf][DSv4] Fuse inverse RoPE + FP8 quant into FlashInfer sparse MLA (#58621) Signed-off-by: Yongye Zhu Co-authored-by: Claude Opus 5.5 --- .../kernels/attention/test_flashmla_sparse.py | 275 ++++++++++++++++++ vllm/models/deepseek_v4/attention.py | 55 ++-- .../deepseek_v4/common/ops/cache_utils.py | 80 ++++- .../deepseek_v4/nvidia/flashinfer_sparse.py | 186 +++++++----- vllm/models/deepseek_v4/nvidia/ops/o_proj.py | 99 ++++++- 5 files changed, 587 insertions(+), 108 deletions(-) diff --git a/tests/kernels/attention/test_flashmla_sparse.py b/tests/kernels/attention/test_flashmla_sparse.py index 7c1a455094..8fcf3f6a58 100644 --- a/tests/kernels/attention/test_flashmla_sparse.py +++ b/tests/kernels/attention/test_flashmla_sparse.py @@ -511,6 +511,86 @@ def test_flashinfer_mixed_sparse_indices_separates_window_and_padded_width(): assert sparse_lens.cpu().tolist() == [padded_width] +def test_flashinfer_mixed_sparse_indices_noncausal_rows_have_no_active_gaps(): + """A DSpark non-causal row keeps -1 out of the kernel's active ranges. + + Each draft token's causal window holds its first visible entries, the rest + of the block starts at column `window`, and the compressed entries follow + directly; `sparse_topk_lens` ends right after them. + """ + from vllm.models.deepseek_v4.common.ops.cache_utils import ( + build_flashinfer_mixed_sparse_indices, + ) + + device = torch.device("cuda") + window, width = 4, 8 + # One request: 1 context token, a 2-token draft block, slots 0..2 visible. + decode_swa = torch.full((2, width), -1, dtype=torch.int32, device=device) + decode_swa[:, :3] = torch.arange(3, dtype=torch.int32, device=device) + compressed = torch.tensor( + [[100, 101, -1, -1]] * 2, dtype=torch.int32, device=device + ) + sparse_indices, sparse_lens = build_flashinfer_mixed_sparse_indices( + decode_swa_indices=decode_swa, + decode_compressed_indices=compressed, + decode_compressed_topk_lens=torch.tensor( + [2, 2], dtype=torch.int32, device=device + ), + prefill_topk_indices=torch.empty((0, 4), dtype=torch.int32, device=device), + query_start_loc=torch.tensor([0, 2], dtype=torch.int32, device=device), + seq_lens=torch.tensor([3], dtype=torch.int32, device=device), + token_to_req_indices=torch.tensor([0, 0], dtype=torch.int32, device=device), + swa_block_table=torch.tensor([[0]], dtype=torch.int32, device=device), + swa_block_size=64, + compressed_block_table=torch.tensor([[0]], dtype=torch.int32, device=device), + compressed_block_size=64, + window_size=window, + compress_ratio=128, + topk=0, + ) + + assert sparse_indices.cpu().tolist() == [ + # pos 1: window [0, 1]; slot 2 spills to column 4; compressed after it. + [0, 1, -1, -1, 2, 100, 101, -1, -1, -1, -1, -1], + # pos 2: window [0, 1, 2]; nothing spills; compressed at column 4. + [0, 1, 2, -1, 100, 101, -1, -1, -1, -1, -1, -1], + ] + assert sparse_lens.cpu().tolist() == [7, 6] + + +def test_flashinfer_mixed_sparse_indices_noncausal_pad_rows_keep_min_length(): + """A CUDA-graph pad token of the DSpark graph (a zero-length request) gets + a causal pad row's shape instead of negative spill arithmetic.""" + from vllm.models.deepseek_v4.common.ops.cache_utils import ( + build_flashinfer_mixed_sparse_indices, + ) + + device = torch.device("cuda") + window, width = 128, 192 + decode_swa = torch.full((3, width), -1, dtype=torch.int32, device=device) + decode_swa[:2, :3] = torch.arange(3, dtype=torch.int32, device=device) + sparse_indices, sparse_lens = build_flashinfer_mixed_sparse_indices( + decode_swa_indices=decode_swa, + decode_compressed_indices=None, + decode_compressed_topk_lens=None, + prefill_topk_indices=torch.empty((0, 0), dtype=torch.int32, device=device), + # Request 1 is the padding: no query tokens, seq_len 0. + query_start_loc=torch.tensor([0, 2, 2], dtype=torch.int32, device=device), + seq_lens=torch.tensor([3, 0], dtype=torch.int32, device=device), + token_to_req_indices=torch.tensor([0, 0, 1], dtype=torch.int32, device=device), + swa_block_table=torch.zeros((2, 1), dtype=torch.int32, device=device), + swa_block_size=64, + compressed_block_table=None, + compressed_block_size=64, + window_size=window, + compress_ratio=1, + topk=0, + ) + + assert sparse_lens.cpu().tolist() == [129, 128, 128] + assert (sparse_indices[2] == -1).all() + + @pytest.mark.parametrize("model", ["deepseek_v4", "deepseek_v41"]) @pytest.mark.parametrize("page_padding", [0, 64]) 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): attn._forward(query, None, cache, metadata, None, True, output) expected = torch.tensor([2, 3, 3], dtype=output.dtype, device=device) torch.testing.assert_close(output, expected[:, None, None].expand_as(output)) + + +def _make_rope_quant_attn(mod, num_rows: int, device: str): + """A bare DSv4 FlashInfer layer set up for RopeQuant: 128 heads, fp8 KV.""" + from types import SimpleNamespace + + heads, groups, rank = 128, 16, 256 + angles = torch.arange(2048, device=device, dtype=torch.float32)[:, None] * 0.05 + angles = angles + torch.arange(32, device=device)[None] * 0.1 + attn = object.__new__(mod.DeepseekV4FlashInferMLAAttention) + attn.compress_ratio = 1 + attn.window_size = 128 + attn.scale = 512**-0.5 + attn.kv_cache_torch_dtype = torch.float8_e4m3fn + attn._flashinfer_fp8_bmm1_scale = attn.scale + attn._flashinfer_fp8_bmm2_scale = 1.0 + attn.attn_sink = torch.linspace(-1, 1, heads, device=device) + attn.padded_heads = heads + attn.topk_indices_buffer = torch.empty( + (num_rows, 0), dtype=torch.int32, device=device + ) + attn.rotary_emb = SimpleNamespace( + cos_sin_cache=torch.cat((angles.cos(), angles.sin()), -1).contiguous() + ) + attn.n_local_heads, attn.n_local_groups, attn.o_lora_rank = heads, groups, rank + attn.nope_head_dim, attn.rope_head_dim = 448, 64 + attn._einsum_recipe, attn._tma_aligned_scales = (1, 1, 128), True + attn.wo_a = SimpleNamespace( + weight=(torch.randn((groups, rank, 4096), device=device) * 0.05).to( + torch.float8_e4m3fn + ), + weight_scale=torch.ones((groups, rank, 32), device=device), + ) + attn.wo_b = lambda z: z + return attn + + +def _rope_quant_vs_unfused(attn, q, cache, metadata, positions, num_tokens): + """Run `_forward` unfused and fused; return (bf16 attention, z_unfused, z_fused).""" + from vllm.models.deepseek_v4.nvidia.ops.o_proj import rope_quant_attn_out + + heads, groups = attn.n_local_heads, attn.n_local_groups + unfused = torch.empty(q.shape, dtype=torch.bfloat16, device=q.device) + attn._forward(q, None, cache, metadata, None, True, unfused) + fused = rope_quant_attn_out(q.shape[0], q.device) + metadata.flashinfer_sparse_index_cache.clear() + attn._forward(q, None, cache, metadata, None, True, fused) + z_unfused = attn._o_proj(unfused[:num_tokens, :heads], positions) + z_fused = attn._o_proj(fused, positions)[:num_tokens] + return unfused[:num_tokens], z_unfused, z_fused.view(num_tokens, groups, -1) + + +def _assert_projects_alike(z_unfused, z_fused): + z_unfused = z_unfused.float().view_as(z_fused) + rel = (z_fused.float() - z_unfused).norm() / z_unfused.norm() + assert rel < 2e-2, rel + + +def _skip_unless_rope_quant(): + from vllm.platforms import current_platform + + if not current_platform.is_device_capability_family(100): + pytest.skip("Requires FlashInfer TRTLLM sparse MLA on SM100/SM103") + + +def test_flashinfer_rope_quant_matches_unfused_o_proj(): + """RopeQuant's fused output must project like the bf16 + inv-RoPE/quant path. + + One mixed step (two decodes, a three-token prefill) plus CUDA-graph pad rows, + so both calls run over the full buffer with un-rebased cum_seq_lens_q. + """ + from vllm.models.deepseek_v4.nvidia import flashinfer_sparse as mod + from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata + + _skip_unless_rope_quant() + torch.manual_seed(0) + device = "cuda" + block_size = 64 + q_lens, seq_lens = [1, 1, 3], [5, 9, 70] + num_tokens, num_rows = sum(q_lens), sum(q_lens) + 3 + cache = torch.randn((3, block_size, 512), device=device).to(torch.float8_e4m3fn) + attn = _make_rope_quant_attn(mod, num_rows, device) + + decode_indices = torch.full((2, 128), -1, dtype=torch.int32, device=device) + for i, seq_len in enumerate(seq_lens[:2]): + decode_indices[i, :seq_len] = torch.arange( + i * block_size, i * block_size + seq_len, device=device + ) + query_start = torch.tensor([0, 1, 2, 5], dtype=torch.int32) + metadata = DeepseekSparseSWAMetadata( + block_table=torch.tensor([[0], [1], [2]], dtype=torch.int32, device=device), + slot_mapping=torch.zeros(num_tokens, dtype=torch.int64, device=device), + block_size=block_size, + seq_lens=torch.tensor(seq_lens, dtype=torch.int32, device=device), + query_start_loc=query_start.to(device), + query_start_loc_cpu=query_start, + token_to_req_indices=torch.tensor( + [0, 1, 2, 2, 2], dtype=torch.int32, device=device + ), + decode_swa_indices=decode_indices, + decode_swa_lens=torch.tensor(seq_lens[:2], dtype=torch.int32, device=device), + decode_swa_width=128, + is_valid_token=torch.ones(num_tokens, dtype=torch.bool, device=device), + num_decodes=2, + num_prefills=1, + num_decode_tokens=2, + num_prefill_tokens=3, + max_decode_query_len=1, + ) + positions = torch.tensor([4, 8, 67, 68, 69], device=device) + q = torch.randn((num_rows, 128, 512), device=device).to(torch.float8_e4m3fn) + _, z_unfused, z_fused = _rope_quant_vs_unfused( + attn, q, cache, metadata, positions, num_tokens + ) + _assert_projects_alike(z_unfused, z_fused) + + +@pytest.mark.parametrize("context_len", [0, 20, 127, 900]) +def test_flashinfer_dspark_noncausal_block_sees_future_tokens(context_len): + """Every DSpark draft token attends to the whole block, fused or not. + + The kernel counts every entry of its active index ranges as a key, so the + builder must place the block around each query's causal window without -1 + gaps, down to an empty context. + """ + from vllm.models.deepseek_v4.nvidia import flashinfer_sparse as mod + from vllm.v1.attention.backends.mla.compressor_utils import ( + get_dspark_swa_index_width, + ) + from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata + + _skip_unless_rope_quant() + torch.manual_seed(0) + device = "cuda" + block_size, window, draft = 64, 128, 4 + width = get_dspark_swa_index_width(window, draft) + contexts = [context_len, context_len + 37] + num_reqs, num_tokens = len(contexts), len(contexts) * draft + num_rows = num_tokens + 2 + pages = -(-(max(contexts) + draft) // block_size) + cache = torch.randn((num_reqs * pages, block_size, 512), device=device) + cache = cache.clamp_(-1, 1).to(torch.float8_e4m3fn) + attn = _make_rope_quant_attn(mod, num_rows, device) + attn.attn_sink = torch.full((128,), -float("inf"), device=device) + + # As the DSpark SWA builder lays them out: the trailing window of context + # and the whole block, contiguous from column 0. + decode_indices = torch.full((num_tokens, width), -1, dtype=torch.int32) + visible, positions = [], [] + for r, ctx in enumerate(contexts): + slots = torch.arange(max(ctx - window, 0), ctx + draft) + r * pages * block_size + for k in range(draft): + decode_indices[r * draft + k, : slots.numel()] = slots.to(torch.int32) + visible.append(slots) + positions.append(ctx + k) + query_start = torch.arange(0, num_tokens + 1, draft, dtype=torch.int32) + seq_lens = torch.tensor([c + draft for c in contexts], dtype=torch.int32) + metadata = DeepseekSparseSWAMetadata( + block_table=torch.arange( + num_reqs * pages, dtype=torch.int32, device=device + ).view(num_reqs, pages), + slot_mapping=torch.zeros(num_tokens, dtype=torch.int64, device=device), + block_size=block_size, + seq_lens=seq_lens.to(device), + query_start_loc=query_start.to(device), + query_start_loc_cpu=query_start, + token_to_req_indices=torch.arange( + num_reqs, dtype=torch.int32, device=device + ).repeat_interleave(draft), + decode_swa_indices=decode_indices.to(device), + decode_swa_lens=torch.tensor( + [v.numel() for v in visible], dtype=torch.int32, device=device + ), + decode_swa_width=width, + is_valid_token=torch.ones(num_tokens, dtype=torch.bool, device=device), + num_decodes=num_reqs, + num_prefills=0, + num_decode_tokens=num_tokens, + num_prefill_tokens=0, + max_decode_query_len=draft, + ) + q = torch.randn((num_rows, 128, 512), device=device).clamp_(-1, 1) + q = q.to(torch.float8_e4m3fn) + attention, z_unfused, z_fused = _rope_quant_vs_unfused( + attn, q, cache, metadata, torch.tensor(positions, device=device), num_tokens + ) + + keys_all = cache.flatten(0, 1).float() + for token, slots in enumerate(visible): + keys = keys_all[slots.to(device)] + weights = torch.softmax(q[token].float() @ keys.T * attn.scale, -1) + torch.testing.assert_close( + attention[token].float(), weights @ keys, atol=0.05, rtol=0.05 + ) + _assert_projects_alike(z_unfused, z_fused) diff --git a/vllm/models/deepseek_v4/attention.py b/vllm/models/deepseek_v4/attention.py index 8a151dba1d..2315bedfeb 100644 --- a/vllm/models/deepseek_v4/attention.py +++ b/vllm/models/deepseek_v4/attention.py @@ -43,6 +43,7 @@ from vllm.distributed import get_tensor_model_parallel_world_size from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase +from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.models.utils import extract_layer_index @@ -171,14 +172,20 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): q: torch.Tensor, kv: torch.Tensor, positions: torch.Tensor, - output: torch.Tensor, + output: "torch.Tensor | QuantizedActivation", ) -> None: """Platform-specific sparse MLA forward; writes attention into ``output``.""" raise NotImplementedError @abstractmethod - def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: - """Inverse-RoPE + wo_a + wo_b output projection (platform-specific).""" + def _o_proj( + self, o: "torch.Tensor | QuantizedActivation", positions: torch.Tensor + ) -> torch.Tensor: + """Inverse-RoPE + wo_a + wo_b output projection (platform-specific). + + ``o`` is the live heads of the bf16 attention output, or the + QuantizedActivation of a layer whose ``_alloc_attn_out`` returns one. + """ raise NotImplementedError def _uses_fp8_ds_mla_layout(self) -> bool: @@ -467,14 +474,9 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): hidden_states: torch.Tensor, llama_4_scaling: torch.Tensor | None = None, ) -> torch.Tensor: - # Pre-allocate attention output with FlashMLA-padded head count. - # The op writes into `o_padded`; we slice to n_local_heads after. - num_tokens = hidden_states.shape[0] - o_padded = torch.empty( - (num_tokens, self.padded_heads, self.head_dim), - dtype=hidden_states.dtype, - device=hidden_states.device, - ) + # The eager attention region writes into a caller-owned buffer + # (breakable_cudagraph needs in-place outputs). + attn_out = self._alloc_attn_out(hidden_states.shape[0], hidden_states) # Keep the attention input preparation in the captured graph. Only the # sparse indexer and MLA attention run in the eager break below. @@ -492,12 +494,27 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): indexer_kv_score, indexer_weights, positions, - o_padded, + attn_out, ) - o = o_padded[:, : self.n_local_heads, :] + if isinstance(attn_out, torch.Tensor): + attn_out = attn_out[:, : self.n_local_heads, :] + return self._o_proj(attn_out, positions) - # Inverse-RoPE + wo_a + wo_b output projection (platform-specific). - return self._o_proj(o, positions) + def _alloc_attn_out( + self, num_tokens: int, hidden_states: torch.Tensor + ) -> "torch.Tensor | QuantizedActivation": + """The buffer ``forward_mqa`` fills. + + A bf16 ``[num_tokens, padded_heads, head_dim]`` buffer by default, whose + padding heads are sliced off before ``_o_proj``. A layer whose kernel + also does the inverse RoPE and the FP8 cast returns a + QuantizedActivation instead and gets it back in ``_o_proj`` whole. + """ + return torch.empty( + (num_tokens, self.padded_heads, self.head_dim), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) def _split_qkv_and_norm( self, qr_kv: torch.Tensor @@ -529,7 +546,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): indexer_kv_score: torch.Tensor, indexer_weights: torch.Tensor, positions: torch.Tensor, - o_padded: torch.Tensor, + o_padded: "torch.Tensor | QuantizedActivation", ) -> None: """Wide eager region: the whole of ``_prepare_and_attn`` runs eagerly. @@ -558,7 +575,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): indexer_kv_score: torch.Tensor, indexer_weights: torch.Tensor, positions: torch.Tensor, - o_padded: torch.Tensor, + o_padded: "torch.Tensor | QuantizedActivation", ) -> None: """Attention input preparation followed by the sparse indexer and MLA. @@ -723,7 +740,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): q: torch.Tensor, kv: torch.Tensor, positions: torch.Tensor, - out: torch.Tensor, + out: "torch.Tensor | QuantizedActivation", ) -> None: if self.indexer is not None and index_q is not None: assert index_weights is not None @@ -736,7 +753,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): ) # MLA attention writes into the pre-allocated `out` buffer - # ([num_tokens, padded_heads, head_dim]). + # (see _alloc_attn_out). self.forward_mqa(q, kv, positions, out) def _fused_qnorm_rope_kv_insert( diff --git a/vllm/models/deepseek_v4/common/ops/cache_utils.py b/vllm/models/deepseek_v4/common/ops/cache_utils.py index 0cc784e348..cea755bd1a 100644 --- a/vllm/models/deepseek_v4/common/ops/cache_utils.py +++ b/vllm/models/deepseek_v4/common/ops/cache_utils.py @@ -1099,6 +1099,7 @@ def build_flashinfer_mixed_sparse_indices( prefill_left_visible: torch.Tensor | None = None, prefill_right_visible: torch.Tensor | None = None, max_image_tokens: int = 0, + num_rows: int | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Build the FlashInfer DSV4 sparse-index matrix for decode-first batches. @@ -1109,6 +1110,9 @@ def build_flashinfer_mixed_sparse_indices( tokens derive their SWA window from the position and translate local compressed indices to global slots via the block tables. + ``num_rows`` (>= ``num_tokens``) sizes both outputs for a kernel call that + spans a CUDA-graph-padded batch; the rows past ``num_tokens`` stay unset. + When ``prefill_left_visible``/``prefill_right_visible`` are given (vision variant), the SWA column region widens by ``max_image_tokens`` and prefill tokens inside an image span get a bidirectionally widened window; decode @@ -1170,13 +1174,15 @@ def build_flashinfer_mixed_sparse_indices( # by ``sparse_topk_lens``, so padding never changes the attention result. padded_topk = max(topk, decode_compressed_topk) padded_topk = (padded_topk + 3) // 4 * 4 + num_rows = num_tokens if num_rows is None else num_rows + assert num_rows >= num_tokens sparse_indices = torch.empty( - (num_tokens, swa_total_width + padded_topk), + (num_rows, swa_total_width + padded_topk), dtype=torch.int32, device=decode_swa_indices.device, ) sparse_topk_lens = torch.empty( - num_tokens, dtype=torch.int32, device=decode_swa_indices.device + num_rows, dtype=torch.int32, device=decode_swa_indices.device ) if num_tokens == 0: return sparse_indices, sparse_topk_lens @@ -1193,9 +1199,10 @@ def build_flashinfer_mixed_sparse_indices( if compressed_block_span is None else compressed_block_span ) + # The launch grid follows the output rows, so hand over only the real ones. _BUILD_FLASHINFER_MIXED_SPARSE_INDICES_KERNEL( - sparse_indices, - sparse_topk_lens, + sparse_indices[:num_tokens], + sparse_topk_lens[:num_tokens], decode_swa_indices, decode_compressed_indices, decode_compressed_topk_lens, @@ -1322,12 +1329,52 @@ class BuildFlashinferMixedSparseIndicesKernel( SWA_TOTAL_WIDTH: tl.constexpr = SWA_INDEX_WIDTH + IMAGE_WIDTH if token_idx < NUM_DECODE_TOKENS: + # A decode row wider than the window is DSpark's non-causal block: + # the block-anchored window of context followed by the whole block, + # contiguous from column 0. The kernel treats every entry of its + # active ranges as a key -- the first min(WINDOW_SIZE, pos + 1) + # columns and [WINDOW_SIZE, sparse_topk_lens) -- so -1 must not + # appear in either. Fill the causal window with the first visible + # entries and put the rest from column WINDOW_SIZE on, followed + # directly by the compressed entries. + NONCAUSAL: tl.constexpr = SWA_INDEX_WIDTH > WINDOW_SIZE + compressed_start = SWA_TOTAL_WIDTH + if NONCAUSAL: + req_idx = tl.load(token_to_req_indices_ptr + token_idx) + query_start = tl.load(query_start_loc_ptr + req_idx) + query_len = tl.load(query_start_loc_ptr + req_idx + 1) - query_start + context_len = tl.load(seq_lens_ptr + req_idx) - query_len + token_in_query = token_idx - query_start + window_len = tl.minimum(context_len + token_in_query + 1, WINDOW_SIZE) + spill_len = ( + tl.minimum(context_len, WINDOW_SIZE) + query_len - window_len + ) + # A CUDA-graph pad token sits outside its empty request, so + # the kernel never reads its row; keep it the shape of a + # causal pad row (all -1, WINDOW_SIZE long) rather than let + # the formulas above go negative. + is_block_token = (token_in_query >= 0) & (token_in_query < query_len) + window_len = tl.where(is_block_token, window_len, 0) + spill_len = tl.where(is_block_token, spill_len, 0) + compressed_start = WINDOW_SIZE + spill_len for i in range(0, SWA_TOTAL_WIDTH, WINDOW_BLOCK_SIZE): offset = i + tl.arange(0, WINDOW_BLOCK_SIZE) mask = offset < SWA_TOTAL_WIDTH + src = offset + src_mask = offset < SWA_INDEX_WIDTH + if NONCAUSAL: + # Columns from compressed_start on belong to the loops below. + mask = mask & (offset < compressed_start) + spill = offset >= WINDOW_SIZE + src = tl.where(spill, offset - WINDOW_SIZE + window_len, offset) + src_mask = tl.where( + spill, + offset < WINDOW_SIZE + spill_len, + offset < window_len, + ) values = tl.load( - decode_swa_indices_ptr + token_idx * decode_swa_stride + offset, - mask=offset < SWA_INDEX_WIDTH, + decode_swa_indices_ptr + token_idx * decode_swa_stride + src, + mask=src_mask, other=-1, ) values = _remap_flashinfer_index(values, swa_block_size, swa_block_span) @@ -1372,12 +1419,27 @@ class BuildFlashinferMixedSparseIndicesKernel( tl.store( sparse_indices_ptr + token_idx * sparse_indices_stride - + SWA_TOTAL_WIDTH + + compressed_start + offset, values, mask=mask, ) + if NONCAUSAL: + # The row's remaining columns, left over by the shifted + # compressed block. + for i in range(0, SWA_TOTAL_WIDTH - WINDOW_SIZE, WINDOW_BLOCK_SIZE): + offset = i + tl.arange(0, WINDOW_BLOCK_SIZE) + tl.store( + sparse_indices_ptr + + token_idx * sparse_indices_stride + + compressed_start + + PADDED_TOP_K + + offset, + tl.full((WINDOW_BLOCK_SIZE,), -1, dtype=tl.int32), + mask=offset < SWA_TOTAL_WIDTH - compressed_start, + ) + if DECODE_COMPRESSED_TOPK == 0: compressed_len = tl.zeros((), dtype=tl.int32) elif not DECODE_COMPRESSED_INDICES_ARE_LOCAL: @@ -1388,7 +1450,9 @@ class BuildFlashinferMixedSparseIndicesKernel( else: compressed_len = tl.full((), DECODE_COMPRESSED_TOPK, dtype=tl.int32) - tl.store(sparse_topk_lens_ptr + token_idx, SWA_TOTAL_WIDTH + compressed_len) + tl.store( + sparse_topk_lens_ptr + token_idx, compressed_start + compressed_len + ) return prefill_idx = token_idx - NUM_DECODE_TOKENS diff --git a/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py b/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py index bb9e42349a..3e83c280f3 100644 --- a/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py +++ b/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py @@ -9,6 +9,8 @@ import torch from vllm.config import VllmConfig from vllm.config.cache import CacheDType from vllm.forward_context import get_forward_context +from vllm.logger import init_logger +from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation from vllm.models.deepseek_v4.attention import DeepseekV4Attention from vllm.models.deepseek_v4.common.ops import ( build_flashinfer_mixed_sparse_indices, @@ -17,6 +19,8 @@ from vllm.models.deepseek_v4.common.ops import ( from vllm.models.deepseek_v4.nvidia.ops.o_proj import ( compute_fp8_einsum_recipe, deep_gemm_fp8_o_proj, + rope_quant_attn_out, + rope_quant_unsupported_reason, ) from vllm.models.deepseek_v4.sparse_mla import ( DeepseekV4FlashMLAMetadata, @@ -37,6 +41,8 @@ from vllm.v1.attention.backends.mla.sparse_swa import ( if TYPE_CHECKING: from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata +logger = init_logger(__name__) + _FLASHINFER_DSV4_WORKSPACE_BUFFER_SIZE = 128 * 1024 * 1024 _flashinfer_dsv4_workspace_by_device: dict[torch.device, torch.Tensor] = {} @@ -203,7 +209,9 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention): def get_padded_num_q_heads(cls, num_heads: int) -> int: return _pad_to_supported_q_heads(num_heads) - def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + def _o_proj( + self, o: torch.Tensor | QuantizedActivation, positions: torch.Tensor + ) -> torch.Tensor: return deep_gemm_fp8_o_proj( o, positions, @@ -219,9 +227,30 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention): tma_aligned_scales=self._tma_aligned_scales, ) - def __init__(self, *args, **kwargs) -> None: - super().__init__(*args, **kwargs) + def _alloc_attn_out( + self, num_tokens: int, hidden_states: torch.Tensor + ) -> torch.Tensor | QuantizedActivation: + if self._fuse_rope_quant: + return rope_quant_attn_out(num_tokens, hidden_states.device) + return super()._alloc_attn_out(num_tokens, hidden_states) + + def __init__(self, vllm_config: VllmConfig, *args, **kwargs) -> None: + super().__init__(vllm_config, *args, **kwargs) self._einsum_recipe, self._tma_aligned_scales = compute_fp8_einsum_recipe() + # RopeQuant derives each query's RoPE position as seq_len - q_len + i, + # which matches `positions` for DSpark's non-causal draft block too, + # since every call passes the real per-request seq_lens. The one + # exception is a draft block running past max_model_len, whose clamped + # positions it does not reproduce; that only costs those drafts. + reason = rope_quant_unsupported_reason(self) + self._fuse_rope_quant = reason is None + if reason is None: + logger.info_once( + "FLASHINFER_MLA_SPARSE_DSV4 fuses the inverse RoPE and FP8 " + "quant of the attention output." + ) + else: + logger.debug_once("FLASHINFER_MLA_SPARSE_DSV4 RopeQuant off: %s", reason) # Per-tensor FP8 scale buffers + precomputed scalar BMM scales. Only the # per-tensor FP8 cache path consumes these; bf16 reads ``self.scale``. if self.kv_cache_torch_dtype != torch.float8_e4m3fn: @@ -253,32 +282,39 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention): q: torch.Tensor, kv: torch.Tensor, positions: torch.Tensor, - output: torch.Tensor, + output: torch.Tensor | QuantizedActivation, ) -> None: - # The TRTLLM-gen kernel requires h_q in {64, 128}, so the output buffer - # is allocated at the padded head count while q arrives at the local - # head count; _forward pads q to match before the launcher. - assert output.shape[0] == q.shape[0] and output.shape[-1] == q.shape[-1], ( - f"output buffer shape {output.shape} incompatible with q shape {q.shape}" - ) - assert output.shape[1] >= q.shape[1], ( - f"output heads {output.shape[1]} must be >= q heads {q.shape[1]}" - ) - # Per-tensor FP8 q produces a bf16 attention output. - expected_output_dtype = ( - torch.bfloat16 if q.dtype == torch.float8_e4m3fn else q.dtype - ) - assert output.dtype == expected_output_dtype, ( - f"output dtype {output.dtype} must match expected {expected_output_dtype} " - f"for q dtype {q.dtype}" - ) + if isinstance(output, QuantizedActivation): + assert output.data.shape[0] == q.shape[0] + else: + # The TRTLLM-gen kernel requires h_q in {64, 128}, so the output + # buffer is allocated at the padded head count while q arrives at + # the local head count; _forward pads q to match before the launcher. + assert output.shape[0] == q.shape[0] and output.shape[-1] == q.shape[-1], ( + f"output buffer shape {output.shape} incompatible with q {q.shape}" + ) + assert output.shape[1] >= q.shape[1], ( + f"output heads {output.shape[1]} must be >= q heads {q.shape[1]}" + ) + # Per-tensor FP8 q produces a bf16 attention output. + expected_output_dtype = ( + torch.bfloat16 if q.dtype == torch.float8_e4m3fn else q.dtype + ) + assert output.dtype == expected_output_dtype, ( + f"output dtype {output.dtype} must match expected " + f"{expected_output_dtype} for q dtype {q.dtype}" + ) forward_context = get_forward_context() attn_metadata = forward_context.attn_metadata if attn_metadata is None: # Warmup dummy run: FlashInfer reads the cache directly and lazily # allocates its workspace, so nothing to reserve here. - output.zero_() + if isinstance(output, QuantizedActivation): + output.data.zero_() + output.scale.zero_() + else: + output.zero_() return assert isinstance(attn_metadata, dict) @@ -313,11 +349,12 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention): swa_metadata: "DeepseekSparseSWAMetadata", attn_metadata: DeepseekV4FlashMLAMetadata | None, swa_only: bool, + num_rows: int | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Build the combined sparse-index tensors for the mixed batch. Returns ``(compressed_kv_cache, seq_lens, sparse_indices, - sparse_topk_lens)``. + sparse_topk_lens)``, sized ``num_rows`` (default: the real tokens). """ num_decodes = swa_metadata.num_decodes num_prefills = swa_metadata.num_prefills @@ -412,6 +449,8 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention): if swa_only else ("c128a" if self.compress_ratio == 128 else "c4a") ) + if num_rows is not None: + cache_key = f"{cache_key}:{num_rows}" cached_sparse = swa_metadata.flashinfer_sparse_index_cache.get(cache_key, None) if cached_sparse is None: swa_block_span = _packed_block_span(swa_k_cache) @@ -439,8 +478,9 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention): prefill_right_visible=swa_metadata.prefill_right_visible, # getattr for tests that bypass __init__ via object.__new__. max_image_tokens=getattr(self, "max_image_tokens", 0), + num_rows=num_rows, ) - if cache_key != "c4a": + if not cache_key.startswith("c4a"): swa_metadata.flashinfer_sparse_index_cache[cache_key] = ( sparse_indices, sparse_topk_lens, @@ -457,7 +497,7 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention): swa_metadata: "DeepseekSparseSWAMetadata", attn_metadata: DeepseekV4FlashMLAMetadata | None, swa_only: bool, - output: torch.Tensor, + output: torch.Tensor | QuantizedActivation, ) -> None: assert self.kv_cache_torch_dtype in (torch.bfloat16, torch.float8_e4m3fn) num_decodes = swa_metadata.num_decodes @@ -468,6 +508,11 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention): num_tokens = num_decode_tokens + num_prefill_tokens if num_tokens == 0: return + # RopeQuant's output is group-major with a group stride of exactly + # sum_q tokens, so both calls span every row of the (CUDA-graph + # padded) buffers and select their tokens through cum_seq_lens_q + # alone; the kernel never touches rows outside it. + rope_quant = isinstance(output, QuantizedActivation) ( compressed_kv_cache, @@ -480,12 +525,12 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention): swa_metadata=swa_metadata, attn_metadata=attn_metadata, swa_only=swa_only, + num_rows=q.shape[0] if rope_quant else None, ) # CUDA graph execution can pad q/output past the scheduled token count; # restrict to the real tokens (the launcher validates sparse indices). - query = q[:num_tokens] - output = output[:num_tokens] + query = q if rope_quant else q[:num_tokens] bmm1_scale: float | torch.Tensor = self.scale bmm2_scale: float | torch.Tensor = 1.0 if self.kv_cache_torch_dtype == torch.float8_e4m3fn: @@ -499,7 +544,9 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention): # The TRTLLM-gen sparse-MLA kernel requires h_q in {64, 128}; zero-pad # the query heads to the allocated output head count. Padded heads attend # to the shared KV and are sliced off downstream (output is padded too). - padded_heads = output.shape[1] + padded_heads = ( + output.shape[1] if isinstance(output, torch.Tensor) else self.padded_heads + ) if query.shape[1] < padded_heads: padded_query = query.new_zeros( (query.shape[0], padded_heads, query.shape[2]) @@ -512,51 +559,60 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention): query_start_loc_cpu = swa_metadata.query_start_loc_cpu assert query_start_loc is not None and query_start_loc_cpu is not None + def run( + rows: slice, reqs: slice, cum_seq_lens_q: torch.Tensor, max_q_len: int + ) -> None: + if rope_quant: + assert isinstance(output, QuantizedActivation) + rows = slice(None) + out_kwargs = dict( + out=output.data, + dsv4_output_scale=output.scale, + dsv4_inv_rope_cos_sin_cache=self.rotary_emb.cos_sin_cache, + ) + else: + assert isinstance(output, torch.Tensor) + out_kwargs = dict(out=output[rows]) + flashinfer_trtllm_batch_decode_sparse_mla_dsv4( + query=query[rows], + swa_kv_cache=swa_k_cache, + workspace_buffer=workspace, + sparse_indices=sparse_indices[rows], + sparse_indices_are_storage_offsets=True, + compressed_kv_cache=compressed_kv_cache, + sparse_topk_lens=sparse_topk_lens[rows], + seq_lens=seq_lens[reqs], + bmm1_scale=bmm1_scale, + bmm2_scale=bmm2_scale, + sinks=self.attn_sink, + cum_seq_lens_q=cum_seq_lens_q, + max_q_len=max_q_len, + **out_kwargs, + ) + # Keep the TRTLLM-gen decode/prefill split: the launcher is tuned for # uniform-q batches, and this avoids flattening mixed batches into one call. if num_decode_tokens > 0: - decode_cu = query_start_loc[: num_decodes + 1] - flashinfer_trtllm_batch_decode_sparse_mla_dsv4( - query=query[:num_decode_tokens], - swa_kv_cache=swa_k_cache, - workspace_buffer=workspace, - sparse_indices=sparse_indices[:num_decode_tokens], - sparse_indices_are_storage_offsets=True, - compressed_kv_cache=compressed_kv_cache, - sparse_topk_lens=sparse_topk_lens[:num_decode_tokens], - seq_lens=seq_lens[:num_decodes], - out=output[:num_decode_tokens], - bmm1_scale=bmm1_scale, - bmm2_scale=bmm2_scale, - sinks=self.attn_sink, - cum_seq_lens_q=decode_cu, - max_q_len=swa_metadata.max_decode_query_len, + run( + slice(0, num_decode_tokens), + slice(0, num_decodes), + query_start_loc[: num_decodes + 1], + swa_metadata.max_decode_query_len, ) if num_prefill_tokens > 0: - # The prefill query view re-anchors at offset 0, so rebase the - # cumulative query offsets to start at 0. - prefill_cu = ( - query_start_loc[num_decodes : num_reqs + 1] - - query_start_loc[num_decodes] - ) + prefill_cu = query_start_loc[num_decodes : num_reqs + 1] + if not rope_quant: + # The prefill query view re-anchors at offset 0, so rebase the + # cumulative query offsets to start at 0. + prefill_cu = prefill_cu - query_start_loc[num_decodes] prefill_cu_cpu = query_start_loc_cpu[num_decodes : num_reqs + 1] prefill_lens_cpu = prefill_cu_cpu[1:] - prefill_cu_cpu[:-1] - flashinfer_trtllm_batch_decode_sparse_mla_dsv4( - query=query[num_decode_tokens:num_tokens], - swa_kv_cache=swa_k_cache, - workspace_buffer=workspace, - sparse_indices=sparse_indices[num_decode_tokens:num_tokens], - sparse_indices_are_storage_offsets=True, - compressed_kv_cache=compressed_kv_cache, - sparse_topk_lens=sparse_topk_lens[num_decode_tokens:num_tokens], - seq_lens=seq_lens[num_decodes:num_reqs], - out=output[num_decode_tokens:num_tokens], - bmm1_scale=bmm1_scale, - bmm2_scale=bmm2_scale, - sinks=self.attn_sink, - cum_seq_lens_q=prefill_cu, - max_q_len=int(prefill_lens_cpu.max().item()), + run( + slice(num_decode_tokens, num_tokens), + slice(num_decodes, num_reqs), + prefill_cu, + int(prefill_lens_cpu.max().item()), ) diff --git a/vllm/models/deepseek_v4/nvidia/ops/o_proj.py b/vllm/models/deepseek_v4/nvidia/ops/o_proj.py index ec36bf763e..1da4ab0965 100644 --- a/vllm/models/deepseek_v4/nvidia/ops/o_proj.py +++ b/vllm/models/deepseek_v4/nvidia/ops/o_proj.py @@ -5,11 +5,15 @@ from collections.abc import Callable import torch import torch.nn as nn +from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kFp8Dynamic128Sym, +) from vllm.models.deepseek_v4.common.ops.fused_inv_rope_fp8_quant import ( fused_inv_rope_fp8_quant, ) from vllm.platforms import current_platform -from vllm.utils.deep_gemm import fp8_einsum +from vllm.utils.deep_gemm import fp8_einsum, get_tma_aligned_size def compute_fp8_einsum_recipe( @@ -29,7 +33,7 @@ def compute_fp8_einsum_recipe( def deep_gemm_fp8_o_proj( - o: torch.Tensor, + o: torch.Tensor | QuantizedActivation, positions: torch.Tensor, cos_sin_cache: torch.Tensor, wo_a: nn.Module, @@ -49,23 +53,31 @@ def deep_gemm_fp8_o_proj( layer selects the recipe at initialization. ``wo_b`` is any callable over the flattened ``z``: the projection module itself, or a wrapper that also reduce-scatters its output (DeepSeek-V4.1 GEMM-RS). + + A QuantizedActivation ``o`` was already rotated and cast by the attention + kernel, one ``wo_a`` group per slot with the live groups first, so only + the padding groups are sliced off. """ use_fp8 = wo_a.weight.dtype == torch.float8_e4m3fn - o_proj_input, o_scale = fused_inv_rope_fp8_quant( - o, - positions, - cos_sin_cache, - n_groups=n_groups, - heads_per_group=heads_per_group, - nope_dim=nope_dim, - rope_dim=rope_dim, - quant_group_size=einsum_recipe[2], - tma_aligned_scales=tma_aligned_scales, - quantize=use_fp8, - ) + if isinstance(o, QuantizedActivation): + assert use_fp8 + o_proj_input, o_scale = o.data[:, :n_groups], o.scale[:, :n_groups] + else: + o_proj_input, o_scale = fused_inv_rope_fp8_quant( + o, + positions, + cos_sin_cache, + n_groups=n_groups, + heads_per_group=heads_per_group, + nope_dim=nope_dim, + rope_dim=rope_dim, + quant_group_size=einsum_recipe[2], + tma_aligned_scales=tma_aligned_scales, + quantize=use_fp8, + ) z = torch.empty( - (o.shape[0], n_groups, o_lora_rank), - device=o.device, + (o_proj_input.shape[0], n_groups, o_lora_rank), + device=o_proj_input.device, dtype=torch.bfloat16, ) if use_fp8: @@ -89,3 +101,58 @@ def deep_gemm_fp8_o_proj( out=z.transpose(0, 1), ) return wo_b(z.flatten(1)) + + +# TRTLLM-gen's DSv4 RopeQuant epilogue (flashinfer-ai/flashinfer#4918) has one +# fixed schedule: 128 query heads in groups of 8 512-wide heads, and FP8 with +# one packed UE8M0 scale per 128 elements. +_ROPE_QUANT_NUM_HEADS = 128 +_ROPE_QUANT_HEADS_PER_GROUP = 8 +_ROPE_QUANT_HEAD_DIM = 512 +_ROPE_QUANT_BLOCK = 128 + + +def rope_quant_unsupported_reason(layer: torch.nn.Module) -> str | None: + """Why ``layer`` cannot take FlashInfer's fused inverse RoPE + FP8 output. + + DSv4 always has 8 heads per ``wo_a`` group at any TP size, and the rest of + the kernel's fixed shape (head/RoPE dims, per-128 einsum scales, FP32 RoPE + cache) is what the SM100 layer uses anyway. What varies is the local head + count and the configs the layer also serves unquantized. + """ + if layer.n_local_heads != _ROPE_QUANT_NUM_HEADS: + # Padding fewer heads up to the kernel's fixed 128 costs more + # attention than the fused epilogue saves. + return f"RopeQuant needs {_ROPE_QUANT_NUM_HEADS} local query heads" + if layer.kv_cache_torch_dtype != torch.float8_e4m3fn: + return "RopeQuant needs the per-tensor FP8 KV cache" + if layer.wo_a.weight.dtype != torch.float8_e4m3fn: + return "RopeQuant needs an FP8 wo_a" + return None + + +def rope_quant_attn_out(num_tokens: int, device: torch.device) -> QuantizedActivation: + """Allocate the output pair the RopeQuant kernel writes. + + Values are group-major, ``[num_tokens, 16, 8 * 512]`` over physical + ``[group, token, K]``. Scales are MN-major packed UE8M0, one int32 per + head over its four 128-wide blocks. Both are what ``fp8_einsum`` consumes. + """ + n_groups = _ROPE_QUANT_NUM_HEADS // _ROPE_QUANT_HEADS_PER_GROUP + group_width = _ROPE_QUANT_HEADS_PER_GROUP * _ROPE_QUANT_HEAD_DIM + data = torch.empty( + (n_groups, num_tokens, group_width), dtype=torch.float8_e4m3fn, device=device + ).transpose(0, 1) + aligned = get_tma_aligned_size(num_tokens, torch.int32.itemsize) + scale = torch.empty( + (n_groups, group_width // (_ROPE_QUANT_BLOCK * 4), aligned), + dtype=torch.int32, + device=device, + ).permute(2, 0, 1)[:num_tokens] + return QuantizedActivation( + data=data, + scale=scale, + orig_dtype=torch.bfloat16, + orig_shape=data.shape, + quant_key=kFp8Dynamic128Sym, + )