From 7705f6c3ff0de1031b499856fdc8c4f15b211f27 Mon Sep 17 00:00:00 2001 From: Hu Di <812791840@qq.com> Date: Wed, 28 Oct 2026 09:42:16 +0800 Subject: [PATCH] [Fix] gdn: Honour the fp32 radix track snapshot in Triton extend When a donated prefill checkpoint is not on the 64-token FLA chunk grid, the radix track path restores the recurrent state from ``h``, which is allocated with the activation dtype (bf16) and is already rounded, while the aligned path copies the fp32 ``last_recurrent_state``. Forward ``track_state`` / ``track_chunk_idx`` through ``chunk_gated_delta_rule`` and declare ``supports_track_state_snapshot`` on ``TritonGDNKernel`` so the chunk kernel writes the fp32 chunk-boundary state into the buffer the backend already allocates. --- .../sglang/kernels/ops/attention/fla/chunk.py | 12 ++ .../attention/linear/kernels/gdn_triton.py | 16 ++ .../gdn/test_gdn_track_state_snapshot.py | 199 ++++++++++++++++++ 3 files changed, 227 insertions(+) create mode 100644 test/registered/attention/unittests/gdn/test_gdn_track_state_snapshot.py diff --git a/python/sglang/kernels/ops/attention/fla/chunk.py b/python/sglang/kernels/ops/attention/fla/chunk.py index 918d3bcd8d8e..662f2fb405e8 100644 --- a/python/sglang/kernels/ops/attention/fla/chunk.py +++ b/python/sglang/kernels/ops/attention/fla/chunk.py @@ -47,6 +47,8 @@ def chunk_gated_delta_rule_fwd( inplace_update: bool = True, chunk_offsets: torch.LongTensor | None = None, o_chunk_indices: torch.LongTensor | None = None, + track_state: Optional[torch.Tensor] = None, + track_chunk_idx: Optional[torch.Tensor] = None, ): g = chunk_local_cumsum( g, chunk_size=CHUNK_SIZE, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices @@ -73,6 +75,8 @@ def chunk_gated_delta_rule_fwd( chunk_indices=chunk_indices, inplace_update=inplace_update, chunk_offsets=chunk_offsets, + track_state=track_state, + track_chunk_idx=track_chunk_idx, ) o = chunk_fwd_o( q=q, @@ -109,6 +113,8 @@ def forward( inplace_update: bool = True, chunk_indices: Optional[torch.LongTensor] = None, chunk_offsets: Optional[torch.LongTensor] = None, + track_state: Optional[torch.Tensor] = None, + track_chunk_idx: Optional[torch.Tensor] = None, ): q_orig = q k_orig = k @@ -136,6 +142,8 @@ def forward( inplace_update=inplace_update, chunk_offsets=chunk_offsets, o_chunk_indices=o_chunk_indices, + track_state=track_state, + track_chunk_idx=track_chunk_idx, ) return o.to(q.dtype), h @@ -156,6 +164,8 @@ def chunk_gated_delta_rule( inplace_update: bool = True, chunk_indices: Optional[torch.LongTensor] = None, chunk_offsets: Optional[torch.LongTensor] = None, + track_state: Optional[torch.Tensor] = None, + track_chunk_idx: Optional[torch.Tensor] = None, ): r""" Args: @@ -275,6 +285,8 @@ def chunk_gated_delta_rule( inplace_update, chunk_indices, chunk_offsets, + track_state, + track_chunk_idx, ) if head_first: o = rearrange(o, "b t h ... -> b h t ...") diff --git a/python/sglang/srt/layers/attention/linear/kernels/gdn_triton.py b/python/sglang/srt/layers/attention/linear/kernels/gdn_triton.py index 433e2704b442..2789680915c0 100644 --- a/python/sglang/srt/layers/attention/linear/kernels/gdn_triton.py +++ b/python/sglang/srt/layers/attention/linear/kernels/gdn_triton.py @@ -42,6 +42,10 @@ class TritonGDNKernel(LinearAttnKernelBase): supports_packed_decode: bool = not is_cpu() and not is_npu() supports_strided_target_verify_qkv: bool = True + # ``chunk_gated_delta_rule`` forwards ``track_state`` to the fused chunk + # kernel, so the radix track path can snapshot the fp32 chunk-boundary + # state instead of rounding it through the bf16 ``h`` output. + supports_track_state_snapshot: bool = not is_cpu() and not is_npu() and not is_xpu() def packed_decode( self, @@ -207,6 +211,17 @@ def extend( # control. Its existing behavior is equivalent to True. inplace_update_args = {} + # The radix track path asks for an fp32 snapshot of the state at a + # chunk boundary that is not the request's last position. Forward it so + # the chunk kernel writes the buffer in fp32 instead of snapping the + # bf16 ``h`` output (which silently rounds the cached state). + track_kwargs = {} + if kwargs.get("track_state") is not None: + track_kwargs = { + "track_state": kwargs["track_state"], + "track_chunk_idx": kwargs.get("track_chunk_idx"), + } + return chunk_gated_delta_rule( q=q, k=k, @@ -220,6 +235,7 @@ def extend( **recurrent_state_indices_args, **inplace_update_args, **graph_table_args, + **track_kwargs, ) def target_verify( diff --git a/test/registered/attention/unittests/gdn/test_gdn_track_state_snapshot.py b/test/registered/attention/unittests/gdn/test_gdn_track_state_snapshot.py new file mode 100644 index 000000000000..55ed259eddc6 --- /dev/null +++ b/test/registered/attention/unittests/gdn/test_gdn_track_state_snapshot.py @@ -0,0 +1,199 @@ +"""The radix track path must snapshot the prefill state in fp32, not in bf16. + +`hybrid_linear_attn_backend._track_mamba_state_extend` restores checkpoints whose +depth is not on the FLA chunk grid from the per-chunk ``h`` tensor, which is +allocated with the activation dtype (bf16) and is "already rounded to the +activation dtype". Snapshot-capable kernels avoid that by writing the fp32 +chunk-boundary state into ``track_state``; the backend only allocates the buffer +when the extend kernel declares ``supports_track_state_snapshot``. + +GDN's Triton kernel used to declare ``supports_packed_decode`` and +``supports_strided_target_verify_qkv`` but not ``supports_track_state_snapshot``, +and ``fla/chunk.py`` dropped the arguments on the way to +``chunk_gated_delta_rule_fwd_h``. These tests pin both halves: the declaration +and the fp32 bits actually landing in the buffer. +""" + +import unittest + +import torch + +from sglang.kernels.ops.attention.fla.chunk import chunk_gated_delta_rule +from sglang.srt.layers.attention.linear.kernels.gdn_triton import TritonGDNKernel +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-large") + +# FLA chunk size: the radix track grid snaps donated checkpoints to this step. +CHUNK_SIZE = 64 + + +@unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA") +class TestGdnTrackStateSnapshot(CustomTestCase): + NUM_HEADS = 2 + HEAD_K_DIM = 64 + HEAD_V_DIM = 64 + # 3 chunks in total, checkpoint taken at the end of the second one. + TOKENS = 3 * CHUNK_SIZE + TRACKED_TOKENS = 2 * CHUNK_SIZE + TRACK_CHUNK_IDX = TRACKED_TOKENS // CHUNK_SIZE + + def _inputs(self, tokens): + torch.manual_seed(0) + device, dtype = "cuda", torch.bfloat16 + q = torch.randn( + 1, tokens, self.NUM_HEADS, self.HEAD_K_DIM, dtype=dtype, device=device + ) + k = torch.randn( + 1, tokens, self.NUM_HEADS, self.HEAD_K_DIM, dtype=dtype, device=device + ) + v = torch.randn( + 1, tokens, self.NUM_HEADS, self.HEAD_V_DIM, dtype=dtype, device=device + ) + g = torch.nn.functional.logsigmoid( + torch.randn(1, tokens, self.NUM_HEADS, dtype=dtype, device=device) + ) + beta = torch.sigmoid( + torch.randn(1, tokens, self.NUM_HEADS, dtype=dtype, device=device) + ) + return q, k, v, g, beta + + def _pool(self): + torch.manual_seed(1) + # Non-zero initial state: the checkpoint must carry it through. + return ( + torch.randn( + 2, + self.NUM_HEADS, + self.HEAD_V_DIM, + self.HEAD_K_DIM, + dtype=torch.float32, + device="cuda", + ) + * 0.1 + ) + + def _cu_seqlens(self, tokens): + return torch.tensor([0, tokens], dtype=torch.long, device="cuda") + + def _reference(self): + """fp32 state after TRACKED_TOKENS tokens, written back by inplace_update.""" + pool = self._pool() + q, k, v, g, beta = self._inputs(self.TOKENS) + indices = torch.tensor([1], dtype=torch.int32, device="cuda") + cut = self.TRACKED_TOKENS + chunk_gated_delta_rule( + q=q[:, :cut], + k=k[:, :cut], + v=v[:, :cut], + g=g[:, :cut], + beta=beta[:, :cut], + initial_state=pool, + initial_state_indices=indices, + cu_seqlens=self._cu_seqlens(cut), + head_first=False, + use_qk_l2norm_in_kernel=True, + ) + return pool[indices[0]] + + def _tracked_buf(self, extend): + pool = self._pool() + q, k, v, g, beta = self._inputs(self.TOKENS) + indices = torch.tensor([1], dtype=torch.int32, device="cuda") + buf = torch.zeros( + 1, + self.NUM_HEADS, + self.HEAD_V_DIM, + self.HEAD_K_DIM, + dtype=torch.float32, + device="cuda", + ) + extend( + q=q, + k=k, + v=v, + g=g, + beta=beta, + initial_state=pool, + initial_state_indices=indices, + cu_seqlens=self._cu_seqlens(self.TOKENS), + head_first=False, + use_qk_l2norm_in_kernel=True, + track_state=buf, + track_chunk_idx=torch.tensor( + [self.TRACK_CHUNK_IDX], dtype=torch.int32, device="cuda" + ), + ) + return buf[0] + + def test_chunk_kernel_writes_fp32_checkpoint(self): + expected = self._reference() + actual = self._tracked_buf(chunk_gated_delta_rule) + self.assertEqual(actual.dtype, torch.float32) + self.assertTrue( + torch.equal(actual, expected), + "track_state must hold the fp32 chunk-boundary state bit-for-bit; " + f"max|diff|={(actual - expected).abs().max().item():.3e}", + ) + + def test_triton_extend_declares_track_state_snapshot(self): + self.assertTrue( + TritonGDNKernel.supports_track_state_snapshot, + "TritonGDNKernel.extend forwards track_state, so it must opt into " + "the fp32 snapshot buffer; otherwise the backend keeps falling back " + "to the bf16 h copy for unaligned checkpoints.", + ) + + expected = self._reference() + + def extend_via_kernel(**kwargs): + kernel = TritonGDNKernel() + return kernel.extend( + kwargs.pop("q"), + kwargs.pop("k"), + kwargs.pop("v"), + kwargs.pop("g"), + kwargs.pop("beta"), + ssm_states=kwargs.pop("initial_state"), + cache_indices=kwargs.pop("initial_state_indices"), + query_start_loc=kwargs.pop("cu_seqlens"), + **kwargs, + ) + + actual = self._tracked_buf(extend_via_kernel) + self.assertTrue(torch.equal(actual, expected)) + + def test_extend_without_track_state_is_unchanged(self): + """The optional arguments must not perturb the ordinary extend path.""" + pool = self._pool() + q, k, v, g, beta = self._inputs(self.TOKENS) + indices = torch.tensor([1], dtype=torch.int32, device="cuda") + TritonGDNKernel().extend( + q, + k, + v, + g, + beta, + ssm_states=pool, + cache_indices=indices, + query_start_loc=self._cu_seqlens(self.TOKENS), + ) + reference = self._pool() + chunk_gated_delta_rule( + q=q, + k=k, + v=v, + g=g, + beta=beta, + initial_state=reference, + initial_state_indices=indices, + cu_seqlens=self._cu_seqlens(self.TOKENS), + head_first=False, + use_qk_l2norm_in_kernel=True, + ) + self.assertTrue(torch.equal(pool[indices[0]], reference[indices[0]])) + + +if __name__ == "__main__": + unittest.main()