[DSV4.1] Optimize prefill GEMMs and compensated mHC for concurrent serving

[BBuf]

Motivation

DeepSeek-V4.1 Flash at concurrency 64 spends avoidable prefill time in untuned dense MXFP8 GEMMs, repeated mHC projections, and output-layout copies. DSpark warmup covers small verify shapes; the large-prefill tactic can also disappear when a subsequent draft autotune phase loads its cache after a server restart.

On the current dsv4.1 port, the code changes improve the unchanged attention TP4 / MoE EP4 serving configuration from 7348.46 to 7955.07 output tokens/s (+8.25%). A separate deployment using existing DPA4 attention reduces TTFT below the recorded vLLM TP4 reference in two independent starts. DPA4 changes attention parallelism; its full gain is not a kernel-only comparison. Original attention TP4 still has higher P99 TTFT.

This remains a draft. Current-source GPU tests and serving measurements are complete; the finite accuracy comparisons below do not establish exact output equivalence or a formal non-inferiority margin. No merge is requested by this update.

Changes

• Warm up the model's resident MXFP8 linear layouts at the per-rank prefill ceiling without constructing request/KV metadata or running draft/LM-head/vision forwards. Include the actual checkpoint's block-FP8-as-MXFP8 consumer. Permit tuned tactics for opted-in DSV4.1 large row counts while retaining small-row and explicit deterministic/batch-invariant behavior.
• Preserve loaded FlashInfer target/draft tactics across autotune phases through the public load/save API. Keep metadata validation, cross-rank cache consensus, timing synchronization, skip policy, and cache-disabled behavior.
• Keep the opt-in SGLANG_DSV41_COMPENSATED_MHC=1 large-batch path. Small supported batches use the high/residual DeepGEMM projections plus fused residual reduction. For DSV4.1 prefill M=4096..65536, three BF16 weight components share one activation read with separate FP32 accumulation and fixed 16-way K partitioning, followed by the existing reduce/Sinkhorn kernel. Unsupported configurations and batch-invariant execution retain fallback.
• Retain original FP32 weights and derived-weight refresh rules. The supported 40-layer target holds about 525 MiB/rank of combined old/new derived mHC buffers; draft buffers are additional when enabled. Online weight updates reject before writes while this optional static-weight mode is enabled.
• Add a large-prefill Q RoPE/store path, direct token-major BF16 WO-A output, and one-row-per-CTA mHC combine/RMSNorm. Preserve intermediate BF16 rounding, padding, strided output, and 64-bit address calculations. Runtime M avoids specializing every prefill token count.
• Allow 65536 query tokens in the compressor plan while retaining the request-count bound that excludes the packed invalid sentinel. Test 65536/65537 boundaries and sentinel rejection.
• Extend the existing bounded candidate graph interface to short-context DSpark target verification. Bound all speculative positions without D2H synchronization, preserve causal top-k, and fall back for unknown bounds. Do not enable the plain-decode candidate-all shortcut during verification.

The candidate graph and compensated projection changes were already part of this standalone PR. This update adds the prefill work and validates the combined runtime on GB300. It does not duplicate #38944/#39305 sparse-indexer or consumer TopK integration.

Current-source validation

Runtime commit: ad3122c811507679b014cde703b9ce91bc603b24. Its 15 changed/new files matched the remote tested files by SHA256. The subsequent commits format one test without changing its AST and move the CUDA cache test into the kernel CI suite; runtime code is identical. Current head: 933cd91263ef45b91e8cb218f639d95ccb134a36.

• 82 tests + 144 subtests passed on GB300, including existing candidate graph, weight-update guard and cross-rank autotune tests; FP64 mHC reference checks, non-aligned/64K shapes, CUDA Graph replay with changed input, Q padding, BF16 WO-A equality, planner overflow, and cache persistence tests.
• Independent all-rank prefill/decode profiles and a 64-request audit: each request returned exactly 1024 output token IDs with matching completion usage and no request failures.
• Ruff 0.15.1 configured lint/format, isort, clang-format 20.1.7, and test registration against the dsv4.1 base pass. The relocated CUDA cache tests also pass again (2 tests). GitHub's prior-head lint failure was formatting in two unrelated tool-parser files in the CI merge tree; it was not a GPU test failure. The draft PR gate skips GitHub GPU jobs. Local GPU validation does not imply all GitHub checks are green.

Real-acceptance GSM8K uses the same 1314 held-out five-shot prompts; both simulation variables were confirmed absent in the server environment. Initial differences triggered a fixed additional control set; every run is reported:

Runtime / topology	Correct out of 1314, runs 1 / 2	Truncated, runs 1 / 2	Errors / empty
Updated code, TP4 / EP4	1278 / 1278	1 / 0	0 / 0 in each run
Updated code, DPA4 / EP4	1276 / 1277	0 / 0	0 / 0 in each run
Prior head + planner boundary fix, DPA4 / EP4	1278 / 1276	0 / 1	0 / 0 in each run

DPA4 scores overlap the prior implementation's repeated-run range; the mean difference is -0.038 percentage points. Parent repeat 1→2 changes 7 wrong answers to correct and 9 correct to wrong; final repeat 1→2 changes 6/5. Corresponding parent→final runs change 5 gained/7 lost and 7 gained/6 lost. This small control set does not prove exact output equivalence or formal non-inferiority.

The TP4 truncation is sample119 (also incorrect in the historical reference); the parent DPA4 truncation is sample1003. Both reach the unchanged 4096 output limit with repeated reasoning. No sample was excluded or given a larger scoring budget. The original stricter historical-score/zero-truncation checks remain recorded as unmet; diagnostic repeats are separate evidence, not a threshold relaxation. AIME was not run in this round.

Fixed serving results

Same 4×GB300 node, checkpoint revision dba1be0a40aa45a94ad051997016db3960a90277, frozen 1024 requests, 4096 input / 1024 output tokens, concurrency 64, static DSpark with five draft tokens and simulated acceptance length 5.5. The SGLang v0.5.19 streaming client uses 64 warmups, flushes the cache, then times 1024 requests including fill/drain; all cases complete 1,048,576 output tokens. TTFT is the first nonempty text SSE chunk.

Current-node configuration	Output tokens/s	Mean TTFT / ms	P50 TTFT / ms	P90 TTFT / ms	P99 TTFT / ms
vLLM attention TP4 / MoE TP4 reference	7272.20	1181.48	1092.64	1962.24	3998.89
Prior PR head + 64K boundary fix, SGLang TP4 / EP4	7348.46	3856.02	3873.08	4790.28	4973.39
Updated code, same TP4 / EP4 and original scheduling	7955.07	3311.67	3140.65	4119.88	4132.11
Updated TP4 / EP4 + existing publication/progress controls	7584.25	1133.42	981.01	1961.99	4404.56
Updated DPA4 / EP4 + controls, independent start 1	9358.31	620.78	432.52	1204.41	2754.89
Updated DPA4 / EP4 + controls, independent start 2	9478.19	629.37	469.76	1205.65	2730.88

The first code-only control already includes the necessary 65536-token planner boundary fix; the unmodified prior head rejected that warmup. DPA4 uses existing --enable-dp-attention --dp-size 4 --enable-dp-lm-head, giving attention TP1 per DP rank, with the same global prefill ceiling of 65536 and per-DP ceiling of 16384. All four ranks flushed their cache before measurement. These comparisons establish improvement against this recorded vLLM deployment, not the best possible vLLM configuration. vLLM offloads Engram to CPU while SGLang keeps it on GPU.

Representative separate profiles confirm dispatch: 64K TP4 dense MXFP8 summed kernel time ~148 → 76 ms, mHC projection ~97.61 → 46.53 ms, and prefill kernel count 3866 → 3786 after the mHC update. Final DPA4 rank-0 prefill span is ~582.85 ms; its head64 sparse-attention family does fewer query rows (~57 ms versus TP4 ~224 ms). Profiled kernel sums include instrumentation and overlap and are not additive TTFT speedups.

Environment: driver 595.71.05, PyTorch 2.13.0+cu130, Triton 3.7.1, FlashInfer 0.6.18, TileLang 0.1.12, NCCL 2.30.7; vLLM 0.1.dev20904+g179dd0fa9. Client source commit 0bcd822377da7b5718e674eaf9c870d349424dd1; frozen corpus SHA256 1f8e67cba16f9fa16573286775db9a2353fc4c42006695ba63c3515c7b2401f3.

Serving command for the DPA4 result

From this checkout, set MODEL_PATH to the pinned checkpoint:

export CUDA_VISIBLE_DEVICES=0,1,2,3
export PYTHONPATH="$PWD/python"
export LIBRARY_PATH=/usr/local/cuda/lib64/stubs
export MAX_JOBS=16
export SGLANG_RAGGED_VERIFY_MODE=static
export SGLANG_SIMULATE_ACC_LEN=5.5
export SGLANG_SIMULATE_ACC_METHOD=match-expected
export SGLANG_DSV41_COMPENSATED_MHC=1
export SGLANG_DISABLE_CONSECUTIVE_PREFILL_OVERLAP=1
python -m sglang.launch_server \
  --model-path "$MODEL_PATH" \
  --served-model-name deepseek-ai/DeepSeek-V4.1-Flash \
  --tp 4 --ep-size 4 --dp-size 4 \
  --enable-dp-attention --enable-dp-lm-head \
  --moe-a2a-backend none --moe-runner-backend flashinfer_mxfp4 \
  --trust-remote-code --mem-fraction-static 0.80 \
  --max-total-tokens 33554432 \
  --chunked-prefill-size 65536 --max-prefill-tokens 65536 \
  --prefill-decode-interval 4 \
  --cuda-graph-bs-decode 1 2 4 8 16 32 64 \
  --max-running-requests 128 \
  --speculative-algorithm DSPARK --speculative-dspark-block-size 5 \
  --skip-server-warmup --reasoning-parser deepseek-v41 \
  --random-seed 42 --decode-log-interval 10 \
  --host 127.0.0.1 --port 31001 --nccl-port 31745

Remove the three DP flags to use the original attention TP4 topology. Also unset SGLANG_DISABLE_CONSECUTIVE_PREFILL_OVERLAP and remove --prefill-decode-interval 4 for the original-schedule code-only control. The DPA result is not attention TP4 despite the outer --tp 4 argument.

Unset both simulation variables before restarting for real accuracy evaluation. The score protocol uses identical pinned five-shot GSM8K prompts, test indices 5..1318, temperature 0, top_p 1, seed 0, max output 4096, concurrency 64, recording complete responses and truncations. Dataset SHA256: 3730d312f6e3440559ace48831e51066acaca737f6eabec99bccb9e4b3c39d14.

Raw serving timings, full accuracy responses, all-rank traces, source hashes and replay scripts are archived with the local experiment report. Regenerating random prompts is a new workload.

Remaining limits

• Fixed attention TP4 P99 TTFT remains above the recorded vLLM reference. DPA4 is a separate validated performance configuration.
• These are finite GSM8K/kernel checks, not a proof of exact output equivalence. AIME, PD disaggregation, other model families and online weight updates in compensated mode are not claimed by these measurements.
• Compensated mHC remains opt-in; no global quantization precision reduction or FP8 WO-A conversion is introduced.

---

CI States

Latest PR Test (Base): ❌ Run #34830110815
Latest PR Test (Extra): ❌ Run #34830110709
Latest PR Test (AMD ROCm 10): ❌ Run #34830111337