Repository navigation
Conversation
added 30 commits
October 22, 2026 06:24
Problem: in linear-chain speculative verify, causal_conv1d_update saves the
intermediate conv window after every draft token with three lane-strided
element stores. MambaPool's deduplicated window view (step stride == window
stride) maps those stores onto one shared row per channel, so the same
addresses are written up to three times by scattered narrow stores. On GB300
(dim 10240, bf16, T=8, bs 128) the verify conv takes 106 us with the
deduplicated window versus 39 us with a dense one.
Fix: when the window view's step stride equals its window stride, step t's
window is columns t..t+2 of the row [state[1], state[2], x_0, ..., x_{T-1}].
For a chain (no tree mask) with kernel width 4, write that row once as a
[BLOCK_N, next_pow2(T+2)] tile before the token loop and skip the per-token
stores. Tree verify, other widths and dense windows keep the per-token
stores. The bytes written are the same; the new test checks them bitwise
against the per-token path.
Effect: kernel time on GB300, dim 10240, bf16, T=8: bs 1 5.4 -> 2.5 us,
bs 64 57.4 -> 10.0 us, bs 128 106.5 -> 17.6 us (T=16, bs 128: 322.6 ->
39.1 us). Every CUDA chain verify on a model with a width-4 conv (GDN,
Mamba2) takes this path. The tile store is SubSir's.
…nce head Confidence-scheduled verification sizes each request's verify window from per-position acceptance probabilities. DSpark drafts get them from a trained confidence head. Draft models without one, such as DFlash2, have no confidence source, so they cannot use the planner. Add OnlineAcceptanceEstimator, a logistic model calibrated while serving: p_k = sigmoid(slope * x_k + bias_k). x_k is logit(q_max) of the temperature-scaled draft distribution at position k (greedy rows at T = 1, clamped to +-40). One slope is shared by all positions and each position has its own intercept, starting from 0.5 and -0.2. After verification, observe() labels the accepted drafts and the first rejected one and adds them to the Newton (IRLS) statistics. Every 100 calls it solves the arrowhead system with ridge 1e-3, damps each update by n / (n + 50), drops non-finite steps and steps that would make the slope negative, and clears the statistics. The feature, the accumulation and the refit run in two Triton kernels; feature_and_predict() also returns the probabilities from the feature launch. The refit counter is a device tensor and the coefficients are updated in place, so no call synchronizes with the host. The estimator therefore works with overlap scheduling and inside CUDA graphs. A single program applies each update, so ranks that see the same inputs keep identical coefficients. Effect: no behavior change on its own; it is the confidence source for confidence-scheduled verification of such drafts. Replayed over recorded static DFlash2 traces of Qwen3.8-27B (about 2M draft positions each), the AUC of the predicted survival against acceptance is 0.898 at T = 1 and 0.921 at T = 0, against 0.886 and 0.920 for the raw q_max product. Feature, predict and observe take about 10 us of GPU time per step in a CUDA graph (14 us at bs = 512).
Problem: the deduplicated conv window tile writes [state[1], state[2], x...],
which is the right row only for kernel width 4, but only the Python wrapper
enforced that; a caller that set the flag for another width would get wrong
window bytes. The new test compiled 32 kernel variants (dim is a constexpr),
and it compared conv outputs of two separately compiled kernels bitwise, which
is not guaranteed on every backend the file runs on. The MambaPool comment
still described the kernel's window write as overlapping per-token stores.
Fix: static_assert KERNEL_WIDTH == 4 inside the tile branch. Test one dim
(1000, a partial channel block) instead of two; T in {2, 4, 8, 16} still covers
every tile width. Keep the bitwise checks on the window and conv state, which
are copies, and compare the conv output with the file's update tolerances.
Update the MambaPool comment.
Effect: no runtime change. The test compiles half as many kernel variants.
build_speculative_verify_target_probs divided the whole [rows, vocab] logits by the temperature and ran torch.topk over the result before building the top-k/top-p distribution. On a 248k vocabulary this costs 3.7 ms per verify step at 1024 rows (128 requests x 8-token blocks) on GB300 and allocates an extra [rows, vocab] fp32 tensor. DFLASH, DSPARK, UNO dense verification and sampling-mask capture all build their target distribution here. Select the top-k on the raw logits and divide only the k selected values by the temperature. A positive temperature keeps the order of the logits, so the probabilities are unchanged. The selection goes through a new kernels/ops/sampling/topk.py op, resolved once per platform: FlashInfer's sorted, deterministic radix top-k when FlashInfer is available, torch.topk otherwise. It accepts any k up to the vocabulary size. The DFlash2 selector's private copy of the same FlashInfer call now uses this op. When several logits tie at a row's top-k boundary, which of them keep the mass is chosen deterministically and may differ from torch.topk; the kept count and the probability values are the same. At 1024 rows the distribution takes 1.0 ms instead of 3.7 ms, at 4096 rows 3.2 ms instead of 14.1 ms, and the peak temporary memory halves.
…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.
… layouts Problem: only the Python wrapper kept tree verify off the deduplicated conv window tile; a caller that set both flags would get the tile plus the tree branch's per-token stores into the same window. The wrapper comment tied the tile to MambaPool's overlapping view, while the real argument is about the strides alone. Kernels that do not take the tile still got a seqlen-dependent NP2_WINDOW_COLS. The test did not cover a num_accept_tokens state offset, MambaPool's axis-0 window layout, or an odd T. Fix: static_assert that the tile branch has no tree mask, next to the width assert. Rewrite the comment: with step stride == window stride == s, the per-token stores write row column t + w at (t + w) * s and the tile writes column c at c * s, the same addresses and values for any layout with those strides. Pass NP2_WINDOW_COLS=1 when the tile is off. Test T = 3, a num_accept_tokens offset, and the axis-0 window (a [T + K - 2, dim] row seen transposed); T = 3 replaces T = 4 in the plain cases, so T + 2 still pads to tiles of 4, 8, 16 and 32. Effect: no runtime change. The test compiles 24 kernel variants instead of 16.
The refit dropped any slope step that would make the slope negative. From a slope well above the optimum the predictions saturate, the curvature is small and the damped Newton step overshoots past zero. The next round starts from the same slope and drops the same step, so the slope never moves: from 3.0 on drafts generated with slope 1.0, and from the default 0.5 on drafts whose feature carries almost no signal (slope 0 or 0.03), it stayed where it started through every refit. Halve such a step until the slope stays positive, at most 8 times, and drop it only if it still crosses zero. The intercepts take their Newton step given the slope step actually taken. The rule runs in the same single program, so the refit stays deterministic, sync-free and graph-capturable. Effect: within ten refits the slope reaches 0.96 in the first case and 0 to 0.06 in the second. Refits whose full step keeps the slope positive are unchanged; from the default start on drafts generated with slope 1.3 no step crosses zero. observe() GPU time is unchanged at the default refit interval and within 0.1 us when every call refits.
…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.
Problem: gdn_replayssm_compact_commit_kernel ran one CTA per (64-column V tile, request, value head x layer), 590k CTAs per step at batch 128 on a 48-layer, 48-value-head GDN model. Every CTA reloaded the request scalars, ring positions and key tile, and folded the history with four chained tensor-core products (key and key residual times the high and low update parts). With a bf16 checkpoint the commit folds every step, and it took 5.40 ms per step at batch 128. Fix: one program per (request, layer, key head) covers the HV / H value heads that read that key head, in 32-column V tiles walked by one flat loop. The history is folded 16 steps at a time: the key and key-residual rows and the high and low update parts of a chunk are stacked along the reduction axis, so each chunk is one tensor-core product with fp32 accumulation, and the tile size no longer depends on the ring length. A ring of one chunk (the default 16) loads its key tile once per program; longer rings visit only the chunks that hold history. fp32 rings carry no low part. Grouping alone and stacking alone do not help (0% and 44% slower); the gain needs both. Fold and track snapshots are computed for every row and masked at the store. Slot and layer strides, null rows, the early exit and the cursor finish are unchanged. The result differs from the previous kernel by fp32 summation order only: at most one bf16 ulp, or 2.8e-6 for fp32 checkpoints with bf16 rings (up to 1.6e-4 with fp32 rings, TF32 operands). Effect: GB300, 48 layers, K = V = 128, ring 16, bf16 checkpoint, 4 history steps, every row folding: - HV 48 / H 16 (TP1): bs128 5.40 -> 3.85 ms, bs8 341 -> 253 us, bs1 39.8 -> 42.0 us. - HV 12 / H 4 (TP4): bs128 1.35 -> 0.97 ms, bs8 90 -> 72 us, bs1 11.5 -> 17.9 us. With half of the rows taking a track snapshot, bs128 TP1 goes from 10.07 to 7.37 ms. Shared memory no longer grows with the ring: at most 37 KB and no spills, where the previous kernel needed 98 KB at ring 128 and spilled with track rows. On a raised --linear-replayssm-cache-len, a 4-step fold at bs128 TP1 goes 5.8 -> 7.5 ms at ring 32, 8.5 -> 8.1 ms at ring 64 and 13.8 -> 8.2 ms at ring 128. Folding a long history there is slower (ring 64, 56 steps: 9.3 -> 33 ms); only fp32 checkpoints fold that much at once.
Problem: the head-group commit kernel replaced the per-head kernel for every ring length. It is faster on the default 16-step ring, but on a raised --linear-replayssm-cache-len it folds the history 16 steps at a time and its time grows with the ring length and the history. At batch 128 with TP1 heads it is 1.2-2.7x slower than the per-head kernel at ring 32, and 3.1-5x slower when a full ring of 64 or 128 steps folds (for example, bf16 checkpoint, head group vs per head: 34.9 vs 9.6 ms at ring 64, 65.4 vs 17.3 ms at ring 128). A nearly full ring (all but one draft window) is what a fold at capacity commits, as with fp32 checkpoints on EAGLE and MTP. Fix: keep the per-head kernel and pick the kernel on the host from the ring length. A ring of one 16-step fold chunk commits per head group, a longer ring per value head. The ring length is a static shape, so the choice needs no sync, is the same on every TP rank, and does not depend on the batch: the two kernels differ in fp32 summation order, so a batch-size rule would make a request's checkpoint depend on the batch it runs in. The circular commit tests run every case through both kernels, and a new test checks which kernel each ring length launches. Effect: GB300 kernel time, 48 layers, K = V = 128, bf16 rings with residuals, every row folding, batch 128, TP1 heads, bf16 checkpoint, before -> after: - ring 16: unchanged, 3.87 ms (per-head kernel: 5.45 ms). - ring 32: 4 history steps 7.58 -> 5.79 ms, full ring 15.8 -> 5.84 ms. - ring 64: 4 steps 8.17 -> 8.47 ms, full ring 34.9 -> 9.62 ms. - ring 128: 4 steps 8.32 -> 13.8 ms, full ring 65.4 -> 17.3 ms. fp32 checkpoints, track rows, batch 1 and 8 and TP4 heads follow the same pattern: full rings 2.4-7x faster, every history on ring 32 1.2-3.9x faster, and 4- to 16-step histories on rings 64 and 128 up to 1.7x slower. Rings longer than 16 fold short histories only when every commit folds, where the ring only has to fit one draft window, or in a track snapshot taken soon after a lazy fold. On the 16-step ring the head-group kernel is still up to 12 us slower at batch 1 (TP4 heads, fp32 checkpoint, track rows: 45.1 vs 33.5 us for the per-head kernel).
Problem: in linear-chain speculative verify, causal_conv1d_update saves the
intermediate conv window after every draft token with three lane-strided
element stores. MambaPool's deduplicated window view (step stride == window
stride) maps those stores onto one shared row per channel, so the same
addresses are written up to three times by scattered narrow stores. On GB300
(dim 10240, bf16, T=8, bs 128) the verify conv takes 106 us with the
deduplicated window versus 39 us with a dense one.
Fix: when the window view's step stride equals its window stride, step t's
window is columns t..t+2 of the row [state[1], state[2], x_0, ..., x_{T-1}].
For a chain (no tree mask) with kernel width 4, write that row once as a
[BLOCK_N, next_pow2(T+2)] tile before the token loop and skip the per-token
stores. Tree verify, other widths and dense windows keep the per-token
stores. The bytes written are the same; the new test checks them bitwise
against the per-token path.
Effect: kernel time on GB300, dim 10240, bf16, T=8: bs 1 5.4 -> 2.5 us,
bs 64 57.4 -> 10.0 us, bs 128 106.5 -> 17.6 us (T=16, bs 128: 322.6 ->
39.1 us). Every CUDA chain verify on a model with a width-4 conv (GDN,
Mamba2) takes this path. The tile store is SubSir's.
Problem: the deduplicated conv window tile writes [state[1], state[2], x...],
which is the right row only for kernel width 4, but only the Python wrapper
enforced that; a caller that set the flag for another width would get wrong
window bytes. The new test compiled 32 kernel variants (dim is a constexpr),
and it compared conv outputs of two separately compiled kernels bitwise, which
is not guaranteed on every backend the file runs on. The MambaPool comment
still described the kernel's window write as overlapping per-token stores.
Fix: static_assert KERNEL_WIDTH == 4 inside the tile branch. Test one dim
(1000, a partial channel block) instead of two; T in {2, 4, 8, 16} still covers
every tile width. Keep the bitwise checks on the window and conv state, which
are copies, and compare the conv output with the file's update tolerances.
Update the MambaPool comment.
Effect: no runtime change. The test compiles half as many kernel variants.
Problem: under ragged (compact) speculative verify each request owns a variable number of packed verify tokens, but causal_conv1d_update only takes a dense [batch, dim, seqlen] window. A caller has to scatter the packed tokens into a zero-padded dense buffer and gather the output back on every layer. Fix: optional query_start_loc and max_seqlen. With them x and out are packed [num_tokens, dim]; sequence i owns rows query_start_loc[i] to query_start_loc[i + 1], at most max_seqlen of them. Each sequence reads its own rows (zeros past its length) and updates the conv state and the intermediate windows, per-token or deduplicated tile, exactly like a dense row zero-padded to max_seqlen; output is written only for its own rows. A sequence with no rows exits right after the slot check, and packed rows past the last sequence are never touched. Linear chains only (no tree mask, no num_accept_tokens). The packed layout and the empty-row exit follow SubSir's packed verify conv. Effect: a packed call matches the zero-padded dense call bit for bit (real output rows, conv state, both window layouts, CUDA-graph replay), so ragged verify needs no per-layer staging around the conv. GB300, dim 10240, bf16, T=8, deduplicated window, bs 128: dense 17.4 us (unchanged), packed with every row full 18.1 us, packed with mixed lengths (mean 4.6) 17.0 us.
Problem: GDN target verify views the packed tokens as a dense [batch, draft tokens] window from the conv on, so compact verify (SGLANG_RAGGED_VERIFY_MODE =compact, which packs a variable number of verify tokens per request) is refused on GDN models and Qwen3.5/3.8 always verify the full draft block. Fix: with --enable-linear-replayssm-spec under compact verify, every target verify reads its packed tokens in place: the conv reads them by query_start_loc, and the ReplaySSM spec verify, which already reads rows by query_start_loc, takes the packed q/k/v views. Nothing is staged to a dense [batch, draft tokens] buffer. - The path is chosen once at startup, not from the batch, so a CUDA graph captured with any layout (including a uniform one) replays every other. A uniform layout gives the dense result bit for bit. - GDN reports ragged verify graph support only for this path. The snapshot verify (FlashInfer MTP or Triton recurrent) stays dense, and the startup refusal of compact verify now points at --enable-linear-replayssm-spec. - The ReplaySSM spec predicate is computed once in the backend constructor instead of in two places. - The ReplaySSM verify writes zeros to packed rows past the last request (CUDA-graph token-tier leftovers), so every output row is finite. Rows of graph padding (mamba slot -1 or no tokens) were already zero or skipped. - The ragged ReplaySSM startup check keeps requiring a ring-writing verify backend on KDA only; the GDN ReplaySSM verify writes the ring under every verify backend. Effect: DSPARK (and DFLASH once it produces compact layouts) can run compact verify on Qwen3.8 with CUDA graphs. A ragged step matches the dense step of the same real tokens bit for bit (outputs, ring records, committed state and conv), in eager mode and in graph replay, including padding rows, uncovered tier leftovers and an extended last row. Static verify outputs for real rows are unchanged; output rows of padding slots in padded dense graphs are now zero.
… layouts Problem: only the Python wrapper kept tree verify off the deduplicated conv window tile; a caller that set both flags would get the tile plus the tree branch's per-token stores into the same window. The wrapper comment tied the tile to MambaPool's overlapping view, while the real argument is about the strides alone. Kernels that do not take the tile still got a seqlen-dependent NP2_WINDOW_COLS. The test did not cover a num_accept_tokens state offset, MambaPool's axis-0 window layout, or an odd T. Fix: static_assert that the tile branch has no tree mask, next to the width assert. Rewrite the comment: with step stride == window stride == s, the per-token stores write row column t + w at (t + w) * s and the tile writes column c at c * s, the same addresses and values for any layout with those strides. Pass NP2_WINDOW_COLS=1 when the tile is off. Test T = 3, a num_accept_tokens offset, and the axis-0 window (a [T + K - 2, dim] row seen transposed); T = 3 replaces T = 4 in the plain cases, so T + 2 still pads to tiles of 4, 8, 16 and 32. Effect: no runtime change. The test compiles 24 kernel variants instead of 16.
Problem: compute_spec_mrope_positions views the target-verify positions as [bs, tokens / bs]. A ragged verify packs verify_lens[i] tokens for request i plus graph padding, so a text-only batch fails the view whenever the token count is not a multiple of bs, and a multimodal batch adds each request's mRoPE delta to the wrong tokens. Fix: keep the positions flat. Text-only batches need no per-request split. Multimodal batches expand the per-request deltas by the ragged layout's verify lengths (graph padding gets delta 0), with a device-side search so no lengths are read back on the host; uniform blocks keep the even split. Effect: mRoPE targets (Qwen3.5/3.8) can run DSpark-style ragged verify. Uniform blocks compute the same positions as before.
… rank Problem: under overlap scheduling the verify budget is planned from the confidence relayed two steps back through a pinned host ring. The relay returned None when that slot's copy event had not completed yet, which is rank-local timing: one TP rank could plan a budget while another planned none, so the ranks entered different collectives and verify graph tiers. The ring copy also read confidence_buf on a side stream with no ordering against the next step's scatter into the same buffer. Fix: wait for the slot's copy instead of skipping it, and make the scatter wait (device side) for the previous ring copy. This is the source change of the open upstream PR sgl-project#42717 (without its new test); drop this commit if that one lands first. Effect: every rank resolves the same lagged confidence, so all ranks plan the same budget. The wait is on a copy issued two steps earlier and is normally already complete.
Problem: DSpark's adaptive verification (SGLANG_RAGGED_VERIFY_MODE cap-accept / compact: a per-step verify budget from a cost table, top-k verify lengths by survival, packed ragged verify graphs) is open to DSPARK only, and it needs a trained confidence head. DFLASH drafts, including DFlash2 selector drafts, have no confidence head, so they always verify the full block. Fix: run DFLASH through the existing DSpark machinery, with an online acceptance estimator as the confidence source. - The DFLASH worker builds the DSpark verify planner (with the worker's own confidence) and the target-verify executor, provides the overlap budget-prepare hook, and publishes its confidence into the existing relay. DSPARK-only gates (supports_ragged_verify, the confidence relay) cover the DFLASH family. - Feature: logit(q_max) of the selector row each draft token was drawn from, at the request's temperature (temperature 1 for greedy rows); it scores the draft distribution, never the drawn token. The draft cuda graph computes it next to the walk. After verify the estimator observes the verified and accepted draft counts. - Every DFLASH accept path stops at the planned verify length: the selector sampling accept passes the cutoff to DSpark's sampling kernel, and the greedy accept caps through AcceptGreedy, which now also takes precomputed argmax ids. The bonus comes from the last verified row. - The ragged window reads the anchor at the block's real row stride (DFLASH blocks are gamma + 1 wide, DSpark's gamma). - TP: like a DSPARK confidence head, the estimator gives every rank the same confidence for the same draft (its updates are deterministic and see the broadcast acceptance), and the relay resolves the same lagged copy on every rank, so every rank plans the same budget and graph tier. The lengths are broadcast from rank 0, as for DSPARK. - Compact verify scatters the packed rows to the block layout row by row, which works on a vocabulary shard as well, so DFLASH keeps its deferred target-verify gather in every mode; the logit adjustments run once, on the gathered logits, for both verify paths. - Ragged modes need a selector (DFlash2) draft and are refused under attention DP. Compact mode is refused at startup when the target attention backend cannot verify packed ragged tokens (supports_ragged_verify_graph), as a GDN hybrid target cannot yet. Effect: SGLANG_RAGGED_VERIFY_MODE=cap-accept and compact work for DFLASH with the same planner, cost table, allocation and compact graphs as DSPARK. The default static mode is unchanged.
Problem: the verify cost table that plans the compact verify budget comes from python/sglang/benchmark/dspark_sps_profiler.py, which only accepts DSPARK servers, and the scheduler only serves the per-step records and the pinned-budget control it reads to DSPARK workers. With a flat (uninitialized) table, compact mode also ignores the pinned budget the profiler's --fracs cells set, so those cells all measure the full verify. Fix: accept DFLASH servers in the profiler, serve the step records and the pinned budget to the DFLASH worker (it records the same steps), and let a pinned budget schedule its lengths without a cost table. This also applies to DSPARK in compact mode, so the scheduler now refuses to pin a budget when SGLANG_SIMULATE_ACC_LEN is above 0 and not 1.0: a simulated acceptance could run past the trimmed lengths. Effect: DFLASH builds its cost table with the same profiler run as DSPARK. Serving without a pinned budget is unchanged.
Compact verify planned, packed and scattered every decode step, even when the plan could only replay the compact graph of the full block: a single request always lands on that graph, and a small batch often does. That per-step host work (budget pricing, top-k lengths and their broadcast, the packed window, the compact-to-block scatter) cost about 1.3 ms at TP4 batch 1 and made the step host-bound, so c1-c8 ran up to 20% slower than static verify. A batch whose shortest plan already replays the full block's graph now prices no budget, and a plan whose graph tier is the full block's takes the cached uniform layout, as verify-all does. A full-block layout skips the packing: its verify rows are the verify window, and its logits and hidden states are already in the block layout. TP4, T=1, decode tok/s at c1/c2/c4/c8: 520/1033/1970/3800 -> 628/1139/2205/4029 (static 652/1229/2343/4336). c64, c128 and c512 are unchanged.
# Conflicts: # python/sglang/kernels/ops/mamba/causal_conv1d_triton.py
This was referenced Oct 22, 2026
jianc99
force-pushed
the
av/integration
branch
2 times, most recently
from
October 23, 2026 00:16
9ef414d to
363b4d0
Compare
jianc99
changed the title
3 of 5 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
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.
Not for merging. This branch merges #1 to #7 so they can be measured and reviewed together; merge those PRs instead.
Motivation
Adaptive verification for DFLASH (DFlash2 block-8 draft) on Qwen3.8-27B, built on upstream DSpark instead of a separate implementation.
Modifications
Merge of, in order:
No other code.
Run with
SGLANG_RAGGED_VERIFY_MODE=compact --enable-linear-replayssm-spec --speculative-dspark-sps-table-path <table>; the table comes from the existing offline profiler:python -m sglang.benchmark.dspark_sps_profiler all --base-url ... --fracs 0.25 0.5 0.75 1.0 --max-batch-size <top concurrency> --input-len 1024(128 at TP1, 512 at TP4 here) against a server started withSGLANG_RAGGED_VERIFY_MODE=compact SGLANG_DSPARK_ENABLE_SPS_RECORD=1 SGLANG_SIMULATE_ACC_LEN=1.0.Speed Tests
Setup: Qwen3.8-27B target, DFlash2 block-8 draft, GB300, bf16 mamba state,
trtllm_mhatarget attention andfa4draft attention, overlap scheduling on. SPEED-Bench first turns, up to 4096 output tokens; T=1 is top_p 0.95, top_k 20, T=0 is greedy. TP1 uses--max-running-requests 136; TP4 uses--tp-size 4 --max-running-requests 528 --cuda-graph-max-bs-decode 512 --mem-fraction-static 0.80. The running cap sits above the top concurrency: every level reaches c running requests (c-1 in the sampled maximum at a few TP4 levels), and the mean queue stays under 0.7% of c (1.9-3.4 requests at TP4 c512, while new requests wait for prefill). The client keeps c requests in flight; requests per level: 16 (c1, c2), 24 (c4), 48 (c8), 96 (c16), 192 (c32), 384 (c64), 640 (c128), 512 (c256), 1024 (c512).Columns:
SGLANG_RAGGED_VERIFY_MODE=compact --enable-linear-replayssm-specwith the profiled cost table.SGLANG_RAGGED_VERIFY_MODE=static --enable-linear-replayssm-spec; every request verifies the full block.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. Each cell is the mean over runs.
Every adaptive and static run used these caps. The upstream column's top levels (TP1 c64-c128, TP4 c128-c512) come from runs with the same caps; its lower levels come from runs with a cap of 128 / 512, which only matters at the top level. The single-PR measurements in #1 to #7 used a cap equal to the top concurrency (128 / 512). At TP1 c128 that keeps about one request queued and nearly doubles TTFT (about 0.5 s to 1.0 s); at TP4 c512 it adds about one queued request and about 10% TTFT. Both read 7-15% higher decode tok/s for every method, so compare their top levels only within each PR.
TP1, T=1, decode tok/s
Runs per cell: adaptive 3, static 3, upstream 6 below c64 and 1 from c64 up.
TP1, T=0, decode tok/s
Runs per cell: adaptive 4, static 4, upstream 6 below c64 and 1 from c64 up.
TP4, T=1, decode tok/s
Runs per cell: adaptive 3, static 2, upstream 2 below c128 and 1 from c128 up.
TP4, T=0, decode tok/s
Runs per cell: adaptive 3, static 2, upstream 2 below c128 and 1 from c128 up.
Run-to-run spread (max minus min, over the mean) is up to 7.3% below c32 and at most 4.1% from c32 up.
All adaptive and static numbers were measured on this branch's head, which carries every review fix.
How to read the gains:
Speedup over no speculation (T=1)
No speculation: the same tree and the same server flags with the speculative flags removed, same client and requests per level.
How the rows were measured:
--max-running-requests 136).--stream-interval 4and--max-running-requests 528 --cuda-graph-max-bs-decode 512.--stream-interval 4: with one streamed chunk per token, no-speculation serving at these levels is limited by the HTTP and detokenizer path, not the GPU. At TP4 c512, time to first token reached 24 s and only about 30% of the c requests were running, so decode tok/s overstated it. With--stream-interval 4every column keeps 78-90% of c requests running at every one of these levels. Time to first token stays under 1 s at TP4 and under 3 s at TP1. The speculative numbers change by at most 3% between the two settings.TP1, T=1, decode tok/s
TP4, T=1, decode tok/s
Static speculation falls below no speculation at the top levels: 0.85x at TP1 c512 and 0.86x at TP4 c512. Adaptive verification keeps decoding at least on par with no speculation at every level, from 3.04x / 2.63x at c1 (TP1 / TP4) down to 1.00x / 0.98x at c512. Speculation can therefore stay on at any load instead of being switched off at high concurrency.
Known gaps
Tests (on this tree)
test_gdn_replayssm_spec_foldcases that fail the same way on av/base.Related upstream work
Each PR states its overlap with open upstream PRs. In short:
Do not merge this branch: merge #1 to #7 instead. It carries the conv tile patches twice (#1 and their copies in #6).