[DeepSeek-V4.1] Bound the dense prefill indexer transient: row-chunk logits, tail-only candidate masks

[kpham-sgl]

Motivation

Serving DeepSeek-V4.1-Flash on 4x GB300 (TP4 + EP4, --chunked-prefill-size 16384, --enable-encoder-swa-bounded-replay --enable-decoder-swa-bounded-replay, --mem-fraction-static 0.80) under an agentic-coding workload (DeepSWE, prompts up to 705K tokens) crashed with CUDA OOM on all ranks in _publish_or_consume_candidates during a 16384-token chunk of a 487,936-token request. The eager dense fp4 prefill indexer (_low_ratio_index_topk_dense, DeepGEMM fp8_fp4_mqa_logits) scores a whole prefill chunk against the full context in one fp32 tensor [T, lc] (16384 x 487936 x 4 B = 29.8 GiB per rank), and the candidate-source layer then builds bool masks [T, lc] in 550-row pieces and torch.cat-copies them (7.45 GiB each). Peak transient is ~6 B x chunk tokens x context tokens per rank, so it grows without bound with context; the prefill CUDA graph path and the torch fallback are already row-chunked, but every context above 16K runs eager.

Reproduction on the unpatched branch: a single 600K-token prompt kills the server within 18 s; even a cold 300K prompt reaches a 49.7 GB transient against 50 GB free after graph capture.

Modifications

In _low_ratio_index_topk_dense:

• Score each request in row chunks bounded by _DENSE_INDEXER_LOGITS_BUDGET_BYTES (2 GiB of fp32 logits), running the ragged top-k per chunk into the shared selected buffer. Short contexts still take a single chunk, so the common path is unchanged.
• Build the candidate masks straight into one preallocated [rows, lc] tensor instead of a list plus a concatenation copy.
• Under decoder SWA bounded replay, publish only each request's tail rows: enter_late_layer_tail slices every mask down to those rows anyway, so the full-row masks were computed and copied for nothing.

The separate _publish_or_consume_candidates helper folds into the loop.

Accuracy / equivalence

Exact in-forward check (debug build, not in this PR): the original and the new implementation ran back to back on identical inputs and their page_indices / raw_indices / published masks were compared bitwise, over cold single prompts (8K, 60K, 200K), a batch of four 7K prompts, and a prefix-hit prefill.

	layer-forwards	new == original	original == original (run twice)	masks equal
this head (0ff9c986c)	208	194	194	208
f3c3e7c7a	176	162	162	176

The residual is top-k tie ordering on ratio-2 layers with heavily repeated prompt text, at the same rate and magnitude as the original path disagrees with itself. DeepGEMM's fp8_fp4_mqa_logits is bitwise row-count invariant (16384 rows in one call vs chunks of 128..4096, all equal) and deterministic across runs, so the chunking itself introduces no numerical change.

Benchmarks

4x GB300, TP4 + EP4, both replays on, chunk 16384, mem-fraction 0.80 (50 GB free after graph capture), cold prompts, greedy 4 new tokens:

Cold prefill	Unpatched	This PR
300K tokens	11.5 s, transient 49.7 GB	10.4 s, transient 10.7 GB
600K tokens	OOM (server exit 137)	18.6 s, transient 9.6 GB
1M tokens	not attempted	35.6 s, transient 9.6 GB

Without decoder replay the published mask is still 1 B per row per context token (16 GiB for a 16K chunk at 1M context); a block-level mask would remove that but changes the consumer format, so it is left for a follow-up.

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ❌ Run #34830063530
Latest PR Test (Extra): ❌ Run #34830063321
Latest PR Test (AMD ROCm 10): ❌ Run #34830063499

[sethforprivacy]

Independent GB10 serving evidence for the logic in cb8dd033ab54af8904733199dbae94270c3395ce.

We backported the row-chunked logits and tail-only candidate-mask changes to SGLang e087e662ba1ac4ef7747537e2a9141085efd4561, retaining that revision's self.candidate_masks interface and _mask_topk_scores helper. This is a backport validation, not a test of the current PR branch as a whole.

Conditions: four DGX Sparks / GB10 (128 GB unified memory per host), TP4/EP4 over a switchless RoCE cycle with SparkRing's patched NCCL 2.30.7; model revision dba1be0a40aa45a94ad051997016db3960a90277; Mia adapter e59e6eb67479aa68f6fa700c600dc90a0729b5ec, node-local packed NVMe Engram; chunk 4096, context 430080, requested shared KV 1500000 (actual 1499904), max running 8, memory fraction 0.90, decoder SWA bounded replay, DSpark block 5 with natural acceptance and no SPS table. PyTorch 2.13.0+cu130, Triton 3.7.1, FlashInfer 0.6.18; b12x MXFP8, fused MoE finalize disabled, expandable segments disabled.

Same qualification sequence	Before	Backport
Minimum sampled host MemAvailable across all four ranks	13.424 GiB	24.631 GiB
Exact retrieval, 406713 actual prompt tokens, TTFT	176.1 s	155.6 s
Nominal 64K prefill throughput	3312 tok/s	3392 tok/s

Both configurations passed seven-key authentication, vision/tools, chunk/tail-boundary and mixed prefill/decode checks. All 20 first-token outputs and all 20 greedy continuations (up to 48 tokens) matched our fixed reference, including four longer prompts. Exact retrieval passed at 130258, 260119 and 406713 prompt tokens. No benchmark request errors were observed. We retained the backport for the additional measured headroom.

Limits: the memory measurement is a five-second host MemAvailable sample, not an isolated CUDA peak-allocation measurement. The baseline was already serving and the candidate was rebooted before qualification; page-cache/boot state remains a confound, and this is one qualification pair rather than repeated independent boot trials. The long retrieval timing is one observation per arm; the small short-prefill change is not a demonstrated speedup. The nominal 64K workload uses fresh random prompts (about 59K actual tokens), one output token, and throughput including request overhead. We have not yet qualified a larger context, tested CP, or independently reproduced your bitwise internal-index comparison.

Our original deployment/replay reference is in SparkRing PR267; this new patch result is separate from the six-hour soak recorded there. Sharing the bounded result because the reduced transient also appears useful on GB10's much tighter unified-memory budget.