Repository navigation
Conversation
added 5 commits
October 22, 2026 08:41
…te slot Problem: the GDN spec-verify ring had no length check. The verify writes the draft window at ring positions [write_pos, write_pos + T) and drops stores past the ring, so a ring shorter than the draft window (for example --speculative-num-draft-tokens 24 with the default --linear-replayssm-cache-len 16) made the commit fold a truncated window. A ring shorter than 16 (for example 8 for a block-8 draft) passed startup and then failed to compile the compact commit at the first verify, because the commit reduces over the whole ring with tl.dot. Separately, the cursor advance decided which rows are live from the request row, while the verify, the compact commit and the finish kernel decide it from the state slot. A committed row with a request row but state slot -1 had write_pos and is_flush advanced, was then skipped by the commit and the finish, and its next verify returned without writing its output. Fix: when the GDN spec ring is allocated, require a power-of-two --linear-replayssm-cache-len of at least 16 that holds the draft window, and name the flag in the error. Pass the state slots to the cursor advance and skip rows whose slot is at or below the null block, as the other three kernels do. The help of --linear-replayssm-cache-len states the constraints. Effect: a ring that cannot hold the draft window or the commit now fails at startup instead of truncating the window or failing at the first verify; this covers every algorithm on the GDN spec ring, EAGLE and NGRAM included. The four kernels agree on the live rows. No served output changes for configurations that fit the ring.
Problem: DSPARK with SGLANG_RAGGED_VERIFY_MODE=compact packs per-request verify lengths into one token list, but GDN target verify views a dense [batch, draft tokens] window from the conv on. Nothing refused that combination: the compact CUDA graphs fall back to eager for GDN, and eager verify misreads the tokens. On Qwen3.6 the first trimmed step was reported to hit an illegal memory access. Fix: refuse SGLANG_RAGGED_VERIFY_MODE=compact when the GDN backend is built for a target whose speculative algorithm builds ragged layouts. Effect: the unsupported combination fails at startup with a message naming SGLANG_RAGGED_VERIFY_MODE=static or cap-accept. Static and cap-accept verify (cap-accept verifies the dense window and only trims acceptance), draft workers and other algorithms are unchanged.
Problem: the ReplaySSM spec verify applied causality by multiplying with zeros: the (I + A)^-1 row update, D = R T^T and the output D F each multiplied later positions by exact zeros. The compact commit loaded the whole ring and gave the rows past the committed history a zero decay the same way. A non-finite activation at a rejected draft therefore turned the accepted outputs before it, and the committed state, into NaN, where the snapshot verify keeps both clean. Fix: select the causal terms with tl.where instead of multiplying them by zero, and mask the commit's ring loads to the committed history. Effect: rejected drafts cannot change accepted outputs, ring records of accepted positions, the state or later steps, whatever values they hold. On finite inputs the selects change the rounding of a few sums: over 12 chained steps at Qwen3.8-27B shapes with a bf16 state, 1e-5 to 3e-5 of the verify outputs differ, by up to 4e-3 of the largest output, and the chain error against a float64 reference is unchanged. The verify takes 4% longer with a bf16 state and 1% longer with an fp32 state (batch 128, GB300). The commit no longer reads the ring rows past the committed history; its load masking is the same as in the open upstream PR sgl-project#42688.
Problem: DSPARK and DFLASH commit verified linear-attention states through HybridLinearAttnBackend.update_mamba_state_after_mtp_verify, which only knew the KDA ReplaySSM fold. On GDN models with --enable-linear-replayssm-spec the verify already writes the compact request ring, but this commit would scatter per-draft snapshots that ReplaySSM does not allocate, so the configurator refused GDN with DSPARK/DFLASH, and those models had to keep a full state snapshot per draft token. Fix: move the GDN ring commit (cursor advance, fold of the accepted window into the checkpoint, conv rollback) into one helper, commit_gdn_replayssm_after_verify, used by the generic spec commit and by the backend commit. The backend commit folds every accepted window, whatever the state dtype, so the checkpoint is the committed state after each step, as in the KDA fold of the same function; radix donation and state copies read the slot directly. Only the lazy fold reads the draft length (for its early flush), so the helper takes max_spec_len only without fold_every_commit. Remove the configurator refusal, and assert before the snapshot scatter that snapshots exist. Compact verify stays refused on GDN. Effect: Qwen3.5/3.8 GDN models run DFLASH and DSPARK (static or cap-accept verify) with --enable-linear-replayssm-spec. The verify reads each state once and the commit reads and writes it once instead of writing a snapshot per draft token, and the snapshot memory is not allocated. The generic spec commit launches the same kernels as before.
… history Problem: the ReplaySSM spec verify replayed committed history for every row: two history dots per K tile and two after the loop, plus the history key, update and gate loads. A row folded at its last commit has no history, and DSPARK, DFLASH and the generic commit of a bf16 state fold every commit, so all of that work added exact zeros. The history decay and scaled updates were also built before the K loop and stayed live across it. Fix: run the history terms only when write_pos is nonzero, and build the decay and scaled updates after the K loop, next to their only use. Effect: per layer at batch 128 on GB300 (Qwen3.8-27B shapes, block 8), the verify without history takes 0.109 ms instead of 0.116 ms with a bf16 state and 0.154 ms instead of 0.188 ms with an fp32 state; with history (the lazy fold of an fp32 state) 0.209 ms instead of 0.250 ms. A bf16 state with history, which no commit path produces, takes 9% longer. Outputs, ring records and states are bit-identical with a bf16 state; with an fp32 state the compiled TF32 dots round differently (up to 8e-4 of the outputs differ, by up to 4.4e-3 of the largest output). A new chain test drives the lazy fold, so verifies replay up to 16 records of history against a float64 reference.
This was referenced Oct 22, 2026
added 2 commits
October 22, 2026 14:29
…selects Problem: the ReplaySSM spec verify keeps rejected drafts out by selecting the causal products, D + where(m, R * T, 0). The select sits between the multiply and the add, so they no longer fuse, and on finite inputs the (I + A)^-1 row update, D = R T^T and the output D F rounded differently than before the masking: over 12 chained steps at Qwen3.8-27B shapes, 1e-5 to 3e-5 of the verify outputs changed with a bf16 state (draft windows of 4 to 16), along with the ring records and the states folded from them. Fix: select the fused update instead, where(m, tl.fma(R, T, D), D), and mask the operand rather than the product in the (I + A)^-1 reduction. An implicit D + R * T inside the select is not enough: at a window of 16 the compiler still splits some of them. Effect: on finite inputs the outputs, ring records, committed states and cursors are bit-identical to the verify before the masking, for draft windows of 4, 8 and 16 with a bf16 state, which folds every commit (GB300). With an fp32 state, which also takes the lazy fold of the generic spec commit, they are bit-identical when measured without the previous commit; on this branch that commit's TF32 rounding change remains. Non-finite rejected drafts stay contained. At batch 128 the verify takes 0 to 7% less time than with the split products: 79 instead of 85 us with a bf16 state at a window of 4, 372 instead of 383 us with an fp32 state at 16. The non-finite rejected draft test now also runs the lazy fold of the generic spec commit, over four steps with committed history before the second poisoned window.
…d_verify Problem: the refusal test named the algorithms that build ragged verify layouts and expected DFLASH to pass the GDN check. The check keys on supports_ragged_verify(), so the test breaks as soon as another algorithm gains ragged verify, although the check itself is unchanged. Fix: give the test runner an algorithm stub whose supports_ragged_verify() returns the case's value, and cover a ragged target, a non-ragged target and a draft worker. Effect: test only; no behaviour change. The test no longer depends on which algorithms support ragged verify.
jianc99
force-pushed
the
av/replayssm-dflash
branch
from
October 23, 2026 00:16
8eeda5a to
bf3c6b7
Compare
Qiaolin-Yu
approved these changes
Oct 23, 2026
This was referenced Oct 24, 2026
Open
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
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.
Motivation
On Qwen3.5/3.8 GDN targets, DFLASH and DSPARK verify by writing a full SSM state snapshot per draft token (FlashInfer MTP or the Triton recurrent verify) and scatter the accepted one afterwards.
--enable-linear-replayssm-specavoids those snapshots with a compact per-request ring, but it was refused for DFLASH/DSPARK on non-KDA models: their commit,HybridLinearAttnBackend.update_mamba_state_after_mtp_verify, only knew the KDA ring and would have scattered snapshots that ReplaySSM does not allocate.Wiring the GDN ring into that commit also turned up problems that affect every algorithm on the GDN spec ring:
Separately, nothing refused DSPARK's compact verify on GDN targets. Compact verify packs per-request verify lengths into one token list, but GDN target verify views a dense
[batch, draft tokens]window, and eager compact verify on Qwen3.6 has been reported to hit an illegal memory access.Modifications
Fixes first, then the feature and the speedup, then two follow-up fixes:
MambaPoolrequires a power-of-two--linear-replayssm-cache-lenof at least 16 that holds the draft window, and the flag's help says so. A ring of 8 used to pass startup and then fail to compile the compact commit at the first verify. The cursor advance skips rows without a state slot, as the verify, commit and finish kernels already do.GDNAttnBackendraises at startup when the target's speculative algorithm builds ragged layouts andSGLANG_RAGGED_VERIFY_MODE=compact. Static and cap-accept verify (cap-accept verifies the dense window and only trims acceptance), draft workers and other algorithms are unchanged. This commit is independent of the others; see "Related upstream work" for how it interacts with [Fix] Honor DSpark SPS profiler budget pins and handle ragged mRoPE positions sgl-project/sglang#38897.tl.whereinstead of multiplying by zero. The commit masks its ring loads to the committed history.commit_gdn_replayssm_after_verify. The generic spec commit and the backend commit (DFLASH/DSPARK) both call it. The backend commit folds every accepted window whatever the state dtype, as the KDA fold in the same function does, sotemporalis the committed state after each step and radix donation and state copies read the slot directly. Only the lazy fold needs the draft length (for its early flush), so the helper takesmax_spec_lenonly withoutfold_every_commit. The configurator refusal is removed, and an assert before the snapshot scatter checks that snapshots exist.write_pos > 0, and they are built after the K loop.tl.where(m, a * b, 0)and added it to the accumulator, so the multiply and the add no longer fused and the(I + A)^-1row update,D = R T^Tand the outputD Frounded differently from main. The(I + A)^-1reduction now masks the operand, and the other two select the fused update,tl.where(m, tl.fma(a, b, acc), acc), never the product alone. An implicitacc + a * binside the select is not enough: at a window of 16 the compiler still splits some of them. Rejected positions still never reach the result.supports_ragged_verify()value, so it keeps passing when another algorithm gains ragged verify.No new flags or environment variables.
Behavior changes:
--linear-replayssm-cache-len 32.SGLANG_RAGGED_VERIFY_MODE=compactnow fails at startup.Memory: the ring replaces the per-draft snapshots. At TP1 with 128 requests on Qwen3.8-27B it takes 3.04 GB instead of 72.6 GB of snapshots, and the KV cache grows from 1.13M to 2.03M tokens. DFLASH and DSPARK fold every commit, so their ring never holds more than one window, but 16 is already the smallest ring the compact commit supports (it reduces over the whole ring with
tl.dot), so the default fits drafts of up to 16 tokens without waste.Accuracy Tests
Validated on GB300 (SM100) only, upstream image
lmsysorg/sglang:v0.5.21with this tree mounted. The GPU test's numeric bounds are relative to the Triton snapshot verify, which runs on every GPU; it has not been run on SM90 yet.test/registered/unit/spec/test_gdn_replayssm_dflash_commit.py(new, CPU) covers:test/registered/attention/unittests/gdn/test_gdn_replayssm_direct_commit.py(new, GPU) runsGDNAttnBackend.forward_extendandupdate_mamba_state_after_mtp_verifyat Qwen3.8 head shapes (HV/H = 3, K = V = 128, block 8, ring 16) for bf16 and fp32 states:test_ngram_mamba_verify_update.pyandtest_linear_attn_config.py, and a regression set oftest_replayssm_ring_accounting.py,test_mamba_unittest.py,test_linear_replayssm_decode.py,test_mamba_track_state_dtype.py,test_gdn_prefill_backend_policy.py,test_gdn_mis_metadata.pyandtest_gdn_cutedsl_ring_verify.py(53 passed on this branch and on main). Every commit passes the tests that exist at that commit.These serving checks predate item 6.
Speed Tests and Profiling
Setup: Qwen3.8-27B (48 GDN + 16 attention layers), DFlash2 block-8 draft, GB300, TP1,
--max-running-requests 128, bf16 mamba state, radix cache off, overlap scheduling on. Client: SPEED-Bench first turns, T=1, top_p 0.95, top_k 20, up to 4096 output tokens.Decode tok/s = per-request decode rate (output tokens after the first / time after the first token), times concurrency; c1 is the single-request decode rate.
The fourth ReplaySSM run at c64 and c128 is the final branch before item 6; the first three ran an earlier revision with the same kernels. No run in this table includes item 6 (the end-to-end numbers in #8 do, on the merged stack); its verify kernel time is in the last table.
At c128 the running cap equals the concurrency, so about one request waits in the queue. #8 shows that this setup reads 7-15% higher at that level than a larger cap, on both sides, so compare c128 only within this table.
Acceptance length is unchanged from c8 to c128: 3.55–3.69 for main and 3.55–3.70 for this PR.
At c1 acceptance is dominated by which trajectories the requests take: 3.67 for main vs 3.53 for this PR on the first 16 requests at T=1, 3.57 vs 3.53 on 48 requests, and 3.96 vs 4.05 greedy. The verify step rate, decode tok/s divided by acceptance length, is within 1% on 48 requests: 93.7 vs 94.3 steps/s at T=1 and 92.2 vs 92.0 greedy (95.0 vs 92.9 on the first 16 requests).
Verify kernel time per layer, batch 128, Qwen3.8-27B shapes, GB300 (measured before item 6):
A bf16 state with history takes 9% longer after 5, but no commit path produces one.
ReplaySSM spec verify kernel time with item 6 (main is the kernel before this PR), batch 128, Qwen3.8-27B shapes, GB300, median:
Item 6 takes 0 to 7% less time than the split products. Applied to commits 1–4 without commit 5, it is within 2.3% of main in every row, so the masking no longer costs time.
Numerics on finite inputs:
Related upstream work
GDNAttnBackendfor every GDN verify path (supports_ragged_verify_graph = True), so compact verify runs on GDN without ReplaySSM. Commit 2 refuses exactly that configuration at startup. Its conv changes handle packed rows (a varlen kernel for CUDA chain verify, a dense scratch otherwise), but it does not change the FlashInfer GDN verify kernel (gdn_flashinfer.py), which still views the verify tokens as[bs, seq_len // bs]; with--linear-attn-verify-backend flashinfer(the verify default whenever the decode backend is flashinfer) a packed batch would be read as uniform. What to do: if [Fix] Honor DSpark SPS profiler budget pins and handle ragged mRoPE positions sgl-project/sglang#38897 lands first, drop commit 2 or narrow it to the verify kernels that cannot read packed rows; if this PR lands first, [Fix] Honor DSpark SPS profiler budget pins and handle ragged mRoPE positions sgl-project/sglang#38897 needs to lift the refusal for the paths it makes safe. [Feat] gdn: run target verify on packed ragged tokens (ReplaySSM) #6 is our version of that lift, for the ReplaySSM path only.