[DSv4.1] Score candidate blocks in paged decode consumer layers

[ormandj]

Motivation

In the dense-blocks candidate scheme, a source layer publishes up to 2,048 blocks of 8 compressed positions per decode row, and consumer layers select their top-k only inside those blocks. On the paged decode path (SM90 today; SM120 with #42371), the source turns the block IDs into a boolean mask over the full CUDA-graph capacity (candidate_block_mask), and every consumer still runs fp4_index_logits_paged over every tile of the visible context, with non-candidate positions masked out. Consumer cost therefore grows with context, even though at most 16,384 positions can be selected.

Modifications

• fp4_index_logits_candidates (dsv4/fp4_indexer.py) fills the visible range with -inf, then scores only the candidate blocks with the same _fp4_index_logits_tile. Tiles whose blocks are all padding or past the visible length are skipped. Visible scores match the masked scan.
• decode_scores takes candidate_blocks instead of candidate_mask; the paged consumer passes the published block IDs. BlockIds.decode_mask and candidate_block_mask are removed (the mask moves into the test as the reference). The non-paged path is unchanged.

Accuracy Tests

• test_fp4_candidate_logits_match_masked_scan replays a CUDA graph with lengths (including 0 and tile boundaries), queries and block IDs changing between replays, and asserts that the candidate scorer's visible scores are bitwise equal to the masked fp4_index_logits_paged (random BF16 queries and weights, 32 and 64 heads, ratios 1 and 2, padded block lists, one row without candidates).
• test_hopper_indexer_backends_replay checks that the source and full top-k use the paged scorer and the consumer uses the candidate scorer, and still checks the selections against the torch reference.

On one RTX PRO 6000 (SM120), with this PR applied to a tree that also carries #42371 and #42861, python -m pytest test/registered/kernels/ops/attention/test_fp4_indexer.py: 104 passed and 36 skipped (test_fp4_logits_visibility_boundaries cases). The passes include the candidate cases [6-32-1-193-None], [6-64-2-8193-None], [4-64-1-32769-None] and [64-64-1-32769-1048580], and test_paged_indexer_backends_replay for ratios 1 and 2, which is #42371's SM120-enabled form of test_hopper_indexer_backends_replay. On this PR alone, test_hopper_indexer_backends_replay is gated to capability 9 and does not run on SM120. Hopper was not tested.

Speed Tests and Profiling

Mask construction plus one scorer call on one RTX PRO 6000 (SM120). Inputs: 64 heads, page size 64, capacity 262,144 positions, scorer ratio=1, dyadic BF16 queries and weights, random FP4 key bytes with E8M0 scales near 1, one request on shuffled pages. Row r sees context - r positions and gets min(2048, ceil(context / 8)) distinct random candidate blocks. Each variant is captured as 20 calls in one CUDA graph and timed over 10 replays after a warm-up replay; the table reports time per call. Visible scores were equal in every case.

"Mask + masked scan" is candidate_block_mask plus masked fp4_index_logits_paged; "Candidates" is fp4_index_logits_candidates. The mask was built once per source layer and shared by its consumers, so the per-consumer saving is smaller than this table shows.

rows	context	mask + masked scan (us)	candidates (us)	speedup
1	1024	13.9	10.9	1.28x
1	4096	16.0	11.5	1.39x
1	16384	21.1	15.7	1.34x
1	32768	24.1	15.8	1.53x
1	100000	46.1	15.9	2.90x
1	200000	78.3	15.8	4.96x
4	1024	18.1	11.4	1.59x
4	4096	21.8	15.8	1.38x
4	16384	34.8	39.2	0.89x
4	32768	51.3	39.0	1.32x
4	100000	127.6	39.1	3.26x
4	200000	230.3	40.2	5.73x

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Developed with AI assistance.

---

CI States

Latest PR Test (Base): ❌ Run #37967581067
Latest PR Test (Extra): ❌ Run #37967580523
Latest PR Test (AMD ROCm 10): ❌ Run #37967581058

[YJR722]

Hi @ormandj, thanks for the work! Just wanted to flag that candidate-only paged FP4 decode scoring is already part of my ongoing PR #40860. Both implementations reuse _fp4_index_logits_tile with candidate block addressing, so there is substantial overlap.
My implementation also keeps the scores compact through Top-K and maps the selected indices back during finalization. Could you take a look at #40860 to help avoid duplicating effort?