Conversation
…ged MLA kernel Mixed-chunk batches currently send all tokens (prefill chunk + one decode token per running request) through the bf16 sparse-prefill workspace path, which dequantizes the fp8 KV cache every layer. Split the batch at the prefill/decode boundary instead: the decode tail uses the fp8-native flash_mla_with_kvcache kernel (no dequant), and the sparse-prefill workspace only covers the prefill requests' windows. Gated behind SGLANG_OPT_MIXED_SPLIT_DECODE_ATTN (default off). 8xB300 DSV4-Pro 8k1k conc2048: 9064 -> 9796 tok/s (+8.1%), TPOT 192.6 -> 174.5 ms. GSM8K under concurrency: 0.970 (on) vs 0.980 (off), within noise.
Mixed-chunk steady state runs fully eager today: launching a ~180 ms step costs ~170 ms of Python launch time, so the GPU queue never builds backlog and every scheduler pause becomes a GPU bubble. Breakable CUDA graph (BCG) cuts the launch cost ~6x, but three integration issues erase the gain: 1. The WAR barrier fast-path event is only published by the decode graph; BCG prefill replays fall back to waiting for the whole in-flight forward, turning batch-assembly D2H syncs into a ~90 ms per-step wall. 2. The mixed-chunk decode/prefill attention split is disabled under replay (MIXED normalizes to EXTEND and the metadata refresh drops the boundary). 3. The eager logits tail sizes its DP-attention gathers from the static batch's capture-time padded token counts, so the vocab all-gather and hidden dp_gather run over every padded token (e.g. 8x2048 rows) instead of the sampling positions. This PR: pre-builds the sparse-prefill chunk cache (sync-free via CPU-side total_swa) and the split boundary in the replay-prepare hook, publishes the WAR read-done event from the prefill runner when the backend opts in, and passes the live logprob token counts to the eager logits tail. BCG for DSV4+megamoe+DP-attention is enabled via SGLANG_EXPERIMENTAL_FORCE_BCG=1, which bypasses the four incompatibility rules pending proper graduation. 8xB300 DSV4-Pro FP4, 8k1k conc 2048, radix forced-miss, num_prompts 6x conc: 9828 -> 10076 tok/s (+2.5% over the eager mixed-split baseline), TPOT 177.2 -> 172.5 ms, GSM8K 0.985 vs 0.980 eager control.
Contributor
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
21 tasks
41 tasks
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.
Status: WIP / experimental draft — not ready to merge. Opened to share the approach and measurements; see the open-problems list below.
Motivation
With mixed-chunk enabled, the steady state is 100% mixed steps running fully eager: launching a ~180 ms step costs ~170 ms of Python launch time (measured), so the GPU queue never builds backlog and every scheduler pause surfaces as a GPU bubble (~7 ms/step). BCG cuts launch CPU ~6x (2156 kernel launches -> 62 graph replays + eager attention breaks), but naively enabling it made throughput worse. Three integration issues were found and fixed.
What this PR does
Scheduler._apply_war_barrier) was only fed by the decode graph; BCG replays fell back towait_stream(forward), turning batch-assembly D2H syncs into a ~90 ms/step wall. The DSV4 backend now pre-builds the sparse-prefill chunk cache from the live batch in the replay-prepare hook (sync-free:total_swacomputed CPU-side), moving all scheduler-shared reads to the snapshot, and the prefill runner publishes the event when the backend opts in (war_reads_done_at_snapshot).forward_mode.is_mixed(), which is normalized to EXTEND at replay; the boundary now flows through the refreshed metadata.global_num_tokens_for_logprob_*, making the DP-attention hidden gather + vocab all-gather + reorder copies run over every padded token (e.g. 16384 rows) instead of the sampling positions (~1.8k rows); the runner now passes the live counts.Enablement is gated behind
SGLANG_EXPERIMENTAL_FORCE_BCG=1, which bypasses the four breakable-CG incompatibility rules (MLA / DSV4 capture-pool pressure / MoE A2A bucket cap / DP-attention) pending proper graduation.Results
8xB300, DeepSeek-V4-Pro FP4, 8k1k, conc 2048 (megamoe + dp-attention + mixed-chunk), radix cache with forced misses, num_prompts = 6x conc, warmup = 2x conc; additionally
--cuda-graph-bs-prefill 512 1024 1536 2048 2560 2816andSGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK=2816:--enable-dp-lm-headAll runs 12288/12288 successful, zero scheduler exceptions. Steady-state traces confirm: launch CPU 171 -> 29 ms/step, graphs replayed on >66k mixed steps under DP attention.
--enable-dp-lm-headis neutral (+0.3%, within run variance) once the logits-gather sizing is fixed — before the fix it would mask the oversized gather entirely (the DP gather path is skipped), at the cost of replicating the LM head (~1.85 GB/GPU). We keep it off by default.Why this is a draft — open problems
NUM_MAX_TOKENS_PER_RANK; DSV4 capture-pool pressure is bounded by a small bucket list (the old default ran to 16384); the MLA rule should exclude the DSV4 backend (its whole attention is an eager break); DP-attention is validated here for one topology only.--cuda-graph-bs-prefill); needs defaults derived fromchunked_prefill_size+ the mixed budget (steady mixed steps land exactly on chunk-size tokens by construction).war_reads_done_at_snapshot); needs review that eviction/allocator interactions cannot invalidate it, plus a debug assertion story.recommended_max_tokens(include_prefill=False)) needs prefill-aware sizing.Depends on #30338.
CI States
Latest PR Test (Base): ❌ Run #28889917824
Latest PR Test (Extra): ❌ Run #28889917584