dsv4.1: candidate indexer library

[hnyls2002]

Summary

• Add the two-level candidate indexer for DeepSeek-V4.1 low-ratio decode: one layer scores every reachable compressed position and publishes its candidate_topk_blocks best blocks of 8, the layers after it score and select only inside those blocks.
• candidate_indexer.py: IndexerInputs as the per-layer operand bundle, make_candidate_indexer, and the mask helpers (CandidateMasks, select_candidate_blocks, mask_topk_scores) the Hopper decode and prefill paths select through.
• candidate_indexer_deep_gemm.py: DeepGemmCandidateIndexer (publish_decode / select_decode), the Blackwell implementation on DeepGEMM's paged sparse MQA logits. Publishing runs the block selection (candidate_row_lens, amax8_varlen, sort_candidate_blocks, DeepGEMM schedule) on a side stream behind the layer's own top-k; the published SparseBlockTable carries a ready event and record_streams its buffers for the main-stream consumers.
• indexer.py: deep_gemm_fp4_paged_mqa_logits and topk_transform_paged_from_metadata join the existing wrapper block; existing definitions unchanged.
• configurer.py: DEEPGEMM_PAGED_SPARSE_MQA_LOGITS probes the installed DeepGEMM for get_paged_sparse_mqa_logits_metadata / fp8_fp4_paged_sparse_mqa_logits; on Blackwell make_candidate_indexer requires them (sgl-deep-gemm 0.2.0, the pinned version, has both) instead of falling back.

Changes to existing kernels

• candidate_blocks.py: the mask-path kernels (candidate_block_logits and its Triton helpers) had no caller and are removed; candidate_row_lens is unchanged.
• indexer_postprocess.py (filter_topk_pages) had no caller and is removed.

Verification

• amax_topk_blocks against a torch block-max reference, and paired verify rows against unpaired rows through DeepGEMM's sparse logits and topk_transform_sparse, both pass on sgl-deep-gemm 0.2.0.
• Every new module imports on a Blackwell device; make_candidate_indexer returns None for models without a candidate source and raises a clear error when the DeepGEMM build lacks the sparse APIs.

---

CI States

Latest PR Test (Base): 🚫 Run #35155455731
Latest PR Test (Extra): ❌ Run #35155455538
Latest PR Test (AMD ROCm 10): 🚫 Run #35155455880

[BBuf]

The abstraction split (CandidateIndexer protocol + DeepGEMM/torch impls, IndexerInputs as the operand bundle) reads well, and valid_lens having both a torch reference and a kernel path is the right shape. Three things.

1. record_stream is applied in one direction only

publish_decode correctly protects the main-allocated tensor that the side stream reads:

self.alt_stream.wait_stream(torch.cuda.current_stream())
# The block-selection chain reads logits after the main stream moves on.
logits.record_stream(self.alt_stream)

but the tensors allocated inside with torch.cuda.stream(self.alt_stream) — nblocks, row_valid_lens, blocks, phys_blocks, schedule — are then consumed on the main stream in select_decode (sparse_logits reads schedule/blocks, topk_transform_sparse reads phys_blocks/valid_lens) with no matching record_stream(torch.cuda.current_stream()).

current_stream().wait_stream(self.alt_stream) orders the production, but it does not address the allocator: a block allocated on alt_stream is returned to alt_stream's pool on free, and the caching allocator may hand it to the next alt_stream allocation without waiting for main-stream kernels that are still reading it. In the current steady-state loop the next publish_decode opens with alt_stream.wait_stream(current_stream()), which probably closes the window — but that is an accident of the call order, not an invariant, and the very next consumer added to the protocol breaks it. Since the reverse direction already does the right thing, could the SparseBlockTable fields get the same treatment (or simply be allocated on the main stream)?

2. need_wait assumes a strict publish→select pairing

self.need_wait is instance state on an object that outlives the forward. If a publish_decode is ever not followed by a select_decode in the same step — a batch where no consumer layer runs, an early return, an exception between the two — the flag stays set and the next step's select_decode consumes it, after which nothing ever waits on the alt stream for that step's work. An event captured on the table itself (SparseBlockTable.ready_event) rather than a flag on the indexer would make this independent of call sequencing. Is publish/select guaranteed 1:1 per forward today?

3. _row_ids can allocate during capture

def _row_ids(rows, device):
    buf = _ROW_IDS.get(device)
    if buf is None or buf.numel() < rows:
        size = max(8192, -(-rows // 8192) * 8192)
        buf = _ROW_IDS[device] = torch.arange(size, ...)

The comment says "cached, no launch", and warmup() exists to pre-build the DeepGEMM workspace. But the cache only grows on demand, so if the captured batch ever has more rows than whatever warmup was called with (and more than 8192), the torch.arange runs inside the graph capture. Is warmup guaranteed to be called with the maximum row count — verify bs * draft_len, not just bs? A cheap assert in _row_ids (assert buf is not None and buf.numel() >= rows) after warmup would turn a subtle capture bug into a loud one.

Tests

No test references candidate_deep_gemm.py, candidate_indexer.py or candidate_torch.py. valid_lens() is explicitly documented as the "Torch reference of candidate_row_lens" and TorchCandidateIndexer is a reference implementation of DeepGemmCandidateIndexer — two free oracles sitting in the diff. A valid_lens vs candidate_row_lens equality test needs no GPU beyond a launch, and torch-vs-deepgemm publish_decode/select_decode agreement would cover the rest.

(Review by Claude Opus 5, run by @BBuf.)

[hnyls2002]

/rerun-test test/registered/kernels/ops/attention/test_dsa_indexer.py test/registered/kernels/ops/attention/test_sm120_paged_mqa_logits.py test/registered/unit/layers/test_dsv4_nonpaged_indexer.py test/registered/kernels/ops/attention/test_dsv4_indexer_quant.py

[github-actions]

Results for /rerun-test test/registered/kernels/ops/attention/test_dsa_indexer.py test/registered/kernels/ops/attention/test_sm120_paged_mqa_logits.py test/registered/unit/layers/test_dsv4_nonpaged_indexer.py test/registered/kernels/ops/attention/test_dsv4_indexer_quant.py:

🚀 1-gpu-h100 (2 tests): ✅ View workflow run

cd test/ && python3 registered/kernels/ops/attention/test_dsa_indexer.py
cd test/ && python3 registered/kernels/ops/attention/test_dsv4_indexer_quant.py

🚀 1-gpu-5090 (1 test): ✅ View workflow run

cd test/ && python3 registered/kernels/ops/attention/test_sm120_paged_mqa_logits.py

🚀 ubuntu-latest (1 test): ✅ View workflow run

cd test/ && python3 registered/unit/layers/test_dsv4_nonpaged_indexer.py

[hnyls2002]

/tag-and-rerun-ci