[Woody-Hu]

Motivation

hybrid_linear_attn_backend._track_mamba_state_extend restores a radix checkpoint whose depth is not on the FLA chunk grid (64 tokens) from the per-chunk h tensor. h is allocated with the activation dtype and is "already rounded to the activation dtype" — bf16 for a BF16 checkpoint. Snapshot-capable kernels avoid that rounding by writing the chunk-boundary state in fp32 into track_state, and gdn_backend allocates that buffer only when the extend kernel declares supports_track_state_snapshot.

All KDA kernels declare it. GDN's Triton kernel declared supports_packed_decode and supports_strided_target_verify_qkv but not this flag, and fla/chunk.py dropped the two arguments on the way to chunk_gated_delta_rule_fwd_h — which already accepts them and asserts fp32. So for GDN with an fp32 state pool, a hit whose cached depth was not a multiple of 64 resumed from a bf16-rounded state.

Measured on A100-SXM4-80GB, BF16 Qwen/Qwen3.5-4B, TP=1, --attention-backend triton --enable-deterministic-inference, prompt 2934 tokens hitting at 2880 (2934 % 64 = 54):

| branch | max|Δlogp| hit vs flushed cold, before | after | output ids |
|---|---:|---:|---|
| A | 3.756e-2 | 0.0 | 32/32 equal |
| B | 1.802e-1 | 0.0 | 32/32 equal |

Controls: aligned donation (prompt 1408 = 22×64) was already bit-exact before the patch, and cold-vs-cold is 0.0 with and without it, so the change is confined to the hit path. Also bit-exact at 8k and 16k prefixes, TP=2, and with Qwen3.5-9B.

Change

• python/sglang/kernels/ops/attention/fla/chunk.py: accept track_state / track_chunk_idx in chunk_gated_delta_rule and its autograd.Function, and pass them to chunk_gated_delta_rule_fwd_h.
• python/sglang/srt/layers/attention/linear/kernels/gdn_triton.py: declare supports_track_state_snapshot, and forward the two kwargs from extend.

Both default to None, so no other caller changes and the untracked path is untouched (covered by a test).

Tests

test/registered/attention/unittests/gdn/test_gdn_track_state_snapshot.py checks that the capability flag is set, that the buffer holds the fp32 chunk-boundary state bit-for-bit (against the state an aligned 128-token run writes back), and that extend without track_state still matches the plain chunk kernel.

On the parent commit the new test fails with TypeError: chunk_gated_delta_rule() got an unexpected keyword argument 'track_state' and with the capability assertion unmet.

Scope

• Requires an fp32 state pool (the Qwen3.5-4B/9B default mamba_ssm_dtype=float32). With --mamba-ssm-dtype bfloat16 both paths round once to bf16 and this divergence is masked.
• Verified on A100/SM80 only. I could not test SM120, NVFP4, TP>2, or release images.
• Related to, but does not explain, [Bug] Potential hybrid GDN Radix-cache selected-logprob drift on repeated branch scoring #41351: on that stack the first divergence reported in the thread is the FP8 bmm_fp8 batch-shape dependency. The state-precision path is orthogonal, and this PR does not claim to fix that report.
• Not measured: step-time cost of the extra buffer (one fp32 (num_tracked_seqs, *state_shape) allocation per tracked prefill layer-forward, the same allocation kda_backend has always made).

Question for maintainers

If you would rather eliminate unaligned checkpoints altogether (keep chunked-prefill cuts on the mamba chunk grid) than make the unaligned restore exact, please say so and I will close this.

---

CI States

Latest PR Test (Base): ❌ Run #40616736653
Latest PR Test (Extra): ❌ Run #40616736503
Latest PR Test (AMD ROCm 10): ❌ Run #40616736604

---

Re-validated after rebasing onto latest main (54be2010d3): A100-SXM4-80GB, BF16 Qwen3.5-4B, TP=1, n_prompt=2059 with cached_tokens=2048 on device — max|Δlogp| hit vs cold = 0.0 for both branches with the patch; the same probe against an unpatched build gives 5.85e-2 / 6.03e-2 (top-1 unchanged). Unit tests: Ran 3 tests in 13.008s — OK.