[DSV4.1] Enable the two-level candidate indexer on DeepGEMM's paged sparse MQA logits - #38944
Conversation
…parse MQA logits - SGLANG_DSV41_DEEP_GEMM_CANDIDATE_INDEXER (default off) routes the paged fp4 decode path through DeepGEMM's sparse indexer: the candidate-source layer publishes an ascending logical block table plus DeepGEMM's schedule, the index-source layers after it score the published blocks only with fp8_fp4_paged_sparse_mqa_logits and select inside that row. - dsv4/candidate_indexer.py (interface, chosen once at backend init), candidate_torch.py (the model code's mask algorithm, default), candidate_deep_gemm.py (the DeepGEMM implementation; level one still torch). - DSV4Metadata.candidate_metadata: one slot for what the source published, in the implementation's own type; replaces the backend-level candidate_masks. - Under the flag the SM100 low-ratio index-K pool pages at 128 slots (512-byte page stride); the JIT paged-logits metadata builder takes 64 or 128. - topk_bf16_small: JIT bf16 top-k transform with a fused page-table transform, used to select inside the sparse row. - Tests: the bf16 top-k against torch.topk; level one against the model code's select_candidate_blocks; the sparse chain against DeepGEMM's dense bf16 logits. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Replace the torch pieces of the source layer's block selection with two JIT kernels, so publish_decode goes from 25 launches / 95 us to 11 / 58 us at bs=8, 128K on B200: * amax8_varlen (deepseek_v4/amax_copy.cuh): the block maximum of every 8 consecutive fp32 scores for the first ceil(seq_len / 8) blocks of a row, the newest block written as +inf so it is always selected, nothing written past the count; rows with at most `topk` blocks may be skipped. 32-byte vectors on Blackwell, PDL. * sort_candidate_blocks (deepseek_v4/sort_idx.cuh): the selected block ids (any order, -1 padded) become, in place, the ascending INT32_MAX-padded table DeepGEMM's sparse schedule reads, plus the same blocks as pool slots / 8 through the row's page table. Counting sort over a 16 KiB bitmap: single-bit words emitted by their owner, denser words drained from a block-wide queue one word per warp step. Rows with at most k blocks get the identity table. transform_candidate_blocks is the page transform alone, for a block top-k that already emits ascending ids. amax_topk_blocks now returns the top-k's unsorted output and publish_decode sorts it in place; physical_blocks() is gone. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…, plan skip * `indexer.fp32_jit_paged_topk`: the three "plain top-k of the dense logits" sites (backend decode, both implementations' publish) chose between the v2 kernel and the v1 fallback on `use_topk_v2 and raw_indices is None`; with v2 emitting raw indices the choice is `use_topk_v2` alone, in one helper. The torch consumer takes its raw positions from the same call. * `candidate_row_lens` (Triton, PDL): per row the block count and the sparse row length in one launch; publish_decode computes them once and stores the lengths on `SparseBlockTable.valid_lens`, so the consumers stop recomputing them (10 elementwise launches per layer). * `TopKKernel::plan` skips its launch when no route of `transform_paged` reads the plan: only the persistent-cluster route does, and it is taken for batches above the pool size only. publish_decode at bs=8, 128K on B200: 11 launches / 58 us -> 8 / 52 us; select_decode 13 / 31 us -> 3 / 16 us. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
publish_decode forks a second stream after the dense logits: the source layer's own top-k stays on the main stream while the block-max, block top-k, block-table sort and DeepGEMM schedule run on the side stream, so they overlap the main stream's top-k and the following layers. `logits` is recorded on the side stream so the allocator does not recycle it under those reads; the first consumer of the forward joins the side stream before its sparse-logits kernel reads the published table. `SGLANG_DSV41_DEEP_GEMM_CANDIDATE_OVERLAP=0` turns the side stream off. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
12c2190 to
5bd0687
Compare
…e schedule DeepGEMM's sparse metadata pairs two consecutive query rows of one request on one KV pass (each row keeps its own block list and output layout; paired rows must share their page-table row). Under speculative verify every draft token is a row of the same request, so pass the per-row request ids from the dispatcher (IndexerInputs.request_ids) instead of numbering rows: verify rows pair, the sparse-logits kernel reads K once per pair. Decode has one row per request and passes none; that case takes a cached int32 arange, no launch. Test: verify-shaped rows (6 per request, lengths L..L+5, shared page-table row) give bitwise-identical sparse logits paired and unpaired, and equal the dense bf16 logits at the published positions. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
|
End-to-end validation of this branch (head 6862d39) on 4x GB300 (TP4/EP4), official DeepGEMM 2.8.0 (66081d4, built from source) and the open gsm8k (
No errors or truncations except two greedy runaways in one flag-on pass (a request looping until the context limit); the same runaway reproduces with the flag off and depends only on batch composition, and the affected prompts answer correctly when sent alone. Pairwise, any two passes agree on ~68% of outputs regardless of configuration, i.e. the differences between configurations are the same as between two passes of one configuration. Long context (the candidate path is only active above 16384 tokens): needle retrieval at 12K / 24K / 48K / 100K / 200K, three needle positions each, greedy 64 tokens: 15/15 retrieved with both flags, generated text byte-identical between flag on and flag off in every case; first-token logprobs agree to 1e-4. Mixed batches of 8 concurrent requests from 4K to 64K (some below and some above the 16384 threshold): 8/8 retrieved, text identical between flags. Speculative decoding ( Module-level equivalence (synthetic inputs, torch vs DeepGEMM implementation, 61 rows from 16K to 1M tokens, batches 1 and 8 including rows below the threshold): the source layer's top-512 and the published candidate blocks are identical for all 61 rows; the consumers' top-512 selections have Jaccard 0.947-1.0 per row, and 717 of the 718 differing slots are within 2 bf16 ulps of the 512th score (the DeepGEMM consumer scores are bf16, the torch path keeps fp32). Every pick lies inside its own candidate set. Unit tests ( CUDA-graph capture of the candidate variants succeeds with the flag on; no scheduler or CUDA errors in any server log across the runs. |
…iton (port of sgl-project#39086) V4.1 prefill on Hopper currently reaches the `_low_ratio_index_topk_torch` fallback (per-request python loop, per-request dequant, [rows, heads, lc] bf16 scores). Add an SM90 Triton FP8 path and route ragged extend batches to it: - `unpack_fp4_index_keys_to_fp8`: decode block-scaled E2M1 index-K directly to E4M3. - `quantize_bf16_index_queries_fp8`: cast each bf16 query head to E4M3. - `fp8_index_logits_prefill`: E4M3 dot with FP32 accumulation, fused relu / head weighting / head reduction, fp32 output padded to 4 for the ragged top-k. - `_low_ratio_index_topk_sm90_extend`: per-request chunked scoring that reuses each converted K row across the chunk's queries, then feeds the existing ragged top-k v2 kernel. Adapted onto the post-sgl-project#38944 candidate-indexer API (`published_masks` / `CandidateMasks`) and gated by the same `SGLANG_DSV41_TORCH_PREFILL_INDEXER` kill-switch as the dense path. Decode path unchanged. Reported cold prefill on 4xH200 (2K..256K): +14.5%..+71.1%; gsm8k 0.900.
1 participant
DeepSeek-V4.1's low-ratio indexer is two-level: the candidate-source layer (layer 20) scores every visible compressed position and publishes, per query row, the
candidate_topk_blocksbestcandidate_block_size-position blocks; the index-source layers after it (24/28/32/36) select their top-k only inside those blocks. Until now the consumers still computed dense logits over the whole context and masked everything outside the candidates. This PR routes the consumers through DeepGEMM's newfp8_fp4_paged_sparse_mqa_logits, which scores only the published blocks.What this enables
SGLANG_DSV41_DEEP_GEMM_CANDIDATE_INDEXER=1(default off) switches the paged fp4 decode path (SM100, decode and target-verify) to the DeepGEMM sparse indexer: layer 20 runs its own top-k and publishes an ascending logical block table plus DeepGEMM's schedule for it; layers 24/28/32/36 score the published blocks only and select inside that row. Requires a DeepGEMM that providesget_paged_sparse_mqa_logits_metadata/fp8_fp4_paged_sparse_mqa_logits(upstream66081d4or later); the user is responsible for that when turning the flag on.topk_bf16_small, a JIT bf16 top-k transform for rows of at most 16384 scores with a fused page-table transform, used to select inside the sparse row.Structure
dsv4/candidate_indexer.py: the interface (CandidateMetadata,IndexerInputs,make_candidate_indexer, chosen once at backend init).dsv4/candidate_torch.py: the model code's algorithm (bool masks over positions), the default.two_level_decode_logitsmoves here from the backend, keeping the Tritoncandidate_block_logitsfast path.dsv4/candidate_deep_gemm.py: the DeepGEMM implementation (SparseBlockTable: block table + schedule + physical blocks). The level-one block selection is still a torch composition (amax_topk_blocks, marked TODO).DSV4Metadata.candidate_metadata: one slot for whatever the source published, in the implementation's own type; written by layer 20, read by the layers after it, never copied from the host. The former backend-levelcandidate_masksstate moves there._low_ratio_index_topk_decodebecomes: consumer ->select_decode; source ->publish_decode; everything else the plain top-k as before. The short-context CUDA-graph variants keep selecting every position and never touch the candidate path.Unchanged: prefill and the Hopper decode path keep the existing mask-based selection (same results), only their state now lives on the forward metadata.
Tests
test/registered/kernels/ops/attention/test_topk_bf16.py: the bf16 top-k againsttorch.topkacross row lengths, ties, page tables and padded outputs.test/registered/attention/unittests/dsv4/test_dsv41_sparse_indexer.py: level one against the model code'sselect_candidate_blocks; with a DeepGEMM that has the sparse kernel, the sparse logits against DeepGEMM's dense bf16 logits at the published positions (bitwise) and the consumer's selection against a torch reference.Follow-ups
Level-one kernel (block amax + top-k) instead of the torch composition; prefill through the same implementations; end-to-end validation once the DeepGEMM dependency carries the sparse kernels.
🤖 Generated with Claude Code
CI States
Latest PR Test (Base): ❌ Run #34632077104
Latest PR Test (Extra): ❌ Run #34632076753
Latest PR Test (AMD ROCm 10): ❌ Run #34632077042