Repository navigation
Conversation
|
/rerun-test test/registered/models_e2e/test_qwen35_fp4_mtp.py::TestQwen35FP4MTPReplaySSM.test_gsm8k test/registered/attention/unittests/gdn/test_gdn_replayssm_spec_fold.py test/registered/attention/unittests/gdn/test_gdn_cutedsl_ring_verify.py test/registered/attention/unittests/gdn/test_linear_replayssm_decode.py |
|
Results for 🚀 🚀 |
|
Added the Replay SSM AIME results. There is no accuracy regression. |
618b9ab to
150c91c
Compare
…35544-conflict-resolution-20260828 # Conflicts: # python/sglang/srt/server_args.py
could we run more pass? maybe repeat 8 or repeat 16. Does the pass@4 mean average of 4 runs? It seems the pass@4 is higher but total is lower? |
| k_raw_tile = k_tile | ||
| q_tile = (q_tile * (q_rnorm * scale)[:, None]).to(h0.dtype.element_ty) | ||
| k_tile = (k_tile * k_rnorm[:, None]).to(h0.dtype.element_ty) | ||
| q_tile = (q_tile * (q_rnorm * scale)[:, None]).to(q.dtype.element_ty) |
There was a problem hiding this comment.
we downgrade the precision of q_tile/k_tile, b_d_scaled, and khist_tile, right? Is this expected?
There was a problem hiding this comment.
Good catch. The FP32-checkpoint path was unintentionally lowering these operands; 9168a3f7dd restores its FP32/TF32 path while leaving the BF16 hot path unchanged. The BF16 and FP32 64-window GPU comparisons both pass.
| elif enable_linear_replayssm_spec and ssm_dtype != torch.float32: | ||
| # Low parts of compact D and normalized K. A BF16 checkpoint | ||
| # materializes these once per accepted verify window. | ||
| replayssm_rawv = torch.zeros( |
There was a problem hiding this comment.
why do we hold low parts for ssm_dtype==bf16 but not for ssm_dtype==fp32?
There was a problem hiding this comment.
The low parts should follow the compact ring dtype, not the checkpoint dtype, because spec-verify replayssm_d/k use the activation/conv dtype. 9168a3f7dd changes allocation and accounting to ring_dtype != torch.float32, including the FP32-checkpoint + BF16-ring case.
…35544-conflict-resolution-20260828 # Conflicts: # python/sglang/srt/mem_cache/mamba_radix_cache.py # python/sglang/srt/mem_cache/memory_pool.py
|
|
|
/tag-and-rerun-ci |
…ernel strided q/k/v views Upstream 9fdb717 (main only, not in v0.5.20). On the 27B's D group every DFLASH verify step ran fused_qkv_split_gdn_prefill on every GDN layer, one copy kernel per layer per step, although the Triton verify kernel (fused_sigmoid_gating_delta_rule_update) reads q/k/v with their token strides (stride_q = q.stride()[1], runtime int, dense head axis). Target verify now passes torch.split views of the post-conv mixed_qkv when the chosen verify kernel opts in; prefill keeps the fused split (FLA chunk kernels want dense). Ported: LinearAttnKernelBase.supports_strided_target_verify_qkv (False), TritonGDNKernel opt-in, GDNKernelDispatcher._get_target_verify_kernel + target_verify_supports_strided_qkv, the routing in forward_extend. Adapted: the fork has no ReplaySSM spec-fold / circular verify routes (sgl-project#28695/sgl-project#35544 not on this line), so the routing is the dispatcher capability alone, decided inline per forward instead of in metadata init. The Triton opt-in is limited to the CUDA/HIP Triton kernel (NPU/CPU/XPU substitute other implementations in gdn_triton.py). Tree verify (topk>1) routes to the Triton tree kernel, which also opts in; FlashInfer/CuTe kernels inherit False. Tests (CPU, incg cgroup): test/registered/unit/layers/attention/test_gdn_verify_strided_qkv_33778.py - dispatcher capability (linear + tree Triton opt in, attribute-less kernel does not); - GDNAttnBackend.forward_extend TARGET_VERIFY: the verify kernel receives views sharing mixed_qkv's storage (token stride = qkv width), the fused split is never called; EXTEND still calls it; - the production verify kernel under TRITON_INTERPRET: strided views vs dense copies (ratio 3, 2 requests x 4 draft tokens, intermediate states cached, state update disabled) -> output and intermediate states bit-identical. 4 passed. On the base sources the capability and verify-routing tests fail.
…ned to the recurrent verify 27B ReplaySSM package, slice S1 of REPLAYSSM_PLAN.md (D-KV: the ring replaces the per-draft "speculative intermediate state"). No wiring, no flag yet -- nothing on any serving path changes. Ported verbatim from upstream main (chain sgl-project#28695 -> sgl-project#32692 -> sgl-project#36970 -> sgl-project#35544): sglang/kernels/ops/attention/fla/gdn_replayssm_spec_decode.py -> this line's srt/layers/attention/fla/ (only the fla.utils import moved): circular verify kernel (output from checkpoint + compact d/k/g history), compact commit (materializes accepted windows, hi/lo low parts for 16-bit checkpoints, track rows snapshot the crossing state), cursor advance/reset. configs/mamba_utils.py: BaseLinearStateParams.replayssm_ring_bytes_per_req (upstream helper; the shape is the rank's own, so uneven GDN TP prices every rank correctly) -- the one formula the mixin and the planner will share (S5). Why the pin: upstream runs the GDN route only behind the EAGLE/MTP commit and refuses DFLASH for non-KDA models (kv_cache_configurator), so GDN + DFLASH (the 27B) never ran upstream. The test drives both production kernels under TRITON_INTERPRET against this line's DFLASH verify kernel (fused_sigmoid_gating_delta_rule_update with per-step intermediate states + the scatter of the last accepted step). Test (CPU, incg cgroup): test/registered/unit/layers/attention/test_gdn_replayssm_spec_equivalence_s1.py - fp32 checkpoint: verify rel 2.6e-7, track snapshot 1.6e-7, second verify from checkpoint + ring history 2.8e-7 vs the recurrent path; cursors [3,0,1]; - 16-bit (fp16 standing in for bf16): verify 4.2e-4 vs recurrent 2.6e-4 of the fp32 truth, committed checkpoint 3.7e-4 = the recurrent 16-bit rounding, track 3.0e-4; cursors back to 0 (fold every commit); other slots bit-equal; - ring price of a 27B rank row (HV 18 / H 6, bf16, L 16): 9.5 MB vs 226 MB of intermediate state for 8 draft steps. 3 passed. Mutant (cursor advance one short): both equivalence tests red. bf16 itself is a metal gate (plan): acceptance, needle, greedy A/B. 27B line (desk/27b-up-replayssm-line-0924): cherry-picked from 420c547 (desk/27b-up-replayssm-0924, on the NF-based staging line fa757eb) onto e15e760, the 27B line (xsn411 829ebd0 + upstream ports + P without draft compute + record filter). Nothing to adapt: the kernel file is new, and configs/mamba_utils.py, fla/ and the DFLASH verify kernel it is pinned against are identical on both lines. Tests on this line (CPU, incg cgroup, S1 tree): test_gdn_replayssm_spec_equivalence_s1.py 3 passed. Interpreter vs the recurrent DFLASH verify: fp32 verify rel 2.63e-7, track snapshot 1.57e-7, second verify from checkpoint + ring history 2.79e-7 (all < 4e-7), cursors [3,0,1]; 16-bit verify 4.18e-4 (recurrent 2.65e-4 of the fp32 truth), committed checkpoint 3.71e-4 == recurrent 3.71e-4, track 2.96e-4, cursors [0,0,0], other slots bit-equal -- the values the NF line reported. (cherry picked from commit 9431003)
Motivation
ReplaySSM speculative decoding currently materializes the full GDN checkpoint state on every target step, even though most speculative iterations only append a small amount of history. Deferring that work removes repeated state materialization from the decode critical path.
Modifications
The review scope is limited to the ReplaySSM kernel, the request/state ownership and accounting needed by that kernel, and one existing accounting test. The larger experimental compatibility framework and additional defensive abstractions are not part of this PR.
Accuracy Tests
0.857 GB/GPUversus0.849 GB/GPUbefore the PR. CUDA-graph startup available memory is36.129-36.131 GBversus36.143-36.145 GB; the earlier approximately3.6 GB/GPUheadroom regression is removed.AIME25 128K x4
The comparison uses Qwen3.5 on one 4-GPU node, aggregate TP4/DP1/PP1/EP1 serving, MTP, ReplaySSM, BF16 SSM state, 128K context, four repeats, and the same 30 problems.
The current run produced all
120/120predictions with zero harness errors and3,019,855completion tokens.AIME25 128K x8
For additional confidence, the old and current ReplaySSM implementations were rerun with the same serving/evaluation contract and eight repeats.
pass@Nis the fraction of the 30 problems with at least one correct answer amongNattempts; it is not the average accuracy across repeats.Both arms produced all
240/240predictions with zero harness errors and passed all 30 problems at pass@8. The current implementation is+1/240in total correct attempts (92.08%versus91.67%).Logits and long-sequence stability
248).99.92%, mean top-5 Jensen-Shannon divergence is3.82e-4, and mean absolute chosen-logprob delta is7.42e-3.100%top-1 agreement and a7.61e-3mean absolute chosen-logprob delta over its common prefix, placing the old-versus-current delta at the same scale as run-to-run variation.0.0194to0.0105and entropy decreases from0.0706to0.0464; there is no evidence of error or variance accumulating with generation length.Speed Tests and Profiling
extra_buffer, radix-cache page size 64, and the same frozen 112-request cohort in both arms.[28, 28, 28, 28], achieved batch 28 on every rank, CUDA graph enabled, queue depth 0, no replenishment, and measurement ending before the first completion.4.528%versus the baseline.-1.254%and mean context by-1.740%; the fixed-AL metric normalizes acceptance to 4.8, while observed throughput remains+3.427%.Output tokens/s/GPUis an engine-side pure-generation metric from the fixed-B28 target-step window, not a service-level AIPerf total-throughput metric.Checklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ciCI States
Latest PR Test (Base): ✅ Run #33840060530
Latest PR Test (Extra): ❌ Run #33840060375
Latest PR Test (AMD ROCm 7.2): ⏳ Run #33840060550