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

python/sglang/srt/layers/attention/deepseek_v4_backend.py

@@ -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(
+                self.forward_metadata.candidate_metadata = PrefillCandidateBlocks(
-                    request_masks=[empty_mask for _ in lc_per_req]
+                    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(
+        candidates = None
-            (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
-        )
         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,

python/sglang/srt/layers/attention/dsv4/candidate_indexer.py

@@ -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:
+) -> torch.return_types.topk:
-    """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."""
     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)
+    return 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
+
+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]

python/sglang/srt/layers/attention/dsv4/dense_prefill_indexer.py

@@ -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)
+        )

test/registered/kernels/ops/attention/test_dense_prefill_indexer.py

@@ -0,0 +1,316 @@
+import unittest
+from unittest.mock import patch
+
+import torch
+
+from sglang.kernels.ops.attention.dsv4.fp4_indexer import quantize_fp4_indexer_tensor
+from sglang.srt.layers.attention.dsv4 import dense_prefill_indexer
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=120, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
+
+
+def make_inputs(request_lengths, ratio=1, zero_queries=False, seed=17):
+    torch.manual_seed(seed)
+    rows = sum(q for q, _ in request_lengths)
+    q = torch.randn((rows, 32, 128), dtype=torch.bfloat16, device="cuda")
+    if zero_queries:
+        q.zero_()
+    packed, scales = quantize_fp4_indexer_tensor(q.flatten(0, 1), rne=True)
+    kv = quantize_fp4_indexer_tensor(
+        torch.randn(
+            (sum(n for _, n in request_lengths), 128),
+            dtype=torch.bfloat16,
+            device="cuda",
+        ),
+        rne=True,
+    )
+    starts, lengths = [], []
+    start = 0
+    for queries, context in request_lengths:
+        starts.extend([start] * queries)
+        lengths.extend(
+            (position + 1) // ratio
+            for position in range(context * ratio - queries, context * ratio)
+        )
+        start += context
+    return dict(
+        q=(packed.view(rows, 32, 64), scales.view(rows, 32)),
+        kv=kv,
+        weights=torch.rand((rows, 32), dtype=torch.float32, device="cuda"),
+        starts=torch.tensor(starts, dtype=torch.int32, device="cuda"),
+        lengths=torch.tensor(lengths, dtype=torch.int32, device="cuda"),
+        request_lengths=request_lengths,
+        topk=512,
+        candidate_topk_blocks=2,
+        candidate_block_size=8,
+    )
+
+
+def dense_scores(inputs):
+    from deep_gemm import fp8_fp4_mqa_logits
+
+    width = (max(n for _, n in inputs["request_lengths"]) + 3) // 4 * 4
+    scores = fp8_fp4_mqa_logits(
+        inputs["q"],
+        inputs["kv"],
+        inputs["weights"],
+        inputs["starts"],
+        inputs["starts"] + inputs["lengths"],
+        False,
+        width,
+    )
+    return scores.masked_fill_(
+        torch.arange(width, device="cuda")[None, :] >= inputs["lengths"][:, None],
+        -torch.inf,
+    )
+
+
+@unittest.skipUnless(
+    torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 10,
+    "requires SM100",
+)
+class TestDensePrefillIndexer(CustomTestCase):
+    def assert_topk(self, inputs, selected, scores):
+        columns = (selected - inputs["starts"][:, None]).long()
+        valid = selected >= 0
+        expected_count = torch.isfinite(scores).sum(-1).clamp_max(inputs["topk"])
+        torch.testing.assert_close(valid.sum(-1), expected_count)
+        self.assertTrue(
+            (~valid | ((columns >= 0) & (columns < inputs["lengths"][:, None]))).all()
+        )
+        actual = scores.gather(1, columns.clamp(0, scores.shape[1] - 1)).masked_fill(
+            ~valid, -torch.inf
+        )
+        expected = scores.topk(min(inputs["topk"], scores.shape[1]), dim=-1).values
+        expected = torch.nn.functional.pad(
+            expected, (0, inputs["topk"] - expected.shape[1]), value=-torch.inf
+        )
+        torch.testing.assert_close(
+            actual.sort(descending=True).values, expected, rtol=1e-5, atol=1e-5
+        )
+        ordered = (
+            columns.masked_fill(~valid, torch.iinfo(torch.int64).max).sort().values
+        )
+        self.assertTrue(
+            (
+                (ordered[:, 1:] != ordered[:, :-1])
+                | (ordered[:, 1:] == torch.iinfo(torch.int64).max)
+            ).all()
+        )
+
+    def test_ragged_source_consumer_and_replay(self):
+        for ratio in (1, 2):
+            for zero_queries in (False, True):
+                with self.subTest(ratio=ratio, zero_queries=zero_queries):
+                    inputs = make_inputs(
+                        [(0, 0), (1, 1), (33, 511), (257, 4097)],
+                        ratio=ratio,
+                        zero_queries=zero_queries,
+                    )
+                    inputs["candidate_topk_blocks"] = 128
+                    scores = dense_scores(inputs)
+                    consumer_inputs = make_inputs(
+                        inputs["request_lengths"],
+                        ratio=ratio,
+                        zero_queries=zero_queries,
+                        seed=29,
+                    )
+                    consumer_inputs["kv"] = inputs["kv"]
+                    consumer_inputs["candidate_topk_blocks"] = 128
+                    consumer_scores = dense_scores(consumer_inputs)
+                    with patch.object(
+                        dense_prefill_indexer, "_SCORE_BUDGET_BYTES", 128 << 10
+                    ):
+                        selected, candidates = dense_prefill_indexer.dense_prefill_topk(
+                            **inputs, publish_candidates=True, candidates=None
+                        )
+                        self.assert_topk(inputs, selected, scores)
+                        row = 0
+                        for (queries, context), blocks in zip(
+                            inputs["request_lengths"], candidates.request_blocks
+                        ):
+                            local = scores[row : row + queries, :context]
+                            if queries and context:
+                                padded = torch.nn.functional.pad(
+                                    local, (0, -context % 8), value=-torch.inf
+                                )
+                                block_scores = padded.unflatten(-1, (-1, 8)).amax(-1)
+                                last = (inputs["lengths"][row : row + queries] - 1) // 8
+                                block_scores.masked_fill_(
+                                    torch.arange(block_scores.shape[1], device="cuda")[
+                                        None, :
+                                    ]
+                                    == last[:, None],
+                                    torch.inf,
+                                )
+                                chosen_scores = block_scores.gather(
+                                    1, blocks.long().clamp_min(0)
+                                ).masked_fill(blocks < 0, -torch.inf)
+                                torch.testing.assert_close(
+                                    chosen_scores.sort(descending=True).values,
+                                    block_scores.topk(blocks.shape[1]).values,
+                                    rtol=1e-5,
+                                    atol=1e-5,
+                                )
+                                columns = torch.arange(context, device="cuda")
+                                member = (
+                                    columns[None, :, None] // 8 == blocks[:, None, :]
+                                ).any(-1)
+                                consumer_scores[
+                                    row : row + queries, :context
+                                ].masked_fill_(
+                                    ~member,
+                                    -torch.inf,
+                                )
+                            row += queries
+                        self.assertGreater(
+                            torch.isfinite(consumer_scores[-1]).sum().item(),
+                            inputs["topk"],
+                        )
+                        self.assertLess(
+                            torch.isfinite(consumer_scores[-1]).sum().item(),
+                            inputs["lengths"][-1].item(),
+                        )
+                        selected, published = dense_prefill_indexer.dense_prefill_topk(
+                            **consumer_inputs,
+                            publish_candidates=False,
+                            candidates=candidates,
+                        )
+                        self.assertIsNone(published)
+                        self.assert_topk(consumer_inputs, selected, consumer_scores)
+                        tail_lengths = [0, 0, 7, 31]
+                        rows, row = [], 0
+                        for (queries, _), tail in zip(
+                            inputs["request_lengths"], tail_lengths
+                        ):
+                            rows.extend(range(row + queries - tail, row + queries))
+                            row += queries
+                        rows = torch.tensor(rows, dtype=torch.int64, device="cuda")
+                        tail_inputs = dict(
+                            consumer_inputs,
+                            q=tuple(t[rows] for t in consumer_inputs["q"]),
+                            weights=consumer_inputs["weights"][rows],
+                            starts=inputs["starts"][rows],
+                            lengths=inputs["lengths"][rows],
+                            request_lengths=list(
+                                zip(
+                                    tail_lengths,
+                                    [n for _, n in inputs["request_lengths"]],
+                                )
+                            ),
+                        )
+                        selected, _ = dense_prefill_indexer.dense_prefill_topk(
+                            **tail_inputs,
+                            publish_candidates=False,
+                            candidates=candidates.tail(tail_lengths),
+                        )
+                        self.assert_topk(tail_inputs, selected, consumer_scores[rows])
+
+    def test_unfiltered_and_zero_length_requests(self):
+        for request_lengths in ([(257, 8192)], [(1, 0), (1, 1), (0, 7)]):
+            for publish in (False, True):
+                with self.subTest(request_lengths=request_lengths, publish=publish):
+                    inputs = make_inputs(request_lengths)
+                    selected, candidates = dense_prefill_indexer.dense_prefill_topk(
+                        **inputs, publish_candidates=publish, candidates=None
+                    )
+                    self.assert_topk(inputs, selected, dense_scores(inputs))
+                    if publish:
+                        self.assertEqual(
+                            [b.shape[0] for b in candidates.request_blocks],
+                            [q for q, _ in request_lengths],
+                        )
+
+    def test_empty_queries_or_context(self):
+        for request_lengths, shape, block_shapes in (
+            ([], (0, 512), []),
+            ([(0, 0)], (0, 512), [(0, 0)]),
+            ([(0, 0), (0, 17)], (0, 512), [(0, 0), (0, 2)]),
+            ([(1, 0)], (1, 512), [(1, 0)]),
+        ):
+            with self.subTest(request_lengths=request_lengths):
+                inputs = make_inputs(request_lengths)
+                selected, candidates = dense_prefill_indexer.dense_prefill_topk(
+                    **inputs, publish_candidates=True, candidates=None
+                )
+                torch.testing.assert_close(
+                    selected, torch.full(shape, -1, dtype=torch.int32, device="cuda")
+                )
+                self.assertEqual(
+                    [tuple(b.shape) for b in candidates.request_blocks], block_shapes
+                )
+
+    def test_score_budget_includes_allocation_padding(self):
+        from deep_gemm import fp8_fp4_mqa_logits
+
+        inputs = make_inputs([(13, 257)])
+        expected = dense_scores(inputs)
+        for budget in (32 << 10, 28 << 10, 8 << 10):
+            with self.subTest(budget=budget):
+                allocations = []
+
+                def checked_logits(*args, **kwargs):
+                    before = torch.cuda.memory_allocated()
+                    logits = fp8_fp4_mqa_logits(*args, **kwargs)
+                    allocations.append(torch.cuda.memory_allocated() - before)
+                    return logits
+
+                with (
+                    patch.object(dense_prefill_indexer, "_SCORE_BUDGET_BYTES", budget),
+                    patch("deep_gemm.fp8_fp4_mqa_logits", new=checked_logits),
+                ):
+                    selected, _ = dense_prefill_indexer.dense_prefill_topk(
+                        **inputs, publish_candidates=True, candidates=None
+                    )
+                self.assertTrue(allocations)
+                self.assertLessEqual(max(allocations), budget)
+                self.assert_topk(inputs, selected, expected)
+
+    def test_score_memory_is_bounded(self):
+        for context, limit_gib in ((65536, 3), (65535, 5)):
+            with self.subTest(context=context):
+                inputs = make_inputs([(16384, context)])
+                inputs["candidate_topk_blocks"] = 2048
+                torch.cuda.synchronize()
+                torch.cuda.reset_peak_memory_stats()
+                baseline = torch.cuda.memory_allocated()
+                selected, candidates = dense_prefill_indexer.dense_prefill_topk(
+                    **inputs, publish_candidates=True, candidates=None
+                )
+                torch.cuda.synchronize()
+                self.assertLess(
+                    torch.cuda.max_memory_allocated() - baseline, limit_gib << 30
+                )
+                self.assertEqual(tuple(selected.shape), (16384, 512))
+                self.assertEqual(
+                    tuple(candidates.request_blocks[0].shape), (16384, 2048)
+                )
+                del selected
+                selected, published = dense_prefill_indexer.dense_prefill_topk(
+                    **inputs, publish_candidates=False, candidates=candidates
+                )
+                self.assertIsNone(published)
+                del selected
+                torch.cuda.synchronize()
+                baseline = torch.cuda.memory_allocated()
+                for _ in range(3):
+                    torch.cuda.reset_peak_memory_stats()
+                    selected, published = dense_prefill_indexer.dense_prefill_topk(
+                        **inputs, publish_candidates=False, candidates=candidates
+                    )
+                    torch.cuda.synchronize()
+                    self.assertIsNone(published)
+                    self.assertLess(
+                        torch.cuda.max_memory_allocated() - baseline, 4 << 30
+                    )
+                    self.assertEqual(tuple(selected.shape), (16384, 512))
+                    del selected
+                    torch.cuda.synchronize()
+                    self.assertEqual(torch.cuda.memory_allocated(), baseline)
+                del inputs, candidates
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/layers/test_dsv41_candidate_blocks.py

@@ -0,0 +1,86 @@
+import unittest
+
+import torch
+
+from sglang.srt.layers.attention.dsv4.candidate_indexer import (
+    PrefillCandidateBlocks,
+    candidate_block_mask,
+    select_candidate_block_ids,
+    select_candidate_blocks,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=3, suite="base-a-test-cpu")
+
+
+class TestPrefillCandidateBlocks(CustomTestCase):
+    def test_causal_partial_blocks_and_forced_newest_block(self):
+        scores = torch.tensor([[100.0] * 8 + [50.0] * 8 + [-10.0] * 3] * 4)
+        lengths = torch.tensor([[0], [1], [9], [19]])
+        scores.masked_fill_(torch.arange(19)[None, :] >= lengths, -torch.inf)
+        original_scores = scores.clone()
+        expected = torch.tensor(
+            [
+                [False] * 19,
+                [True] * 8 + [False] * 11,
+                [True] * 16 + [False] * 3,
+                [True] * 8 + [False] * 8 + [True] * 3,
+            ]
+        )
+        blocks = select_candidate_block_ids(
+            logits=scores, compress_lens=lengths, topk_blocks=2, block_size=8
+        )
+        self.assertEqual(blocks.dtype, torch.int32)
+        self.assertEqual(tuple(blocks.shape), (4, 2))
+        torch.testing.assert_close(blocks[0], torch.tensor([-1, -1], dtype=torch.int32))
+        torch.testing.assert_close(
+            candidate_block_mask(blocks=blocks, width=19, block_size=8), expected
+        )
+        torch.testing.assert_close(
+            select_candidate_blocks(
+                logits=scores, compress_lens=lengths, topk_blocks=2, block_size=8
+            ),
+            expected,
+        )
+        torch.testing.assert_close(scores, original_scores)
+
+    def test_underfilled_and_empty_candidates(self):
+        for width in (0, 1, 7, 8, 9):
+            with self.subTest(width=width):
+                scores = torch.zeros((2, width))
+                blocks = select_candidate_block_ids(
+                    logits=scores, compress_lens=width, topk_blocks=2048, block_size=8
+                )
+                torch.testing.assert_close(
+                    candidate_block_mask(blocks=blocks, width=width, block_size=8),
+                    torch.ones_like(scores, dtype=torch.bool),
+                )
+        blocks = torch.full((3, 2), -1, dtype=torch.int32)
+        self.assertFalse(
+            candidate_block_mask(blocks=blocks, width=19, block_size=8).any()
+        )
+
+    def test_replay_tail_keeps_request_boundaries_and_empty_tails(self):
+        requests = [torch.arange(n * 2).reshape(n, 2) for n in (5, 0, 3)]
+        candidates = PrefillCandidateBlocks(request_blocks=requests)
+        tail = candidates.tail([2, 0, 0])
+        self.assertEqual(
+            [tuple(b.shape) for b in tail.request_blocks], [(2, 2), (0, 2), (0, 2)]
+        )
+        torch.testing.assert_close(tail.request_blocks[0], requests[0][3:])
+        self.assertEqual(tail.request_blocks[0].data_ptr(), requests[0][3:].data_ptr())
+        self.assertEqual([b.shape[0] for b in candidates.request_blocks], [5, 0, 3])
+
+    def test_nonfinite_blocks_match_mask_selection(self):
+        logits = torch.tensor([[float("nan")] * 8 + [1.0] * 8 + [-torch.inf] * 8])
+        kwargs = dict(logits=logits, compress_lens=24, topk_blocks=3, block_size=8)
+        blocks = select_candidate_block_ids(**kwargs)
+        torch.testing.assert_close(
+            candidate_block_mask(blocks=blocks, width=24, block_size=8),
+            select_candidate_blocks(**kwargs),
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()