[DeepSeek-V4] Support Batched and Round-Robin CP Prefill in the Non-Paged Indexer

[YMbmzy]

Motivation

PR #29619 introduced the non-paged DeepSeek-V4 C4 indexer for long-context prefill, and PR #30140 enabled it by default when the local query length is at least 8,192 tokens per rank. The benchmarks in those PRs showed that, once the fixed KV-gather cost is amortized at this scale, gathering the indexer KV cache into contiguous storage and using the non-paged DeepGEMM kernel can improve prefill throughput.

However, the existing fast path is restricted to CP=1 and batch_size=1, which excludes two important serving scenarios:

1. Prefill workers in prefill/decode-disaggregated deployments commonly use context parallelism for long-context requests. CP distributes the prefill workload across ranks and limits per-rank activation memory, but the existing restriction forces these requests back to the paged indexer even when the per-rank query length reaches the profitable non-paged region.
2. With chunked prefill, the scheduler can co-schedule multiple requests as long as their combined prefill tokens fit within --chunked-prefill-size. Such batch_size>1 batches should also benefit from the non-paged path when their aggregate local query length reaches the threshold.

This PR extends the non-paged indexer to batched prefill and round-robin context-parallel prefill while preserving the existing query-length threshold and fail-closed fallbacks for unsupported configurations.

Modifications

• Extend the non-paged indexer from single-request prefill to batched prefill.
• Support CP prefill only with the interleave (round-robin) strategy. Other CP strategies, such as zigzag, explicitly fall back to the paged indexer; support for additional CP layouts is left to follow-up work.
• Build rank-local gather plans that concatenate the selected requests' C4 KV sequences and provide request-aware ragged [ks, ke) ranges to DeepGEMM.
• Trim CP/DP padding before DeepGEMM and top-k, keep dispatch consistent across CP ranks, and rebuild the top-k v2 plan when the logical query shape changes.
• Add unit tests for the new batched and round-robin CP paths and their fallback boundaries.

Accuracy Tests

Test design

• Hardware: 4x HZZ2
• Workload: GSM8K, repeated three times per configuration
• Paged baseline: SGLANG_OPT_DSV4_NONPAGED_INDEXER=0
• Optimized path: SGLANG_OPT_DSV4_NONPAGED_INDEXER=1

The two configurations used the same model, server arguments, dataset, and evaluation parameters. The environment variable above was the only difference between them.

Reproduction commands

Run the server once for each configuration, changing only SGLANG_OPT_DSV4_NONPAGED_INDEXER:

# Use 0 for the paged baseline and 1 for the optimized non-paged run.
export SGLANG_OPT_DSV4_NONPAGED_INDEXER=0

python3 -m sglang.launch_server \
  --model-path <model-path> \
  --tp-size 4 \
  --ep-size 4 \
  --moe-runner-backend flashinfer_mxfp4 \
  --enable-prefill-cp \
  --cp-strategy interleave \
  --max-prefill-tokens 40960 \
  --chunked-prefill-size 40960 \
  --disable-radix-cache

Run the same evaluation against each server:

python3 -m sglang.test.run_eval \
  --base-url http://127.0.0.1:30000 \
  --model <served-model-name> \
  --eval-name gsm8k \
  --api completion \
  --gsm8k-data-path <gsm8k-data-path> \
  --num-threads 128 \
  --repeat 3 \
  --temperature 0 \
  --top-p 1

Runtime (NONPAGED=1)

A runtime log screenshot confirms that the local sequence length exceeded 8,192 tokens per rank, verifying that the non-paged indexer threshold was met and that the optimized code path was exercised during the accuracy evaluation.

Results

Configuration	Run 1	Run 2	Run 3	Mean
Paged baseline (NONPAGED=0)	0.952	0.950	0.956	0.9527
Non-paged (NONPAGED=1)	0.951	0.954	0.956	0.9537
Delta (non-paged - baseline)	-0.001	+0.004	0.000	+0.0010

The paged and non-paged configurations achieved comparable GSM8K scores. The optimized run's mean score was +0.0010 (+0.10 percentage points) above the baseline, which is within the observed run-to-run variation. No accuracy regression was observed.

Checklist

• Format the code with pre-commit.
• Add or update unit tests.
• Update documentation where needed.
• Provide relevant accuracy and performance results.
• Follow the SGLang code-style guidance.

---

CI States

Latest PR Test (Base): Not run yet
Latest PR Test (Extra): ⚠️ Not enabled -- add run-ci-extra label to opt in.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[Qiaolin-Yu]

/tag-and-rerun-ci

[YMbmzy]

/rerun-failed-ci