[DeepSeek-V4.1] Bound dense prefill indexer memory (#40217)

2026-09-20 14:43:19 -07:00
parent c2c3629f2d
commit 95521da18d
5 changed files with 641 additions and 94 deletions
@@ -19,7 +19,6 @@ import msgspec
 import torch
 import torch.nn.functional as F

-from sglang.kernels.ops.attention.dsv4 import topk_transform_ragged_v2
 from sglang.kernels.ops.attention.dsv4.decode_attention_sm100 import (
    can_use_swapab_attention,
 )
@@ -61,6 +60,7 @@ from sglang.srt.layers.attention.dsv4.candidate_indexer import (
    CandidateMasks,
    CandidateMetadata,
    IndexerInputs,
+    PrefillCandidateBlocks,
    make_candidate_indexer,
    mask_topk_scores,
    published_masks,
@@ -71,6 +71,7 @@ from sglang.srt.layers.attention.dsv4.compressor_v2 import (
    FusedCompressMetadata,
    create_paged_compressor_data,
 )
+from sglang.srt.layers.attention.dsv4.dense_prefill_indexer import dense_prefill_topk
 from sglang.srt.layers.attention.dsv4.dsv41_sparse import (
    _rope_fq4,
    token_req_indices,
@@ -313,21 +314,6 @@ def _has_dense_fp4_indexer() -> bool:
    return hasattr(deep_gemm, "fp8_fp4_mqa_logits")


-def _dense_fp4_mqa_logits(
-    q_fp4: Tuple[torch.Tensor, torch.Tensor],
-    kv_fp4: Tuple[torch.Tensor, torch.Tensor],
-    weights: torch.Tensor,
-    ks: torch.Tensor,
-    ke: torch.Tensor,
-    max_seqlen_k: int,
-) -> torch.Tensor:
-    from deep_gemm import fp8_fp4_mqa_logits as fn
-
-    # q (int8 [T, H, 64], int32 [T, H]) x kv (int8 [L, 64], int32 [L]) -> fp32
-    # [T, max_seqlen_k]; row t column j is k[ks_t + j], garbage past ke_t - ks_t.
-    return fn(q_fp4, kv_fp4, weights, ks, ke, False, max_seqlen_k)
-
-
 def _low_ratio_source_projections(layer, x, q_lora, positions, bufs):
    from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
        get_tc_piecewise_forward_context,
@@ -1762,7 +1748,9 @@ class DeepseekV4AttnBackend(
        # TODO(candidate): goes away once the source publishes its tail rows straight
        # onto the tail metadata (publish_prefill); until then cut the full masks.
        full_masks = self.forward_metadata.candidate_metadata
-        if isinstance(full_masks, CandidateMasks) and full_masks.request_masks:
+        if isinstance(full_masks, PrefillCandidateBlocks):
+            tail_metadata.candidate_metadata = full_masks.tail(tail_lens_cpu)
+        elif isinstance(full_masks, CandidateMasks) and full_masks.request_masks:
            tail_metadata.candidate_metadata = CandidateMasks(
                request_masks=[
                    mask[mask.shape[0] - t :]
@@ -3252,13 +3240,14 @@ class DeepseekV4AttnBackend(
                // ratio
            )
            start += lc
-        empty_mask = torch.zeros(0, 0, dtype=torch.bool, device=device)
        num_tokens = pos.shape[0]
-        # TODO(candidate): move this to candidate indexer
        if not slot_chunks or num_tokens == 0:
            if indexer.is_candidate_source:
-                self.forward_metadata.candidate_metadata = CandidateMasks(
-                    request_masks=[empty_mask for _ in lc_per_req]
+                self.forward_metadata.candidate_metadata = PrefillCandidateBlocks(
+                    request_blocks=[
+                        torch.empty((length, 0), dtype=torch.int32, device=device)
+                        for length in q_lens_cpu
+                    ]
                )
            return
        k_slots = torch.cat(slot_chunks)
@@ -3276,26 +3265,26 @@ class DeepseekV4AttnBackend(
            q_lens.to(torch.int64),
            output_size=num_tokens,
        )
-        logits = _dense_fp4_mqa_logits(
-            (q_fp4, q_sf),
-            (k_fp4, k_sf),
-            weights,
-            ks,
-            ks + compress_lens,
-            # the fused top-k reads score rows through 16-byte vectors
-            ceil_align(max(lc_per_req), 4),
-        )
-        if indexer.is_candidate_source or indexer.uses_candidates:
-            self._publish_or_consume_candidates(
-                indexer, logits, compress_lens, lc_per_req, q_lens_cpu, empty_mask
-            )
-        topk = indexer.index_topk
-        selected = torch.empty((num_tokens, topk), dtype=torch.int32, device=device)
-        topk_transform_ragged_v2(
-            logits, compress_lens, out_offsets=ks, out_indices=selected
-        )
+        candidates = None
        if indexer.uses_candidates and not indexer.is_candidate_source:
-            selected = mask_topk_scores(logits, selected, ks)
+            candidates = self.forward_metadata.candidate_metadata
+            assert isinstance(candidates, PrefillCandidateBlocks)
+        topk = indexer.index_topk
+        selected, published = dense_prefill_topk(
+            q=(q_fp4, q_sf),
+            kv=(k_fp4, k_sf),
+            weights=weights,
+            starts=ks,
+            lengths=compress_lens,
+            request_lengths=list(zip(q_lens_cpu, lc_per_req)),
+            topk=topk,
+            candidate_topk_blocks=indexer.candidate_topk_blocks,
+            candidate_block_size=indexer.candidate_block_size,
+            publish_candidates=indexer.is_candidate_source,
+            candidates=candidates,
+        )
+        if published is not None:
+            self.forward_metadata.candidate_metadata = published
        # ascending positions, padding last: the layout the consumers expect
        unselected = torch.iinfo(torch.int32).max
        selected = selected.masked_fill(selected < 0, unselected).sort(dim=-1).values
@@ -3308,50 +3297,6 @@ class DeepseekV4AttnBackend(
                chosen, selected - ks[:, None], -1
            )

-    # TODO(candidate): dense-prefill level one / level two inline with masks; move
-    # into the candidate indexer as publish_prefill / select_prefill.
-    def _publish_or_consume_candidates(
-        self, indexer, logits, compress_lens, lc_per_req, q_lens_cpu, empty_mask
-    ) -> None:
-        publish = [] if indexer.is_candidate_source else None
-        consume = (
-            None
-            if publish is not None
-            else published_masks(self.forward_metadata.candidate_metadata)
-        )
-        j = torch.arange(logits.shape[1], device=logits.device)
-        tok_start = 0
-        for b, (lc, t_len) in enumerate(zip(lc_per_req, q_lens_cpu)):
-            rows = slice(tok_start, tok_start + t_len)
-            tok_start += t_len
-            if lc == 0 or t_len == 0:
-                if publish is not None:
-                    publish.append(empty_mask)
-                continue
-            scores = logits[rows, :lc]
-            if publish is None:
-                scores.masked_fill_(~consume.request_masks[b], -torch.inf)
-                continue
-            lens = compress_lens[rows, None]
-            # the block selection tells unreachable positions apart by -inf
-            scores.masked_fill_(j[None, :lc] >= lens, -torch.inf)
-            # the block selection pads and pools a copy of its rows; bound that copy
-            step = max(1, _TORCH_INDEXER_SCORE_BUDGET_BYTES // (lc * 4))
-            masks = [
-                select_candidate_blocks(
-                    scores[start : start + step],
-                    lens[start : start + step],
-                    topk_blocks=indexer.candidate_topk_blocks,
-                    block_size=indexer.candidate_block_size,
-                )
-                for start in range(0, t_len, step)
-            ]
-            publish.append(masks[0] if len(masks) == 1 else torch.cat(masks))
-        if publish is not None:
-            self.forward_metadata.candidate_metadata = CandidateMasks(
-                request_masks=publish
-            )
-
    def _low_ratio_index_topk_prefill_graph(self, layer, pos, q, w) -> None:
        from sglang.kernels.ops.attention.dsv4.fp4_indexer import (
            quantize_fp4_indexer_tensor,
@@ -3,6 +3,7 @@ from __future__ import annotations
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, List, Optional, Union

+import msgspec
 import torch
 import torch.nn.functional as F

@@ -70,6 +71,18 @@ class CandidateMasks(CandidateMetadata):
    request_masks: Optional[List[torch.Tensor]] = None  # prefill: [rows_b, lc_b] each


+class PrefillCandidateBlocks(CandidateMetadata, msgspec.Struct):
+    request_blocks: List[torch.Tensor]
+
+    def tail(self, lengths: List[int]) -> PrefillCandidateBlocks:
+        return PrefillCandidateBlocks(
+            request_blocks=[
+                blocks[blocks.shape[0] - length :]
+                for blocks, length in zip(self.request_blocks, lengths)
+            ]
+        )
+
+
 def published_masks(candidate) -> CandidateMasks:
    assert isinstance(candidate, CandidateMasks), "candidate masks missing"
    return candidate
@@ -91,18 +104,15 @@ def mask_topk_scores(
    return indices.masked_fill(~valid, -1)


-def select_candidate_blocks(
+def _candidate_block_topk(
    logits: torch.Tensor,
    compress_lens: Union[torch.Tensor, int],
    topk_blocks: int,
    block_size: int,
-) -> torch.Tensor:
-    """Level one of the two-level top-k: a bool mask over positions keeping the
-    topk_blocks best-scoring blocks per query. Unreachable positions are already -inf
-    in logits, so an all -inf block means not reachable yet; the block holding the
-    query's newest position is always kept."""
+) -> torch.return_types.topk:
    width = logits.size(-1)
-    scores = F.pad(logits, (0, -width % block_size), value=-torch.inf)
+    padding = -width % block_size
+    scores = F.pad(logits, (0, padding), value=-torch.inf) if padding else logits
    scores = scores.unflatten(-1, (-1, block_size)).amax(dim=-1)
    num_blocks = scores.size(-1)

@@ -111,8 +121,50 @@ def select_candidate_blocks(
        torch.arange(num_blocks, device=logits.device) == last, torch.inf
    )

-    top = scores.topk(min(topk_blocks, num_blocks), dim=-1)
-    keep = torch.zeros_like(scores, dtype=torch.bool).scatter_(
-        -1, top.indices, top.values > -torch.inf
+    return scores.topk(min(topk_blocks, num_blocks), dim=-1)
+
+
+def select_candidate_block_ids(
+    logits: torch.Tensor,
+    compress_lens: Union[torch.Tensor, int],
+    topk_blocks: int,
+    block_size: int,
+) -> torch.Tensor:
+    top = _candidate_block_topk(
+        logits=logits,
+        compress_lens=compress_lens,
+        topk_blocks=topk_blocks,
+        block_size=block_size,
    )
+    return top.indices.to(torch.int32).masked_fill_(~(top.values > -torch.inf), -1)
+
+
+def candidate_block_mask(
+    blocks: torch.Tensor, width: int, block_size: int
+) -> torch.Tensor:
+    num_blocks = (width + block_size - 1) // block_size
+    keep = torch.zeros(
+        (*blocks.shape[:-1], num_blocks + 1), dtype=torch.bool, device=blocks.device
+    )
+    keep.scatter_(-1, blocks.to(torch.int64).masked_fill(blocks < 0, num_blocks), True)
+    return keep[..., :num_blocks].repeat_interleave(block_size, dim=-1)[..., :width]
+
+
+def select_candidate_blocks(
+    logits: torch.Tensor,
+    compress_lens: Union[torch.Tensor, int],
+    topk_blocks: int,
+    block_size: int,
+) -> torch.Tensor:
+    top = _candidate_block_topk(
+        logits=logits,
+        compress_lens=compress_lens,
+        topk_blocks=topk_blocks,
+        block_size=block_size,
+    )
+    width = logits.shape[-1]
+    num_blocks = (width + block_size - 1) // block_size
+    keep = torch.zeros(
+        (*logits.shape[:-1], num_blocks), dtype=torch.bool, device=logits.device
+    ).scatter_(-1, top.indices, top.values > -torch.inf)
    return keep.repeat_interleave(block_size, dim=-1)[..., :width]
@@ -0,0 +1,148 @@
+from __future__ import annotations
+
+import torch
+
+from sglang.srt.layers.attention.dsv4.candidate_indexer import (
+    PrefillCandidateBlocks,
+    candidate_block_mask,
+    mask_topk_scores,
+    select_candidate_block_ids,
+)
+from sglang.srt.layers.attention.mqa_logits_utils import (
+    mqa_logits_row_bytes,
+    mqa_logits_rows_per_chunk,
+)
+from sglang.srt.utils.common import ceil_align
+
+# TODO: use a per-forward mqa_logits_budget_bytes() budget that also
+# leaves room for candidate masks and block-selection scratch.
+_SCORE_BUDGET_BYTES = 2 << 30
+
+
+def dense_prefill_topk(
+    *,
+    q: tuple[torch.Tensor, torch.Tensor],
+    kv: tuple[torch.Tensor, torch.Tensor],
+    weights: torch.Tensor,
+    starts: torch.Tensor,
+    lengths: torch.Tensor,
+    request_lengths: list[tuple[int, int]],
+    topk: int,
+    candidate_topk_blocks: int,
+    candidate_block_size: int,
+    publish_candidates: bool,
+    candidates: PrefillCandidateBlocks | None,
+) -> tuple[torch.Tensor, PrefillCandidateBlocks | None]:
+    selected = torch.full(
+        (q[0].shape[0], topk), -1, dtype=torch.int32, device=weights.device
+    )
+    published = (
+        PrefillCandidateBlocks(request_blocks=[]) if publish_candidates else None
+    )
+    request_ranges = []
+    row = 0
+    for query_length, context_length in request_lengths:
+        request_ranges.append((row, row + query_length, context_length))
+        if published is not None:
+            num_blocks = (
+                context_length + candidate_block_size - 1
+            ) // candidate_block_size
+            published.request_blocks.append(
+                torch.empty(
+                    (query_length, min(candidate_topk_blocks, num_blocks)),
+                    dtype=torch.int32,
+                    device=weights.device,
+                )
+            )
+        row += query_length
+    width = ceil_align(max((n for _, n in request_lengths), default=0), 4)
+    if row == 0 or width == 0:
+        return selected, published
+    row_alignment = 128 // q[0].shape[1]
+    rows_per_chunk = mqa_logits_rows_per_chunk(
+        num_rows=ceil_align(row, row_alignment),
+        row_bytes=mqa_logits_row_bytes(width),
+        budget_bytes=_SCORE_BUDGET_BYTES,
+    )
+    if rows_per_chunk is None:
+        rows_per_chunk = row
+    else:
+        rows_per_chunk = max(
+            row_alignment, rows_per_chunk // row_alignment * row_alignment
+        )
+    for offset in range(0, row, rows_per_chunk):
+        rows = slice(offset, min(offset + rows_per_chunk, row))
+        _select_tile(
+            q=(q[0][rows], q[1][rows]),
+            kv=kv,
+            weights=weights[rows],
+            starts=starts[rows],
+            lengths=lengths[rows],
+            width=width,
+            selected=selected[rows],
+            block_size=candidate_block_size,
+            row_offset=offset,
+            request_ranges=request_ranges,
+            publish=published,
+            consume=candidates,
+        )
+    return selected, published
+
+
+def _select_tile(
+    *,
+    q: tuple[torch.Tensor, torch.Tensor],
+    kv: tuple[torch.Tensor, torch.Tensor],
+    weights: torch.Tensor,
+    starts: torch.Tensor,
+    lengths: torch.Tensor,
+    width: int,
+    selected: torch.Tensor,
+    block_size: int,
+    row_offset: int,
+    request_ranges: list[tuple[int, int, int]],
+    publish: PrefillCandidateBlocks | None,
+    consume: PrefillCandidateBlocks | None,
+) -> None:
+    from deep_gemm import fp8_fp4_mqa_logits
+
+    from sglang.kernels.ops.attention.dsv4 import topk_transform_ragged_v2
+
+    logits = fp8_fp4_mqa_logits(q, kv, weights, starts, starts + lengths, False, width)
+    if publish is not None or consume is not None:
+        for request, (start, end, context_length) in enumerate(request_ranges):
+            begin, stop = max(start, row_offset), min(end, row_offset + logits.shape[0])
+            if begin >= stop or context_length == 0:
+                continue
+            rows = slice(begin - row_offset, stop - row_offset)
+            request_rows = slice(begin - start, stop - start)
+            scores = logits[rows, :context_length]
+            if publish is not None:
+                lens = lengths[rows, None]
+                scores.masked_fill_(
+                    torch.arange(context_length, device=logits.device)[None, :] >= lens,
+                    -torch.inf,
+                )
+                blocks = publish.request_blocks[request][request_rows]
+                blocks.copy_(
+                    select_candidate_block_ids(
+                        logits=scores,
+                        compress_lens=lens,
+                        topk_blocks=blocks.shape[1],
+                        block_size=block_size,
+                    )
+                )
+            else:
+                scores.masked_fill_(
+                    ~candidate_block_mask(
+                        blocks=consume.request_blocks[request][request_rows],
+                        width=context_length,
+                        block_size=block_size,
+                    ),
+                    -torch.inf,
+                )
+    topk_transform_ragged_v2(logits, lengths, out_offsets=starts, out_indices=selected)
+    if consume is not None:
+        selected.copy_(
+            mask_topk_scores(scores=logits, indices=selected, offsets=starts)
+        )