[DSV4.1] Combine DSpark decode and prefill kernel optimizations - #39370
Merged
BBuf merged 22 commits intoSep 15, 2026
Merged
Conversation
`accept_greedy` takes the target prediction with `torch.argmax` over the verify logits. At DSpark block 5 that tensor is [6, 129280] fp32 -- 3.1 MB, six very wide rows -- and `at::native::reduce_kernel` gives each row a single block, so it needs 23.6 us on GB300 for a reduction that reads 3.1 MB. The accept sits on the critical path between the target verify and the KV commit, so the whole cost is exposed. A flat two-stage split (64 partials per row, then a 64-wide final) does the same reduction in 3.5 us. Ties resolve to the lowest index, matching `ArgMaxOps`' strict `>`, so the result is bitwise identical to `torch.argmax`: 0 mismatches over 300 randomized [6, 129280] cases including exact ties, -Inf and +Inf. The helper only takes over for few-row/wide-vocab fp32 rows and falls back to `torch.argmax` everywhere else; like the tensor it is given, it takes the row stride rather than assuming one. -46 us per verify cycle on DeepSeek-V4.1-Flash, 4xGB300, BS=1.
`vision_topk` is the router a DeepSeek-V4.1 checkpoint with a vision tower takes for *all* of its MoE layers, not just the vision ones. It calls `moe_fused_gate` without the `packed_out` argument the text path passes, so every MoE layer paid a separate `PackTopkIds` launch to rebuild what the gate already had in registers. Pass the buffer and return `StandardTopKOutputPacked`, under the same admission `_fused_gate_emits_packed_ids` uses on the text path: only the flashinfer_mxfp4 runner consumes the packed form, and nothing may rewrite ids or weights after the router -- which rules out the fused-shared-expert slots `_scale_fused_shared_weights` rescales just below. Same ids, one fewer launch per layer. -64 us per verify cycle on DeepSeek-V4.1-Flash, 4xGB300, BS=1.
On the cross-layer mHC path the collapsed, normalized attention input is written BF16 by `hc_combine_norm` and then immediately re-read by a standalone FlashInfer `mxfp8_quantize` inside the first projection. At BS=1 the two launches cost about the same, even though the second one only re-reads what the first just wrote. Add an MXFP8 epilogue to the combine+norm kernel and hand the result to the attention as the `Mxfp8SwizzledInput` it already accepts. The scale factors use FlashInfer's 128x4 swizzle and the same UE8M0 conversion, and the values come from the same BF16 rounding the standalone pair produces, so both outputs are bitwise identical to `hc_combine_norm` followed by `mxfp8_quantize(..., is_sf_swizzled_layout=True)`: 0 mismatches over 200 randomized cases across five magnitude decades, against both the `cuda` and the `cute-dsl` FlashInfer quantizer backends. The epilogue is only taken when the attention's first projection actually consumes a swizzled MXFP8 tuple (`accepts_mxfp8_swizzled_input`) and the existing fused combine+norm fast path applies; otherwise nothing changes. -56 us per verify cycle on DeepSeek-V4.1-Flash, 4xGB300, BS=1.
…05 -> 7.44 us) `sparse_decode_fwd` builds its tile-scheduler metadata itself whenever it is handed none, in a `<<<1, 32>>>` kernel whose entire partition loop runs on thread 0 and stores each 32-byte entry straight to global memory. The loop is over `num_sm_parts` = `num_sms / s_q`, so a BS=1 step on GB300 walks 152 partitions at ~156 ns each: 28.05 us measured in the decode graph, where it is fully exposed on the critical path. It cannot be hoisted out of the graph. `on_after_cuda_graph_warmup` clears the scheduler before capture on purpose, because the schedule depends on the step's `topk_length` and a frozen one would leave the tail of a longer row unattended. So the fix is to make the schedule cheap, not to compute it less often. `decoding_sched_meta` runs the same algorithm over a whole block: * the per-request pass and the final write-out are spread over 256 threads; * the request being consumed changes only `batch_size` times across the walk, so its three values stay in registers instead of being re-read from shared memory once per partition -- that chain of dependent loads with one thread and no ILP to hide them is what the 156 ns actually was; * the walk stops when the last request is consumed. Every partition after that describes the same empty range, and the block fills them in parallel. At BS=1 that is most of them: 152 partitions over 5 or 6 query rows leaves fewer than 50 doing any work. The result is fed to FlashMLA as the cached metadata, so it skips its own kernel. Only computed where FlashMLA would have computed -- when the scheduler holds no buffers yet -- so the calls that already reuse a schedule are untouched. In the decode graph: 80 launches per 20 steps either way, 28.05 us each before and 7.44 us after, and no `get_mla_metadata_kernel` left in the trace. Bit-identical to what it replaces, with a new registered test: the schedule itself over batch 1..64, s_q 1 and 6, top-k 512 and 2048 and five `topk_length` shapes including all-zero and all-one rows; the attention output compared as raw bits (a random fp8 cache decodes to NaN in places); and the extra-cache schedule, which rounds the main length up to a block before adding the extra one. FlashMLA's `DecodingSchedMeta` ends in a `_pad` word it never writes, so the comparison covers the seven defined fields. DeepSeek-V4.1-Flash, 4xGB300, BS=1. Isolated A/B over the env switch, 3 launches x 7 runs per arm: 911.30 -> 922.20 tok/s (+1.20%), every launch pair positive. (That branch also carried a benchmark-only change since dropped, so the absolute figures sit above this PR's; the delta is the same either way.)
Split FP32 projection weights into TF32 high and residual components, then fuse per-slice compensation into Sinkhorn. Preserve small-row and batch-invariant paths and reject online updates while derived graph weights are active. Validated on GB300: 35 coefficient, replay, dispatch, epilogue and update-guard tests; 1314 real-acceptance GSM8K requests, 1277 correct, no request errors and two length-limited outputs.
BBuf
requested review from
AniZpZ,
Edwardf0t1,
FlamingoPg,
Fridge003,
HaiShaw,
JustinTong0323,
OrangeRedeng,
Qiaolin-Yu,
Ying1123,
b8zhong,
ch-wan,
hnyls2002,
ispobock,
merrymercy,
mickqian,
yctseng0211,
yhyang201 and
yuan-luo
as code owners
September 14, 2026 05:20
BBuf
requested review from
DarkSharpness,
HydraQYH,
celve and
hebiao064
as code owners
September 14, 2026 05:20
Collaborator
Author
|
Server launch commands used for the new AL=5.4, random 4K/1K, concurrency 1–64 curve in this PR. These are the recorded commands, with the original local paths retained. Run the two engines sequentially on the same four GPUs, using their separate environments. Versions and paths
SGLangexport MODEL_PATH=/scratch/dsv41/checkpoint
export CUDA_VISIBLE_DEVICES=0,1,2,3
export MAX_JOBS=16
export LIBRARY_PATH=/usr/local/cuda/lib64/stubs${LIBRARY_PATH:+:$LIBRARY_PATH}
export PYTHONPATH=/scratch/dsv41-combined-al54/sglang/python
export SGLANG_RAGGED_VERIFY_MODE=static
export SGLANG_SIMULATE_ACC_LEN=5.4
export SGLANG_SIMULATE_ACC_METHOD=match-expected
python3 -m sglang.launch_server \
--trust-remote-code \
--model-path "$MODEL_PATH" \
--served-model-name deepseek-ai/DeepSeek-V4.1-Flash \
--tp 4 --ep-size 4 \
--mem-fraction-static 0.8 \
--speculative-algorithm DSPARK \
--speculative-dspark-block-size 5 \
--cuda-graph-max-bs-decode 64 \
--reasoning-parser auto \
--tool-call-parser auto \
--host 127.0.0.1 --port 31001This is attention TP4 / MoE EP4, with compensated mHC enabled. No DPA flags or scheduling overrides were added. Both starts resolved to:
The provided command did not fix the server seed. SGLang generated vLLMRun this after stopping the SGLang server, in the separate vLLM environment: export MODEL_PATH=/scratch/dsv41/checkpoint
export CUDA_VISIBLE_DEVICES=0,1,2,3
export MAX_JOBS=16
export LIBRARY_PATH=/usr/local/cuda/lib64/stubs${LIBRARY_PATH:+:$LIBRARY_PATH}
unset PYTHONPATH
export VLLM_ENGINE_READY_TIMEOUT_S=7200
export VLLM_SERVER_DEV_MODE=1
export VLLM_USE_RUST_FRONTEND=1
/scratch/dsv41-combined-al54/vllm-venv/bin/python \
-m vllm.entrypoints.cli.main serve "$MODEL_PATH" \
--served-model-name deepseek-ai/DeepSeek-V4.1-Flash \
--tokenizer-mode deepseek_v41 \
--tensor-parallel-size 4 \
--tool-call-parser deepseek_v41 \
--enable-auto-tool-choice \
--reasoning-parser deepseek_v41 \
--mm-encoder-tp-mode data \
--speculative-config '{"method":"dspark","num_speculative_tokens":5,"draft_sample_method":"probabilistic","rejection_sample_method":"synthetic","synthetic_acceptance_length":5.4,"enable_adaptive_verification":false}' \
--host 127.0.0.1 --port 31002This retains the supplied vLLM command's attention TP4 / MoE TP4 topology ( Curve measurement protocol
At c=64, the two-start median total output throughput is 7487.70 vs 6567.17 tokens/s (+14.02% for SGLang), while median TTFT is 2663.57 vs 667.56 ms. The throughput curve should not be interpreted as TTFT parity. These commands use synthetic acceptance for performance measurement. |
BBuf
requested review from
alphabetc1,
hanming-lu,
huangtingwei9988,
hzh0425,
xiezhq-hermann and
yizhang2077
as code owners
September 14, 2026 07:54
Collaborator
Author
|
GSM8K and AIME accuracy comparison for the restored compensated-mHC head Tested current PR head
All four evaluations completed on both revisions with zero request errors. All samples, including output-limit hits and empty answers, remain in the denominators. Observed differences:
Reproduction conditions:
GPQA Diamond comparison is still running and will be reported separately. These accuracy measurements are separate from the PR's synthetic-acceptance throughput results. |
…v4.1 Integrate NVLink vocabulary gather and ratio-1/2 stream scheduling. Use compact greedy collectives and model-specific metadata plans. Compute only the 16 consumed attention heads for small SM10x batches, fuse inverse RoPE and Q normalization/quantization, and preserve the large-batch compensated mHC path and deterministic fallbacks.
Collaborator
Author
|
Reproducing the 1070.80 output tokens/s BS=1 result after removing small-batch attention This result is tied to The metric is decode/output throughput at simulated speculative acceptance 5.5, excluding prefill. It is not real-acceptance throughput or an accuracy result. The previous concurrency-curve comment uses AL=5.4 and EP4; those are different measurement settings. Tested environment
From an existing SGLang git clone, create an isolated checkout: git fetch https://github.com/BBuf/sglang.git bbuf/dsv41-combined-optimizations
git worktree add --detach ../sglang-pr39370-1070 7322accddbc1d302213342c97d1b516fda209af1
cd ../sglang-pr39370-1070
git rev-parse HEADUse the pinned runtime above and a clean shell without leftover Server commandThe original outer harness supplied #!/usr/bin/env bash
set -euo pipefail
export MODEL_PATH=/scratch/dsv41/checkpoint
export SGLANG_SPEC_TP_SYNC=rng
export SGLANG_RAGGED_VERIFY_MODE=static
export SGLANG_SIMULATE_ACC_LEN=5.5
export SGLANG_SIMULATE_ACC_METHOD=match-expected
CUDA_VISIBLE_DEVICES=0,1,2,3 PYTHONPATH="$PWD/python" MAX_JOBS=16 \
python -m sglang.launch_server \
--model-path "$MODEL_PATH" \
--served-model-name deepseek-ai/DeepSeek-V4.1-Flash \
--tp 4 --ep-size 1 --trust-remote-code \
--moe-a2a-backend none --moe-runner-backend flashinfer_mxfp4 \
--mem-fraction-static 0.80 --max-total-tokens 33554432 \
--chunked-prefill-size 4096 \
--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 30021Exact workload and timing clientOne fixed synthetic 4096-token input, 1024 output tokens, sequential requests (concurrency 1), The following generator reconstructs the exact recorded input IDs. The three ranges are the pinned tokenizer's sorted ordinary vocabulary IDs after excluding declared special tokens (128,051 IDs). The SHA256 assertion verifies the generated sequence; it does not depend on another machine's tokenizer version. Run in a second terminal, in an empty directory for the reproduction files: cat > make_prompt.py <<'PYGEN'
import hashlib
import json
import random
from pathlib import Path
vocab = list(range(3, 128000)) + list(range(128799, 128846)) + list(range(129264, 129271))
assert len(vocab) == 128051
rng = random.Random(42)
ids = [vocab[rng.randrange(len(vocab))] for _ in range(4096)]
digest = hashlib.sha256(json.dumps(ids, separators=(',', ':')).encode()).hexdigest()
assert digest == 'c37b49585df8fb22e125ae467003a00fff94a3f010cc7cddc03a3b463d047614'
Path('prompt.json').write_text(json.dumps({
'input_ids': ids, 'input_ids_sha256': digest,
'max_new_tokens': 1024, 'ignore_eos': True,
}, indent=2) + '\n')
PYGEN
python3 make_prompt.pySave this as benchmark.py — the recorded timing client, included in full"""Synthetic random-token BS=1 timing; DSpark target acceptance is simulated at 5.5."""
import argparse
import collections
import hashlib
import json
import statistics
import time
from pathlib import Path
import requests
def sha(data):
return hashlib.sha256(data).hexdigest()
def save(path,value):
path.write_text(json.dumps(value,ensure_ascii=False,indent=2)+'\n')
class Client:
def __init__(self,url):self.url=url.rstrip('/')
def post(self,route,body=None):
r=requests.post(self.url+route,json=body or {},timeout=(30,1800));r.raise_for_status();return r
def wait(self):
end=time.monotonic()+2400
while time.monotonic()<end:
try:
if requests.get(self.url+'/v1/models',timeout=3).status_code==200:break
except requests.RequestException:pass
time.sleep(3)
else:raise TimeoutError('Server not ready')
self.post('/freeze_gc')
def run(self,prompt,max_tokens=2048):
self.post('/flush_cache?timeout=30')
body={'input_ids':prompt['input_ids'],'sampling_params':{
'temperature':0,'max_new_tokens':max_tokens,'ignore_eos':prompt.get('ignore_eos',False),'stream_interval':1},
'stream':True}
start=time.perf_counter();first=None;first_count=None;last=None;last_time=None
with requests.post(self.url+'/generate',json=body,stream=True,timeout=(30,1800)) as r:
r.raise_for_status()
for line in r.iter_lines():
if not line.startswith(b'data: '):continue
if line[6:]==b'[DONE]':break
chunk=json.loads(line[6:]);now=time.perf_counter()
if 'error' in chunk:raise RuntimeError(chunk)
count=chunk.get('meta_info',{}).get('completion_tokens',0)
if count and first is None:first=now;first_count=count
last=chunk;last_time=now
assert first is not None and last is not None
meta=last['meta_info'];count=meta['completion_tokens']
assert meta['prompt_tokens']==len(prompt['input_ids'])
assert meta.get('cached_tokens',0)==0,meta
if prompt.get('ignore_eos',False):assert count==max_tokens,meta
text=last.get('text','')
words=text.split();ngrams=collections.Counter(tuple(words[i:i+16]) for i in range(max(0,len(words)-15)))
return {'input_tokens':meta['prompt_tokens'],'output_tokens':count,
'accept_length':meta.get('spec_accept_length'), 'verify_steps':meta.get('spec_verify_ct'),
'ttft_s':first-start,'elapsed_s':last_time-start,'first_event_tokens':first_count,
'output_tps':(count-first_count)/(last_time-first) if last_time>first else None,
'finish_reason':meta.get('finish_reason'),'max_repeated_16gram':max(ngrams.values(),default=0),
'request':body,'response':last}
def main():
p=argparse.ArgumentParser()
p.add_argument('mode', choices=['bench'])
p.add_argument('--url', default='http://127.0.0.1:30021')
p.add_argument('--prompt', type=Path, default=Path(__file__).with_name('prompt.json'))
p.add_argument('--out', type=Path, required=True)
p.add_argument('--repeat', type=int, default=6)
p.add_argument('--max-tokens', type=int, default=None,
help='Override the output length saved in the prompt JSON.')
a=p.parse_args()
if a.repeat < 1: p.error('--repeat must be positive')
a.out.mkdir(parents=True, exist_ok=True)
prompt=json.loads(a.prompt.read_text());values=[]
max_tokens=a.max_tokens if a.max_tokens is not None else prompt.get('max_new_tokens',1024)
if max_tokens < 1: p.error('--max-tokens must be positive')
assert sha(json.dumps(prompt['input_ids'],separators=(',',':')).encode())==prompt['input_ids_sha256']
c=Client(a.url);c.wait()
for rep in range(a.repeat+1):
r=c.run(prompt,max_tokens);r.update(repeat=rep,warmup=rep==0)
values.append(r);save(a.out/'measurements.json',values)
print('BENCH',rep,'tps',r['output_tps'],'accept',r['accept_length'],'output',r['output_tokens'],flush=True)
measured=values[1:]
summary={'runs':a.repeat,'output_tps_median':statistics.median(r['output_tps'] for r in measured),
'output_tps_min':min(r['output_tps'] for r in measured),'output_tps_max':max(r['output_tps'] for r in measured),
'accept_length_median':statistics.median(r['accept_length'] for r in measured) if measured[0]['accept_length'] else None,
'input_tokens':len(prompt['input_ids']),'output_tokens':[r['output_tokens'] for r in measured],
'requested_output_tokens':max_tokens,'ignore_eos':prompt.get('ignore_eos',False),
'accept_length_min':min(r['accept_length'] for r in measured) if measured[0]['accept_length'] else None,
'accept_length_max':max(r['accept_length'] for r in measured) if measured[0]['accept_length'] else None,
'prompt_file_sha256':sha(a.prompt.read_bytes()),'acceptance_mode':'Read server launch manifest: simulated target 5.5 for DSpark, not applicable for ordinary decode'}
save(a.out/'summary.json',summary);print('BENCH_DONE',json.dumps(summary),flush=True)
if __name__=='__main__':main()python3 benchmark.py bench --url http://127.0.0.1:30021 --prompt prompt.json --max-tokens 1024 --repeat 6 --out run-aThe client executes one warmup, then six measured requests. It stores every request in python3 benchmark.py bench --url http://127.0.0.1:30021 --prompt prompt.json --max-tokens 1024 --repeat 6 --out run-b
python3 - <<'PYRESULT'
import json
import statistics
from pathlib import Path
values = [r['output_tps'] for name in ('run-a', 'run-b')
for r in json.loads((Path(name) / 'measurements.json').read_text())
if not r['warmup']]
assert len(values) == 12
print('pooled median output tps:', statistics.median(values))
print('range:', min(values), max(values))
PYRESULTFor each request, the calculation is: This excludes both prefill time and the tokens already delivered in the first event. It is not Recorded results
Pooled median: 1070.79684 output tokens/s (1070.80 rounded). All 12 measured requests produced exactly 1024 output tokens; range 1051.52–1099.00, observed acceptance-length median 5.505376. Startup medians were 1068.93 and 1072.48. This describes the recorded distribution, not a guarantee that every request exceeds 1070 TPS. The input generator above was checked byte-for-byte at the token-ID sequence level against the frozen input. The timing client is the original client (SHA256 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
DeepSeek-V4.1 Flash spends substantial decode time in communication, small normalization/quantization launches and speculative metadata handling. This combines #39301/#39336 with the current
dsv4.1improvements and automatically selects supported paths. Current head:9623e9040af55bc2fdd2a05fddae362011be4359, merged with upstream1e9679db59ea7cf72a3ce2f2a547bc6e54c09ebc.SGLANG_OPT_DEEPGEMM_HC_PRENORMswitch.8bd60a9913.Latest upstream merge smoke check (
9623e9040a)Resolved the relocated communication-test runner conflict using the upstream B200 runner and kept the upstream deletion of the development-only
bench_c2_verify.py. The mHC graph cases remain intact. No subsequent optimization experiments are included.On 4×GB300, with the upstream-required
sglang-kernel 0.4.7andsgl-deep-gemm 0.2.0in an isolated environment:Previous 4×GB300 validation (
7322accddb)Same checkpoint revision
dba1be0a40aa45a94ad051997016db3960a90277, torch 2.13.0+cu130, Triton 3.7.1, FlashInfer 0.6.18, sglang-kernel 0.4.6.post1, sgl-deep-gemm 0.1.7, Transformers 5.12.1 and driver 580.105.08. No dependency overlay. Each arm uses a clean immutable checkout and a fresh server; GPU jobs are serial.8bd60a99137322accddbFinal BS1 starts have medians 1068.93 and 1072.48; all 12 measurements range 1051.52–1099.00. The upstream BS1 control has one fresh start with six measurements. Relative to this fresh upstream control, the complete PR improves observed C64 throughput by 34.66%. Earlier
c495C64 throughput was 7626.35 on the same node; that is historical evidence and is not substituted for the fresh upstream comparison.BS1: TP4/EP1, static DSpark block 5, fixed 4096 input/1024 output, simulated acceptance 5.5, temperature 0, ignore_eos; one warmup and six measured requests per start. Output tps excludes prefill and is timed from the first to last output event. Two independent final starts are pooled with every sample retained. C64: TP4/EP4, simulated acceptance 5.4, frozen 4096/1024 corpus, 64 warmups and 1024 measured requests, same pinned client. These are simulated-acceptance performance measurements, not real-acceptance throughput.
The 40 attention-combine kernels use two streams at BS1 (first layer, then the remaining 39 layers), and a single stream at C64. The same pattern appears upstream. The final C64 target graph contains 160 compensated TF32 projection calls and 80 Sinkhorn calls per replay on one statistics stream, confirming that large-row compensated mHC remains active. The final 64K check exercises the filtered graph beyond the 16K candidate budget (target graph ID 41, versus ID 38 at 4K); upstream uses the same full-filter graph at both lengths. No per-layer stream proliferation is observed in these traces.
Stream counts come from original Torch profiler
args.stream, without track remapping. BS1/4K has both 20-step GPU and 5-step CPU+GPU captures; BS1/64K has a 5-step CPU+GPU capture; C64 captures 5 steps only after the scheduler reports 64 requests decoding together, with CPU annotations confirming the batch size. Profiling follows and is separate from throughput timing. Counts cover all four TP ranks and distinguish all decode kernels from the 40-layer target graph.Correctness checks: changed-file pre-commit passed; 122 selected metadata/pool, wo_a and compensated-mHC regression tests passed. Three existing RoPE-policy mock cases fail on both final and unmodified upstream because the fixture lacks
q_head_norm; these are not reported as passing. No production workaround was added.Real-acceptance functional check: final and upstream each score 64/64 on the same held-out five-shot GSM8K samples, concurrency 1, temperature 0, seed 0, output cap 4096 and simulation disabled. This is a smoke check, not a full accuracy comparison. Full GSM8K/AIME/GPQA were not rerun for this cleanup; earlier results are tied to their tested commits. Removal of the experimental attention path is not presented as proof of accuracy parity across all retained optimizations. Hosted CI is separate from these dedicated-node checks.
CI States
Latest PR Test (Base): 🚫 Run #34941353724
Latest PR Test (Extra): ❌ Run #34941353254
Latest PR Test (AMD ROCm 10): ❌ Run #34941353597