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 878b17fde..bfcda1965 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py @@ -10,6 +10,7 @@ from sglang.kernels.ops.mamba.mamba_state_indices_triton import ( fused_replay_state_indices, ) from sglang.kernels.ops.mamba.mamba_state_scatter_triton import ( + fused_conv_window_scatter_with_mask, scatter_mamba_states_after_mtp_verify, track_mamba_states_all_layers, track_mamba_states_if_needed, @@ -50,6 +51,11 @@ _validate_mamba_replay_state_indices = ( class MambaAttnBackendBase(AttentionBackend): + # Per-slot accept lengths for the KDA fused-accept spec path; allocated only + # by KDAAttnBackend where `_can_fuse_accept_state` holds. None everywhere + # else — update_mamba_state_after_mtp_verify keys the fused branch on it. + accept_lens_pool: Optional[torch.Tensor] = None + def __init__(self, model_runner: ModelRunner): super().__init__() self.pad_slot_id = PAD_SLOT_ID @@ -1282,6 +1288,31 @@ class HybridLinearAttnBackend(AttentionBackend): ) return + # KDA fused-accept: the next verify seeds itself in-kernel from the + # accepted checkpoint slot (recurrent_kda's num_accepted_tokens), so the + # SSM state never round-trips through `temporal` and only the conv + # windows still need the accept rollback. Recording this round's accept + # length is what selects that seed next round; chain layout only (see + # above), so accept_lens == last_correct_step_indices + 1. The pool + # exists only where KDAAttnBackend found the contract satisfied, which + # includes mamba radix tracking being off. + accept_lens_pool = self.linear_attn_backend.accept_lens_pool + if accept_lens_pool is not None: + assert mamba_track_indices is None, "fused-accept runs with radix off" + for conv_states, intermediate_conv_window in zip( + mamba_caches.conv, mamba_caches.intermediate_conv_window + ): + fused_conv_window_scatter_with_mask( + conv_states, + intermediate_conv_window, + state_indices_tensor, + last_correct_step_indices, + ) + accept_lens_pool[state_indices_tensor.to(torch.int64)] = ( + last_correct_step_indices.to(torch.int32) + 1 + ) + return + scatter_mamba_states_after_mtp_verify( mamba_caches, state_indices_tensor, diff --git a/python/sglang/srt/layers/attention/linear/kda_backend.py b/python/sglang/srt/layers/attention/linear/kda_backend.py index f2c13b49c..699971f08 100644 --- a/python/sglang/srt/layers/attention/linear/kda_backend.py +++ b/python/sglang/srt/layers/attention/linear/kda_backend.py @@ -10,6 +10,9 @@ from sglang.kernels.ops.mamba.causal_conv1d_triton import ( ) from sglang.srt.environ import envs from sglang.srt.layers.attention.hybrid_linear_attn_backend import MambaAttnBackendBase +from sglang.srt.layers.attention.linear.kernels.kda_flashinfer import ( + build_fused_accept_indices, +) from sglang.srt.layers.attention.linear.kernels.kda_triton import TritonKDAKernel from sglang.srt.layers.attention.linear.utils import ( LinearAttnKernelBackend, @@ -33,6 +36,9 @@ elif is_cpu(): from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.model_executor.model_runner import ModelRunner from sglang.srt.runtime_context import ( + get_disagg, + get_exec, + get_memory, get_spec, ) @@ -435,6 +441,80 @@ class KDAAttnBackend(MambaAttnBackendBase): model_runner.device, ) ) + # Fused-accept spec path (flashinfer recurrent_kda): the next verify + # seeds itself in-kernel from the accepted checkpoint slot + # (num_accepted_tokens), so the per-round SSM commit scatter is skipped. + # accept_lens_pool holds last round's accept length per mamba slot; + # extend stages fresh requests with 1 (read slot 0). Its presence is the + # signal that switches the post-verify commit to conv-only. + if self._can_fuse_accept_state(verify_backend): + self.accept_lens_pool = torch.ones( + self.req_to_token_pool.size + 1, + dtype=torch.int32, + device=model_runner.device, + ) + + @staticmethod + def _can_fuse_accept_state(verify_backend) -> bool: + """Whether the verify kernel can seed itself from the accepted checkpoint. + + The seed comes from recurrent_kda's ``num_accepted_tokens``, which makes + the previous round's accepted state addressable in-kernel so the SSM + state never has to round-trip through the committed pool (`temporal`). + Hence the keying on the verify backend -- the one that selects the + target_verify kernel -- and not on decode, which is set separately. + + `temporal` then goes stale between verifies, which is what the remaining + conditions rule out: each is a reader of the committed state that the + skipped scatter would starve. Falling short of the contract falls back to + the commit scatter rather than raising -- this is a capability, not a + mode. Every condition reads its namespace bag rather than the record: + these are fields resolution decides, so the record would answer with + what the operator typed instead of what was decided. + """ + if not verify_backend.is_flashinfer(): + return False # only recurrent_kda takes num_accepted_tokens + if get_spec().speculative_algorithm is None: + return False # no verify round, and no intermediate scratch + if not get_memory().disable_radix_cache: + return False # mamba radix tracking snapshots `temporal` + if get_exec().mamba.enable_linear_replayssm_spec: + return False # the ring already owns the verify-round commitment + if get_disagg().disaggregation_mode != "null": + return False # the PD hand-off transfers `temporal` + return True + + def _fused_accept_indices( + self, + *, + cache_indices: torch.Tensor, + query_start_loc: torch.Tensor, + intermediate_state_cache: torch.Tensor, + draft_token_num: int, + ): + """Slot-indexed verify rows + accept lengths for this forward. + + Every KDA layer of a forward verifies the same requests over the same + draft window, so the build is hoisted onto the shared forward metadata: + layer 0 builds, the rest reuse. Under cuda graph the capture then holds + a single build reading the static slot buffer, instead of one per layer. + """ + metadata = self.forward_metadata + if metadata.fused_accept_state_indices is None: + batch_size = query_start_loc.shape[0] - 1 + ( + metadata.fused_accept_state_indices, + metadata.fused_accept_num_accepted, + ) = build_fused_accept_indices( + slots=cache_indices[:batch_size], + scratch_steps=intermediate_state_cache.shape[1], + draft_token_num=draft_token_num, + accept_lens_pool=self.accept_lens_pool, + ) + return ( + metadata.fused_accept_state_indices, + metadata.fused_accept_num_accepted, + ) def init_forward_metadata(self, forward_batch: ForwardBatch): super().init_forward_metadata(forward_batch) @@ -732,6 +812,19 @@ class KDAAttnBackend(MambaAttnBackendBase): forward_batch, h, ssm_states, self.forward_metadata ) + if ( + self.accept_lens_pool is not None + and not forward_batch.forward_mode.is_draft_extend_v2() + ): + # Fused-accept staging: the extend kernel just wrote this request's + # committed state; copy it into scratch slot 0 and reset the accept + # length to 1 so the first verify reads slot 0. Runs once per KDA + # layer (the nat write is idempotent; the scratch copy is per-layer). + slots = cache_indices.to(torch.int64) + intermediate_ssm = mamba_cache_params.intermediate_ssm + intermediate_ssm[slots, 0] = ssm_states[slots].to(intermediate_ssm.dtype) + self.accept_lens_pool[slots] = 1 + return core_attn_out def _forward_target_verify( @@ -951,6 +1044,21 @@ class KDAAttnBackend(MambaAttnBackendBase): retrieve_parent_token=retrieve_parent_token, lower_bound=layer.lower_bound, **ring_kwargs, + **( + dict( + zip( + ("fused_accept_state_indices", "fused_accept_num_accepted"), + self._fused_accept_indices( + cache_indices=cache_indices, + query_start_loc=query_start_loc, + intermediate_state_cache=intermediate_state_cache, + draft_token_num=draft_token_num, + ), + ) + ) + if self.accept_lens_pool is not None + else {} + ), ) if dense_token_indices is not None: # Kernel output is empty-allocated and the capped qsl skips the diff --git a/python/sglang/srt/layers/attention/linear/kernels/kda_flashinfer.py b/python/sglang/srt/layers/attention/linear/kernels/kda_flashinfer.py index 561c4b71e..180c3482b 100644 --- a/python/sglang/srt/layers/attention/linear/kernels/kda_flashinfer.py +++ b/python/sglang/srt/layers/attention/linear/kernels/kda_flashinfer.py @@ -57,6 +57,33 @@ def _get_flashinfer_kda_kernel(): return _flashinfer_kda_available, _flashinfer_recurrent_kda +def build_fused_accept_indices( + *, + slots: torch.Tensor, + scratch_steps: int, + draft_token_num: int, + accept_lens_pool: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Slot-indexed verify indices + accept lengths for fused-accept mode. + + Row n of the returned ``[N, T]`` index tensor addresses the scratch slots of + the request holding mamba slot ``slots[n]``. A padded row (``slots[n] < 0``) + yields ONLY negative indices (``-scratch_steps + step`` with + ``step < scratch_steps``), which recurrent_kda treats as inactive — the + padding contract must survive any refactor of this arithmetic. The nat + gather clamps padded slots to row 0 of the pool; their value is never + consumed (inactive rows). All ops are device-side and capture-safe. + """ + step = torch.arange(draft_token_num, device=slots.device, dtype=torch.int32) + ssm_state_indices = ( + slots.to(torch.int32)[:, None] * scratch_steps + step[None, :] + ).contiguous() # [N, T] + num_accepted_tokens = accept_lens_pool.index_select( + 0, slots.clamp(min=0).to(torch.int64) + ) + return ssm_state_indices, num_accepted_tokens + + class FlashInferKDAKernel(LinearAttnKernelBase): """FlashInfer KDA kernel: SM100 decode + MTP (target_verify), topk=1. @@ -235,6 +262,8 @@ class FlashInferKDAKernel(LinearAttnKernelBase): cache_steps: int, retrieve_parent_token: torch.Tensor, lower_bound: Optional[float] = None, + fused_accept_state_indices: Optional[torch.Tensor] = None, + fused_accept_num_accepted: Optional[torch.Tensor] = None, **kwargs, ) -> torch.Tensor: if retrieve_parent_token is not None: @@ -272,36 +301,54 @@ class FlashInferKDAKernel(LinearAttnKernelBase): f"but intermediate_ssm only has {scratch_steps}." ) - base_rows = intermediate_state_indices[:batch_size] - cache_key = ( - id(intermediate_state_indices), - batch_size, - draft_token_num, - scratch_steps, - ) - ssm_state_indices = self._verify_idx_cache.get(cache_key) - if ssm_state_indices is None: - # The fast seed copy below assumes row n in scratch belongs to request n. - expected = torch.arange( - batch_size, device=base_rows.device, dtype=base_rows.dtype + if fused_accept_state_indices is not None: + # Fused-accept mode: rows are the requests' mamba SLOTS (stable for + # the request lifetime, unlike batch positions), so last round's + # checkpoints are addressable this round. The kernel seeds each row + # from slot[nat - 1] (nat = last round's accept length, gathered + # from accept_lens_pool; fresh requests were staged with nat = 1 at + # extend) and overwrites all T slots in place — no committed-pool + # seed copy here and no SSM commit scatter after verify. Padded + # graph rows carry slot -1: every derived index stays negative, + # which recurrent_kda treats as inactive. Both tensors are built + # once per forward by the backend (see KDAAttnBackend), so the 20 + # KDA layers of a step share one build. + ssm_state_indices = fused_accept_state_indices + num_accepted_tokens = fused_accept_num_accepted + else: + num_accepted_tokens = None + base_rows = intermediate_state_indices[:batch_size] + cache_key = ( + id(intermediate_state_indices), + batch_size, + draft_token_num, + scratch_steps, ) - if not torch.equal(base_rows, expected): - raise RuntimeError( - "FlashInfer KDA verify requires an identity intermediate row-map " - "(verify_intermediate_state_indices must be arange)." + ssm_state_indices = self._verify_idx_cache.get(cache_key) + if ssm_state_indices is None: + # The fast seed copy below assumes row n in scratch belongs to + # request n. + expected = torch.arange( + batch_size, device=base_rows.device, dtype=base_rows.dtype ) - step = torch.arange(draft_token_num, device=q.device, dtype=torch.int32) - ssm_state_indices = ( - base_rows.to(torch.int32)[:, None] * scratch_steps + step[None, :] - ).contiguous() # [N, T] - self._verify_idx_cache[cache_key] = ssm_state_indices + if not torch.equal(base_rows, expected): + raise RuntimeError( + "FlashInfer KDA verify requires an identity intermediate " + "row-map (verify_intermediate_state_indices must be arange)." + ) + step = torch.arange(draft_token_num, device=q.device, dtype=torch.int32) + ssm_state_indices = ( + base_rows.to(torch.int32)[:, None] * scratch_steps + step[None, :] + ).contiguous() # [N, T] + self._verify_idx_cache[cache_key] = ssm_state_indices - # Seed step 0 from committed state, then recurrent_kda overwrites it with - # token-0 post-state. Padded graph rows clamp to slot 0; their output is ignored. - base_state = ssm_states.index_select( - 0, cache_indices[:batch_size].clamp(min=0).to(torch.int64) - ) - scratch[:batch_size, 0].copy_(base_state) + # Seed step 0 from committed state, then recurrent_kda overwrites it + # with token-0 post-state. Padded graph rows clamp to slot 0; their + # output is ignored. + base_state = ssm_states.index_select( + 0, cache_indices[:batch_size].clamp(min=0).to(torch.int64) + ) + scratch[:batch_size, 0].copy_(base_state) # Same storage as scratch, flattened over the allocated step stride. state_pool = scratch.view( @@ -325,6 +372,7 @@ class FlashInferKDAKernel(LinearAttnKernelBase): cu_seqlens=query_start_loc.to(torch.int32), ssm_state_indices=ssm_state_indices, num_spec_tokens=num_spec_tokens, + num_accepted_tokens=num_accepted_tokens, ) return output_fi.view(1, seq_len, num_v_heads, head_v_dim) diff --git a/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py b/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py index 2cb3dd4f4..5ebd6f367 100644 --- a/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py +++ b/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py @@ -62,6 +62,13 @@ class ForwardMetadata: is_target_verify: bool = False draft_token_num: int = 1 + # KDA fused-accept: the [N, T] slot-indexed scratch rows and the per-request + # accept length that seed the verify kernel. Every KDA layer of a forward + # sees the same slots and draft window, so these are built once and shared: + # a cuda-graph capture then holds one build instead of one per layer. + fused_accept_state_indices: Optional[torch.Tensor] = None + fused_accept_num_accepted: Optional[torch.Tensor] = None + has_mamba_track_mask: bool = False mamba_track_mask_indices: Optional[torch.Tensor] = None conv_states_mask_indices: Optional[torch.Tensor] = None diff --git a/test/registered/attention/unittests/hybrid_linear/test_kda_fused_accept_indices.py b/test/registered/attention/unittests/hybrid_linear/test_kda_fused_accept_indices.py new file mode 100644 index 000000000..afc37719e --- /dev/null +++ b/test/registered/attention/unittests/hybrid_linear/test_kda_fused_accept_indices.py @@ -0,0 +1,141 @@ +"""Padding-safety invariants of the KDA fused-accept verify index builder. + +``build_fused_accept_indices`` produces the slot-indexed ``[N, T]`` +``ssm_state_indices`` and the per-row ``num_accepted_tokens`` gather that +flashinfer ``recurrent_kda`` consumes in fused-accept mode. The kernel's +padding contract is: a row is inactive iff its raw slot index is negative. +Padded sglang rows carry mamba slot ``-1``, so EVERY derived index +``-1 * scratch_steps + step`` must stay negative for all +``step < scratch_steps`` — an arithmetic reorder (e.g. adding the step before +the multiply) would silently activate padded rows and corrupt neighbor state. +The nat gather must clamp padded slots in-bounds (their value is never +consumed) and keep real slots' accept lengths intact. + +CPU tensors only — the invariants are pure index arithmetic. +""" + +import unittest + +import torch + +from sglang.srt.layers.attention.linear.kernels.kda_flashinfer import ( + build_fused_accept_indices, +) +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=2, stage="base-b", runner_config="1-gpu-large") + + +class TestBuildFusedAcceptIndices(CustomTestCase): + def test_real_slots_address_their_scratch_rows(self): + for scratch_steps, draft_token_num in ((4, 4), (8, 4), (33, 33)): + slots = torch.tensor([0, 3, 7], dtype=torch.int32) + pool = torch.arange(10, dtype=torch.int32) + 1 # accept len = slot + 1 + indices, nat = build_fused_accept_indices( + slots=slots, + scratch_steps=scratch_steps, + draft_token_num=draft_token_num, + accept_lens_pool=pool, + ) + self.assertEqual(indices.shape, (3, draft_token_num)) + self.assertEqual(indices.dtype, torch.int32) + step = torch.arange(draft_token_num, dtype=torch.int32) + expected = slots[:, None] * scratch_steps + step[None, :] + self.assertTrue(torch.equal(indices, expected)) + self.assertEqual(nat.dtype, torch.int32) + self.assertTrue(torch.equal(nat, slots + 1)) + + def test_padded_slot_rows_stay_fully_negative(self): + # T == scratch_steps is the tight case: the largest step must still + # land below zero for slot -1. + for scratch_steps, draft_token_num in ((4, 4), (8, 8), (8, 4), (33, 33)): + slots = torch.tensor([2, -1, 5, -1], dtype=torch.int32) + pool = torch.full((8,), 3, dtype=torch.int32) + indices, nat = build_fused_accept_indices( + slots=slots, + scratch_steps=scratch_steps, + draft_token_num=draft_token_num, + accept_lens_pool=pool, + ) + padded_rows = indices[slots < 0] + self.assertTrue( + (padded_rows < 0).all(), + f"padded row leaked a non-negative index " + f"({scratch_steps=}, {draft_token_num=}): {padded_rows.tolist()}", + ) + real_rows = indices[slots >= 0] + self.assertTrue((real_rows >= 0).all()) + # nat gather clamps padded slots in-bounds (value unused). + self.assertEqual(nat.shape[0], 4) + + def test_nat_gather_reads_pool_values(self): + slots = torch.tensor([1, 4, -1], dtype=torch.int32) + pool = torch.tensor([9, 2, 9, 9, 5, 9], dtype=torch.int32) + _, nat = build_fused_accept_indices( + slots=slots, + scratch_steps=4, + draft_token_num=4, + accept_lens_pool=pool, + ) + self.assertEqual(nat[0].item(), 2) + self.assertEqual(nat[1].item(), 5) + # Padded row clamps to pool row 0; the value is never consumed but the + # gather itself must stay in-bounds. + self.assertEqual(nat[2].item(), 9) + + +class TestFusedAcceptPerForwardCache(CustomTestCase): + """The verify indices are built once per forward and shared by every KDA + layer. That sharing is only sound while the cache dies with the forward: a + cache that outlived it would seed the next batch from the previous batch's + mamba slots, which is a silent wrong-state bug (no shape or index error). + """ + + @staticmethod + def _build(slots, pool_values, draft_token_num=4, scratch_steps=4): + device = "cuda" if torch.cuda.is_available() else "cpu" + return build_fused_accept_indices( + slots=torch.tensor(slots, dtype=torch.int32, device=device), + scratch_steps=scratch_steps, + draft_token_num=draft_token_num, + accept_lens_pool=torch.tensor( + pool_values, dtype=torch.int32, device=device + ), + ) + + def test_shared_build_matches_a_per_layer_build(self): + """What every layer reuses must equal what it would have built itself.""" + pool = [1] * 8 + pool[3], pool[5] = 2, 4 + first_idx, first_nat = self._build([3, 5], pool) + second_idx, second_nat = self._build([3, 5], pool) + self.assertTrue(torch.equal(first_idx, second_idx)) + self.assertTrue(torch.equal(first_nat, second_nat)) + + def test_a_different_batch_builds_different_rows(self): + """Guards the staleness mode: reusing a previous forward's tensor would + address the previous forward's slots, and the values must differ so the + cache-reset is observable rather than accidentally correct.""" + pool = [1] * 8 + pool[3], pool[5], pool[6] = 2, 4, 3 + idx_a, nat_a = self._build([3, 5], pool) + idx_b, nat_b = self._build([6, 5], pool) + self.assertFalse(torch.equal(idx_a, idx_b)) + self.assertFalse(torch.equal(nat_a, nat_b)) + + def test_metadata_starts_uncached(self): + """A forward's metadata must arrive with no indices carried over: the + backend keys 'build once' on these being None.""" + from sglang.srt.layers.attention.mamba.mamba2_metadata import ForwardMetadata + + metadata = ForwardMetadata( + query_start_loc=torch.zeros(2, dtype=torch.int32), + mamba_cache_indices=torch.zeros(1, dtype=torch.int32), + ) + self.assertIsNone(metadata.fused_accept_state_indices) + self.assertIsNone(metadata.fused_accept_num_accepted) + + +if __name__ == "__main__": + unittest.main()