[shikicloud]

Motivation

The DeepSeek-V4 C4 indexer scores every query row of a prefill chunk against every compressed (c4) position of the context into one fp32 logits matrix, [query_rows, align256(max_compressed_seq_len)], allocated inside DeepGEMM (fp8_paged_mqa_logits on the paged path, fp8_mqa_logits on the nonpaged path) once per layer. Nothing sized by --mem-fraction-static accounts for it, and it grows linearly with context:

shape	logits
4096 rows x 372K ctx (issue repro, H100 TP4)	1.42 GiB
4096 rows x 800K ctx (last chunk of the same request)	3.05 GiB
8192 rows x 1M ctx (default chunked-prefill, nonpaged path)	~8 GiB

So a long prefill deterministically OOMs mid-request on every TP rank once the KV pool has consumed the headroom. The DSA (V3.2) indexer already solves the same problem in _get_topk_ragged by chunking the query rows under a free-memory budget; the DSV4 indexer had no equivalent on either path, while SM120 already chunked the paged path (fixed 4096 rows, DeepGEMM metadata kernel cap) and ROCm FP4 chunked into a pooled buffer.

Modifications

• dsa/utils.py: DSA budget logic (_get_mqa_logits_budget_bytes, _should_chunk_mqa_logits) extracted into shared helpers; dsa_indexer.py delegates, behavior unchanged. Knob stays SGLANG_DSA_MQA_LOGITS_FREE_MEM_FRACTION; no new env var.
• dsv4/metadata.py: PagedIndexerMetadata.__post_init__ decides once per forward: one mem_get_info, rows_per_chunk = min(SM120 cap, rows the budget allows), then a DeepGEMM schedule and a top-k v2 plan per chunk. The SM120 list mechanism becomes the generic one. No chunking under prefill CUDA graph, stream capture, BCG/piecewise, or below 8M elements.
• dsv4/indexer.py: the SM120 list loop steps by rows_per_chunk and uses per-chunk top-k plans (slices previously reused the full-batch plan on CUDA); the nonpaged path gathers K once and chunks fp8_mqa_logits by rows.

Accuracy Tests

• CPU unit tests: 21 cases pass (test_dsv4_indexer_chunk.py, test_dsv4_nonpaged_indexer.py, existing test_dsa_mqa_logits_chunking.py).
• GB200 kernel check at the issue shape (4096 x 93000 c4, cluster-path rows included): chunked vs unchunked fp8_paged_mqa_logits / fp8_mqa_logits logits bit-identical; per-chunk top-k v2 plans select the same pages as the full plan.
• E2E, DeepSeek-V4-Flash-0731 TP4 GB200: forced chunking (fraction 0.0005) vs default, prompt logprob deviation inside the run-to-run noise floor of two default launches (mean |dlogprob| ~0.08, top-1 agreement ~0.9), on both the paged (cps 4096) and nonpaged (cps 16384) paths; chunk logs 24/16 vs 0.
• Issue shape: 692K-token prompt, --mem-fraction-static 0.95 (3.7 GB free), prefill completes in 49 s with no OOM (late steps chunk to 411 rows x 173K cols).

Speed Tests and Profiling

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.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ✅ Run #35405295077
Latest PR Test (Extra): ✅ Run #35405873102
Latest PR Test (AMD ROCm 10): ❌ Run #35405295029

[Fridge003]

Extra tests passed https://github.com/sgl-project/sglang/actions/runs/35192680570