From fddfc1fb5ef58fdb746616103aea01ea482370a6 Mon Sep 17 00:00:00 2001 From: YAMY <74099316+YAMY1234@users.noreply.github.com> Date: Wed, 29 Jul 2026 09:47:35 -0700 Subject: [PATCH] [GDN] Support FlashInfer GDN prefill with extra-buffer radix cache (#29735) --- .../advanced_features/attention_backend.mdx | 2 +- .../attention/hybrid_linear_attn_backend.py | 8 +- .../layers/attention/linear/gdn_backend.py | 25 +++- .../linear/kernels/gdn_flashinfer.py | 120 +++++++++++++----- .../linear/kernels/kernel_backend.py | 2 + .../layers/attention/mamba/mamba2_metadata.py | 3 + .../attention_methods/gdn_attention.py | 19 ++- .../unittests/gdn/test_flashinfer.py | 75 +++++++++++ .../test_gdn_prefill_backend_policy.py | 102 ++++++++++----- 9 files changed, 282 insertions(+), 74 deletions(-) diff --git a/docs_new/docs/advanced_features/attention_backend.mdx b/docs_new/docs/advanced_features/attention_backend.mdx index eb1afedb9..ef7fcb6f3 100644 --- a/docs_new/docs/advanced_features/attention_backend.mdx +++ b/docs_new/docs/advanced_features/attention_backend.mdx @@ -338,7 +338,7 @@ GDN (Gated Delta Network) is a linear attention mechanism with O(n) complexity, The GDN linear attention layers have their own kernel backends, selected via `--linear-attn-backend` (default: `triton`). You can override the kernel per phase with `--linear-attn-decode-backend` and `--linear-attn-prefill-backend`. -On SM100/SM103 with CUDA 13+, SGLang automatically selects FlashInfer for GDN prefill when the per-phase override is unset, the base linear-attention backend is Triton, recurrent state is BF16, key/value head dimensions are 128, dynamic chunking and page-major KV layout are disabled, and `--chunked-prefill-size` is between 1 and 8192. Radix caching may be disabled or use the `no_buffer` strategy; extra-buffer strategies require state checkpoint support. +On SM100/SM103 with CUDA 13+, SGLang automatically selects FlashInfer for GDN prefill when the per-phase override is unset, the base linear-attention backend is Triton, recurrent state is BF16, key/value head dimensions are 128, dynamic chunking and page-major KV layout are disabled, and `--chunked-prefill-size` is between 1 and 8192. Radix caching may be disabled or use `no_buffer`, `extra_buffer`, or `extra_buffer_lazy`; the extra-buffer paths use state checkpoints. diff --git a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py index 61a45f42e..de43cce59 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py @@ -724,16 +724,18 @@ class MambaAttnBackendBase(AttentionBackend): def _track_mamba_state_extend( self, forward_batch: ForwardBatch, - h: torch.Tensor, + h: Optional[torch.Tensor], ssm_states: torch.Tensor, forward_metadata: ForwardMetadata, ): """Copy extend SSM state at the last chunk boundary to track slots (source depends on chunk alignment; see `_init_track_ssm_indices`).""" if forward_metadata.has_mamba_track_mask: - h = h.squeeze(0) - + # Triton always returns h; FlashInfer returns it only when checkpoints + # were requested. Aligned-only tracking reads the final state below. if forward_metadata.track_ssm_h_src.numel() > 0: + assert h is not None + h = h.squeeze(0) ssm_states[forward_metadata.track_ssm_h_dst] = h[ forward_metadata.track_ssm_h_src ].to(ssm_states.dtype, copy=False) diff --git a/python/sglang/srt/layers/attention/linear/gdn_backend.py b/python/sglang/srt/layers/attention/linear/gdn_backend.py index 61943b885..dde215538 100644 --- a/python/sglang/srt/layers/attention/linear/gdn_backend.py +++ b/python/sglang/srt/layers/attention/linear/gdn_backend.py @@ -75,10 +75,6 @@ def maybe_set_default_flashinfer_gdn_prefill(model_runner: ModelRunner) -> None: ): return - # Extra-buffer strategies need intermediate state checkpoints. - if args.uses_mamba_radix_cache and args.mamba_radix_cache_strategy != "no_buffer": - return - cuda_version = torch.version.cuda chunk_size = args.chunked_prefill_size config = hybrid_gdn_config(model_runner.model_config) @@ -204,6 +200,10 @@ class GDNKernelDispatcher: f"packed_decode={self.supports_packed_decode}" ) + @property + def extend_uses_state_checkpoints(self) -> bool: + return self.extend_kernel.uses_state_checkpoints + def packed_decode( self, mixed_qkv: torch.Tensor, @@ -362,6 +362,14 @@ class GDNAttnBackend(MambaAttnBackendBase): self.forward_metadata.mamba_track_mask_indices ] ) + if self.kernel_dispatcher.extend_uses_state_checkpoints: + from sglang.srt.layers.attention.linear.kernels.gdn_flashinfer import ( + maybe_build_flashinfer_checkpoint_plan, + ) + + maybe_build_flashinfer_checkpoint_plan( + forward_batch, self.forward_metadata, self.device + ) def forward_decode( self, @@ -649,6 +657,13 @@ class GDNAttnBackend(MambaAttnBackendBase): ssm_states=ssm_states_contig, cache_indices=state_cache_indices, query_start_loc=query_start_loc, + state_checkpoint_cu_starts=( + forward_metadata.state_checkpoint_cu_starts + ), + num_state_checkpoints=forward_metadata.num_state_checkpoints, + state_checkpoint_every_n_tokens=( + forward_metadata.state_checkpoint_every_n_tokens + ), ) if is_npu() and last_recurrent_state is not None: @@ -663,7 +678,7 @@ class GDNAttnBackend(MambaAttnBackendBase): conv_states[cache_indices] = conv_states_contig ssm_states[cache_indices] = ssm_states_contig - if h is not None: + if forward_metadata.has_mamba_track_mask: self._track_mamba_state_extend( forward_batch, h, ssm_states, forward_metadata ) diff --git a/python/sglang/srt/layers/attention/linear/kernels/gdn_flashinfer.py b/python/sglang/srt/layers/attention/linear/kernels/gdn_flashinfer.py index 7b7a89801..f7f4d974d 100644 --- a/python/sglang/srt/layers/attention/linear/kernels/gdn_flashinfer.py +++ b/python/sglang/srt/layers/attention/linear/kernels/gdn_flashinfer.py @@ -5,20 +5,27 @@ Both SM90 and SM100 use the same pool layout: [pool, HV, V, K] (K-last). SM90 (Hopper): full support — decode, prefill, MTP. State dtype: fp32. SM100 (Blackwell): full support — decode, prefill, MTP. -Requires flashinfer >= 0.6.7. +Requires flashinfer >= 0.6.14. """ +from __future__ import annotations + import logging import os -from typing import Optional +from typing import TYPE_CHECKING, Optional import torch from sglang.srt.layers.attention.linear.kernels.kernel_backend import ( LinearAttnKernelBase, ) +from sglang.srt.runtime_context import get_server_args from sglang.srt.utils import is_cuda +if TYPE_CHECKING: + from sglang.srt.layers.attention.mamba.mamba2_metadata import ForwardMetadata + from sglang.srt.model_executor.forward_batch_info import ForwardBatch + logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- @@ -31,6 +38,47 @@ _flashinfer_gated_delta_rule_decode = None _flashinfer_gated_delta_rule_mtp_bf16 = None +def maybe_build_flashinfer_checkpoint_plan( + forward_batch: ForwardBatch, + forward_metadata: ForwardMetadata, + device: str, +) -> None: + """Populate packed FlashInfer checkpoint metadata when tracking requires it.""" + if ( + forward_metadata.track_ssm_h_src is None + or forward_metadata.track_ssm_h_src.numel() == 0 + ): + return + + checkpoint_every_n_tokens = get_server_args().mamba_cache_chunk_size + extend_seq_lens = forward_batch.extend_seq_lens.to(device="cpu", dtype=torch.int64) + track_mask = forward_batch.mamba_track_mask.to(device="cpu", dtype=torch.bool) + relative_track_lens = forward_batch.mamba_track_seqlens.to( + device="cpu", dtype=torch.int64 + ) - forward_batch.extend_prefix_lens.to(device="cpu", dtype=torch.int64) + + checkpoint_counts = extend_seq_lens // checkpoint_every_n_tokens + checkpoint_cu_starts = torch.zeros(checkpoint_counts.numel() + 1, dtype=torch.int64) + checkpoint_cu_starts[1:] = torch.cumsum(checkpoint_counts, dim=0) + + use_checkpoint = track_mask & (relative_track_lens % checkpoint_every_n_tokens != 0) + track_checkpoint_src = checkpoint_cu_starts[:-1][use_checkpoint] + ( + relative_track_lens[use_checkpoint] // checkpoint_every_n_tokens - 1 + ) + if track_checkpoint_src.numel() and track_checkpoint_src.min() < 0: + raise ValueError("Tracked GDN state precedes the first FlashInfer checkpoint.") + assert track_checkpoint_src.numel() == forward_metadata.track_ssm_h_dst.numel() + + forward_metadata.track_ssm_h_src = track_checkpoint_src.to( + device, non_blocking=True + ) + forward_metadata.state_checkpoint_cu_starts = checkpoint_cu_starts.to( + device, non_blocking=True + ) + forward_metadata.num_state_checkpoints = int(checkpoint_cu_starts[-1]) + forward_metadata.state_checkpoint_every_n_tokens = checkpoint_every_n_tokens + + def _get_flashinfer_gdn_kernels(): """Lazy import for FlashInfer GDN prefill, decode and verify (MTP) kernels. @@ -89,9 +137,11 @@ class FlashInferGDNKernel(LinearAttnKernelBase): SM90 (Hopper): decode uses gather/scatter; prefill and MTP verify supported. SM100 (Blackwell): decode uses gather/scatter; prefill and MTP verify supported. - Requires flashinfer >= 0.6.7. + Requires flashinfer >= 0.6.14. """ + uses_state_checkpoints = True + def __init__(self): ( available, @@ -231,6 +281,9 @@ class FlashInferGDNKernel(LinearAttnKernelBase): ssm_states: torch.Tensor, cache_indices: torch.Tensor, query_start_loc: torch.Tensor, + state_checkpoint_cu_starts: Optional[torch.Tensor] = None, + num_state_checkpoints: int = 0, + state_checkpoint_every_n_tokens: int = 0, **kwargs, ) -> tuple: from sglang.kernels.ops.attention.fla.l2norm import l2norm_fwd @@ -253,23 +306,7 @@ class FlashInferGDNKernel(LinearAttnKernelBase): # slot) so the FlashInfer kernel never reads out-of-bounds state. ssm_cache_indices = cache_indices.clamp(min=0).to(torch.int64) initial_state_fi = ssm_states[ssm_cache_indices].contiguous() - # Pre-allocate bf16 output_state so the kernel compiles and writes the - # bf16 state path directly, avoiding a fp32 allocation and a subsequent - # fp32->bf16 conversion in the scatter step. - output_state_fi = torch.empty_like(initial_state_fi) - output_fi, output_state_fi = self._prefill_fn( - q=q_fi, - k=k_fi, - v=v_fi, - g=alpha_fi, - beta=beta_fi, - scale=None, - initial_state=initial_state_fi, - output_final_state=True, - cu_seqlens=query_start_loc, # already int32 - use_qk_l2norm_in_kernel=False, - output_state=output_state_fi, - ) + cu_seqlens = query_start_loc # already int32 else: # SM90: preserve original negative-index handling (remap to last slot). ssm_cache_indices = torch.where( @@ -279,18 +316,33 @@ class FlashInferGDNKernel(LinearAttnKernelBase): ).to(torch.int64) # State must be float32; kernel requires int64 cu_seqlens. initial_state_fi = ssm_states[ssm_cache_indices].to(torch.float32) - output_fi, output_state_fi = self._prefill_fn( - q=q_fi, - k=k_fi, - v=v_fi, - g=alpha_fi, - beta=beta_fi, - scale=None, - initial_state=initial_state_fi, - output_final_state=True, - cu_seqlens=query_start_loc.to(torch.int64), - use_qk_l2norm_in_kernel=False, + cu_seqlens = query_start_loc.to(torch.int64) + + # Keep final state and checkpoints in the same kernel state dtype. + output_state_fi = torch.empty_like(initial_state_fi) + state_checkpoints = ( + initial_state_fi.new_empty( + (num_state_checkpoints, *initial_state_fi.shape[1:]) ) + if num_state_checkpoints > 0 + else None + ) + output_fi, output_state_fi = self._prefill_fn( + q=q_fi, + k=k_fi, + v=v_fi, + g=alpha_fi, + beta=beta_fi, + scale=None, + initial_state=initial_state_fi, + output_final_state=True, + cu_seqlens=cu_seqlens, + use_qk_l2norm_in_kernel=False, + output_state=output_state_fi, + state_checkpoints=state_checkpoints, + checkpoint_cu_starts=state_checkpoint_cu_starts, + checkpoint_every_n_tokens=state_checkpoint_every_n_tokens, + ) # Write back state to pool ssm_states.index_copy_( @@ -302,9 +354,9 @@ class FlashInferGDNKernel(LinearAttnKernelBase): # Output: [seq, HV, V] -> [1, seq, HV, V] core_attn_out = output_fi.view(1, total_seq_len, num_v_heads, head_v_dim) - # Return (output, last_recurrent_state, h) to match Triton kernel interface. - # h=None since FlashInfer doesn't provide intermediate states. - return core_attn_out, None, None + # Match Triton's [1, checkpoints, H, V, K] intermediate-state layout. + h = state_checkpoints.unsqueeze(0) if state_checkpoints is not None else None + return core_attn_out, None, h # ---- target_verify (MTP) ---- diff --git a/python/sglang/srt/layers/attention/linear/kernels/kernel_backend.py b/python/sglang/srt/layers/attention/linear/kernels/kernel_backend.py index 5e8018699..2a840ec35 100644 --- a/python/sglang/srt/layers/attention/linear/kernels/kernel_backend.py +++ b/python/sglang/srt/layers/attention/linear/kernels/kernel_backend.py @@ -10,6 +10,8 @@ class LinearAttnKernelBase(ABC): and provides decode/extend/target_verify methods with a unified interface. """ + uses_state_checkpoints: bool = False + @abstractmethod def decode( self, diff --git a/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py b/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py index bffe8d100..977b7f7cc 100644 --- a/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py +++ b/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py @@ -55,6 +55,9 @@ class ForwardMetadata: track_ssm_h_dst: Optional[torch.Tensor] = None track_ssm_final_src: Optional[torch.Tensor] = None track_ssm_final_dst: Optional[torch.Tensor] = None + state_checkpoint_cu_starts: Optional[torch.Tensor] = None + num_state_checkpoints: int = 0 + state_checkpoint_every_n_tokens: int = 0 is_target_verify: bool = False draft_token_num: int = 1 diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py index 409a92767..743b4b314 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py @@ -55,6 +55,7 @@ class GDNAttentionCase: page_size: int prefix_lens: tuple[int, ...] extend_lens: tuple[int, ...] = () + linear_attn_prefill_backend: str | None = None @property def batch_size(self) -> int: @@ -245,7 +246,7 @@ class MockGDNModelRunner(ModelRunner): enable_mis=False, linear_attn_backend="triton", linear_attn_decode_backend=None, - linear_attn_prefill_backend=None, + linear_attn_prefill_backend=case.linear_attn_prefill_backend, max_running_requests=None, revision=None, speculative_algorithm=None, @@ -268,10 +269,16 @@ class MockGDNModelRunner(ModelRunner): state_size=head_k_dim, conv_kernel=2, ) + temporal_state_dtype = ( + dtype + if case.linear_attn_prefill_backend == "flashinfer" + and torch.cuda.get_device_capability()[0] >= 10 + else torch.float32 + ) cache_params = Mamba2CacheParams( shape=cache_shape, layers=[0], - dtype=Mamba2StateDType(conv=dtype, temporal=torch.float32), + dtype=Mamba2StateDType(conv=dtype, temporal=temporal_state_dtype), ) self.req_to_token_pool = HybridReqToTokenPool( size=pool_batch_size, @@ -591,6 +598,14 @@ def build_gdn_attention_fixture( initialize_linear_attn_config(runner.server_args) linear_backend = GDNAttnBackend(runner) + if case.linear_attn_prefill_backend == "flashinfer": + from sglang.srt.layers.attention.linear.kernels.gdn_flashinfer import ( + FlashInferGDNKernel, + ) + + testcase.assertIsInstance( + linear_backend.kernel_dispatcher.extend_kernel, FlashInferGDNKernel + ) backend = HybridLinearAttnBackend(full_backend, linear_backend, full_attn_layers=[]) actual_module = ProjectedGDNAttention( num_k_heads=case.num_k_heads, diff --git a/test/registered/attention/unittests/gdn/test_flashinfer.py b/test/registered/attention/unittests/gdn/test_flashinfer.py index 7c8c5c62b..a1f69dcca 100644 --- a/test/registered/attention/unittests/gdn/test_flashinfer.py +++ b/test/registered/attention/unittests/gdn/test_flashinfer.py @@ -10,11 +10,14 @@ from sglang.test.test_utils import CustomTestCase sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +from sglang.srt.layers.attention.linear.kernels.gdn_triton import TritonGDNKernel from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.kits.attention_unittest.attention_methods.gdn_attention import ( GDNAttentionCase, + build_gdn_attention_fixture, make_gdn_cases, run_gdn_attention_case, + run_gdn_fixture_eager, ) from sglang.test.kits.attention_unittest.runner_modes.cuda_graph_decode_runner import ( run_gdn_cuda_graph_decode_case, @@ -30,6 +33,12 @@ from sglang.test.kits.attention_unittest.runner_modes.split_op_runner import ( register_cuda_ci(est_time=20, stage="base-b", runner_config="4-gpu-b200") register_cuda_ci(est_time=20, stage="base-b", runner_config="1-gpu-large") +_cuda_major = int(torch.version.cuda.split(".")[0]) if torch.version.cuda else 0 +_sm_major = torch.cuda.get_device_capability()[0] if torch.cuda.is_available() else 0 +_supports_flashinfer_linear_gdn = _sm_major == 9 or ( + _sm_major == 10 and _cuda_major >= 13 +) + @unittest.skipIf( not torch.cuda.is_available() or not is_flashinfer_available(), @@ -322,5 +331,71 @@ class TestFlashInferGDNBackendCorrectness(CustomTestCase): ) +@unittest.skipUnless( + torch.cuda.is_available() + and is_flashinfer_available() + and _supports_flashinfer_linear_gdn, + "FlashInfer linear GDN requires SM90 or SM100/SM103 with CUDA 13+", +) +class TestFlashInferLinearGDNBackendCorrectness(CustomTestCase): + # FlashInfer's DSL prefill kernels require head size 128 on SM90 and SM100. + HEAD_DIM = 128 + CHECKPOINT_CASE = GDNAttentionCase( + name="flashinfer_gdn_prefill_state_checkpoints", + backend="triton", + linear_attn_prefill_backend="flashinfer", + forward_mode=ForwardMode.EXTEND, + num_k_heads=2, + num_v_heads=4, + page_size=16, + prefix_lens=(0, 64, 128), + extend_lens=(64, 65, 129), + ) + + def test_prefill_tracked_state_checkpoints(self): + fixture = build_gdn_attention_fixture( + self, + self.CHECKPOINT_CASE, + head_k_dim=self.HEAD_DIM, + head_v_dim=self.HEAD_DIM, + max_context_len=320, + runner_batch_size=6, + ) + batch = fixture.forward_batch + # Simulate the tracking metadata produced by the extra-buffer scheduler. + # This test covers checkpoint mapping and state copies, not scheduler setup. + batch.mamba_track_mask = torch.ones(3, dtype=torch.bool, device="cuda") + batch.mamba_track_indices = torch.tensor( + [4, 5, 6], dtype=torch.int64, device="cuda" + ) + batch.mamba_track_seqlens = torch.tensor( + # The final entry selects the second checkpoint at absolute S256. + [64, 129, 257], + dtype=torch.int64, + device="cuda", + ) + + cache = fixture.runner.req_to_token_pool.mamba2_layer_cache(0) + initial_conv = cache.conv[0].clone() + initial_ssm = cache.temporal.clone() + flashinfer_output = run_gdn_fixture_eager(fixture) + flashinfer_tracked = cache.temporal[batch.mamba_track_indices].clone() + + cache.conv[0].copy_(initial_conv) + cache.temporal.copy_(initial_ssm) + fixture.backend.linear_attn_backend.kernel_dispatcher.extend_kernel = ( + TritonGDNKernel() + ) + triton_output = run_gdn_fixture_eager(fixture) + triton_tracked = cache.temporal[batch.mamba_track_indices] + + torch.testing.assert_close( + flashinfer_output, triton_output, atol=3e-2, rtol=3e-2 + ) + torch.testing.assert_close( + flashinfer_tracked, triton_tracked, atol=3e-2, rtol=3e-2 + ) + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/layers/attention/test_gdn_prefill_backend_policy.py b/test/registered/unit/layers/attention/test_gdn_prefill_backend_policy.py index 1381981b2..95ea659e8 100644 --- a/test/registered/unit/layers/attention/test_gdn_prefill_backend_policy.py +++ b/test/registered/unit/layers/attention/test_gdn_prefill_backend_policy.py @@ -4,11 +4,18 @@ from unittest.mock import MagicMock, patch, sentinel import torch +from sglang.srt.layers.attention.hybrid_linear_attn_backend import ( + MambaAttnBackendBase, +) from sglang.srt.layers.attention.linear import gdn_backend from sglang.srt.layers.attention.linear.gdn_backend import ( + GDNAttnBackend, GDNKernelDispatcher, maybe_set_default_flashinfer_gdn_prefill, ) +from sglang.srt.layers.attention.linear.kernels.gdn_flashinfer import ( + maybe_build_flashinfer_checkpoint_plan, +) from sglang.srt.layers.attention.linear.kernels.gdn_triton import TritonGDNKernel from sglang.srt.layers.attention.linear.utils import LinearAttnKernelBackend from sglang.test.ci.ci_register import register_cpu_ci @@ -32,16 +39,13 @@ def make_runner( enable_dynamic_chunking=False, chunked_prefill_size=8192, ) - for name, value in arg_overrides.items(): - setattr(args, name, value) - # The policy routes its load-time default through the audited mutation entry # (server_args.override); mirror that on the stub so the write lands. - def _override(source, **fields): - for _field, _value in fields.items(): - setattr(args, _field, _value) - - args.override = _override + args.override = MagicMock( + side_effect=lambda _source, **fields: vars(args).update(fields) + ) + for name, value in arg_overrides.items(): + setattr(args, name, value) return SimpleNamespace( server_args=args, @@ -87,14 +91,21 @@ class TestFlashInferGDNPrefillBackendPolicy(unittest.TestCase): return runner.server_args.linear_attn_prefill_backend def test_selects_flashinfer_for_supported_sm100_gdn(self): - self.assertEqual(self.apply_policy(make_runner()), "flashinfer") - - def test_selects_flashinfer_for_no_buffer_radix_cache(self): - runner = make_runner( - uses_mamba_radix_cache=True, - mamba_radix_cache_strategy="no_buffer", - ) + runner = make_runner() self.assertEqual(self.apply_policy(runner), "flashinfer") + runner.server_args.override.assert_called_once_with( + "gdn_backend.sm100_flashinfer_default", + linear_attn_prefill_backend="flashinfer", + ) + + def test_selects_flashinfer_for_radix_cache_strategies(self): + for strategy in ("no_buffer", "extra_buffer", "extra_buffer_lazy"): + with self.subTest(strategy=strategy): + runner = make_runner( + uses_mamba_radix_cache=True, + mamba_radix_cache_strategy=strategy, + ) + self.assertEqual(self.apply_policy(runner), "flashinfer") def test_preserves_explicit_prefill_override(self): for backend in ("triton", "flashinfer", "cutedsl"): @@ -128,20 +139,6 @@ class TestFlashInferGDNPrefillBackendPolicy(unittest.TestCase): cases = ( ("non_triton_base", {"linear_attn_backend": "cutedsl"}), ("page_major_kv", {"enable_page_major_kv_layout": True}), - ( - "extra_buffer", - { - "uses_mamba_radix_cache": True, - "mamba_radix_cache_strategy": "extra_buffer", - }, - ), - ( - "extra_buffer_lazy", - { - "uses_mamba_radix_cache": True, - "mamba_radix_cache_strategy": "extra_buffer_lazy", - }, - ), ("dynamic_chunk", {"enable_dynamic_chunking": True}), ("unchunked", {"chunked_prefill_size": -1}), ("unknown_chunk", {"chunked_prefill_size": None}), @@ -151,6 +148,53 @@ class TestFlashInferGDNPrefillBackendPolicy(unittest.TestCase): with self.subTest(name=name): self.assertIsNone(self.apply_policy(make_runner(**runner_args))) + def test_builds_compact_checkpoint_plan_for_packed_sequences(self): + forward_batch = SimpleNamespace( + extend_seq_lens=torch.tensor([63, 64, 65, 127, 128, 129]), + mamba_track_mask=torch.tensor([False, True, True, True, True, True]), + # 65 on the 128-token sequence represents an interior S64 + # boundary encoded as S64 + 1 by the scheduler. + mamba_track_seqlens=torch.tensor([63, 64, 65, 127, 65, 129]), + extend_prefix_lens=torch.zeros(6, dtype=torch.int64), + ) + metadata = SimpleNamespace( + track_ssm_h_src=torch.empty(4), + track_ssm_h_dst=torch.empty(4), + ) + + with patch( + "sglang.srt.layers.attention.linear.kernels.gdn_flashinfer." + "get_server_args", + return_value=SimpleNamespace(mamba_cache_chunk_size=64), + ): + maybe_build_flashinfer_checkpoint_plan(forward_batch, metadata, "cpu") + + torch.testing.assert_close( + metadata.state_checkpoint_cu_starts, + torch.tensor([0, 0, 1, 2, 3, 5, 7]), + ) + torch.testing.assert_close(metadata.track_ssm_h_src, torch.tensor([1, 2, 3, 6])) + self.assertEqual(metadata.num_state_checkpoints, 7) + self.assertEqual(metadata.state_checkpoint_every_n_tokens, 64) + + def test_decode_tracking_without_h_source_skips_checkpoint_plan(self): + backend = object.__new__(GDNAttnBackend) + backend.device = "cpu" + backend.kernel_dispatcher = SimpleNamespace(extend_uses_state_checkpoints=True) + metadata = SimpleNamespace(has_mamba_track_mask=True, track_ssm_h_src=None) + forward_batch = SimpleNamespace( + mamba_track_mask=torch.tensor([True]), + mamba_track_indices=torch.tensor([7]), + ) + + def init_base(instance, _forward_batch): + instance.forward_metadata = metadata + + with patch.object(MambaAttnBackendBase, "init_forward_metadata", init_base): + backend.init_forward_metadata(forward_batch) + + torch.testing.assert_close(metadata.conv_states_mask_indices, torch.tensor([7])) + def test_tree_verify_uses_triton_kernel(self): flashinfer_kernel = MagicMock(supports_target_verify=True) with (