diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index 7fa4c045e..7d63d9d81 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -873,6 +873,7 @@ class Indexer(MultiPlatformOp): shared_index_buffer: Optional[torch.Tensor] = None, shared_block_tables: Optional[torch.Tensor] = None, actual_seq_q_tensor: Optional[torch.Tensor] = None, + actual_seq_q_cu_tensor: Optional[torch.Tensor] = None, ) -> torch.Tensor: if TYPE_CHECKING: assert isinstance(forward_batch.token_to_kv_pool, NSATokenToKVPool) @@ -1064,6 +1065,7 @@ class Indexer(MultiPlatformOp): clean_logits=False, ) topk_indices_offset_override = None + cu_seqlens_q_topk_override = None if ( getattr(getattr(metadata, "topk_transform_method", None), "name", None) == "RAGGED" @@ -1074,8 +1076,13 @@ class Indexer(MultiPlatformOp): # post-MQA metadata kernels entirely. topk_indices_offset_override = ks actual_seq_q_tensor = None + elif valid_q_count == actual_seq_q and actual_seq_q_cu_tensor is not None: + cu_seqlens_q_topk_override = actual_seq_q_cu_tensor elif actual_seq_q_tensor is None or valid_q_count != actual_seq_q: actual_seq_q_tensor = ke_offset.new_full((1,), valid_q_count) + cu_seqlens_q_topk_override = ke_offset.new_empty((2,)) + cu_seqlens_q_topk_override[0] = 0 + cu_seqlens_q_topk_override[1] = actual_seq_q_tensor.reshape(-1)[0] elif actual_seq_q_tensor.ndim == 0: actual_seq_q_tensor = actual_seq_q_tensor.reshape(1) valid_topk_result = metadata.topk_transform( @@ -1085,6 +1092,7 @@ class Indexer(MultiPlatformOp): cu_seqlens_q=actual_seq_q_tensor, ke_offset=ke_offset, topk_indices_offset_override=topk_indices_offset_override, + cu_seqlens_q_topk_override=cu_seqlens_q_topk_override, ) if valid_q_count == actual_seq_q: topk_result = valid_topk_result @@ -1154,6 +1162,7 @@ class Indexer(MultiPlatformOp): shared_index_buffer=shared_index_buffer, shared_block_tables=shared_block_tables, actual_seq_q_tensor=forward_batch.nsa_cp_metadata.actual_seq_q_prev_tensor, + actual_seq_q_cu_tensor=forward_batch.nsa_cp_metadata.actual_seq_q_prev_cu_tensor, ) topk_result_next = self._get_topk_ragged_with_cp( @@ -1168,6 +1177,7 @@ class Indexer(MultiPlatformOp): shared_index_buffer=shared_index_buffer, shared_block_tables=shared_block_tables, actual_seq_q_tensor=forward_batch.nsa_cp_metadata.actual_seq_q_next_tensor, + actual_seq_q_cu_tensor=forward_batch.nsa_cp_metadata.actual_seq_q_next_cu_tensor, ) return torch.cat([topk_result_prev, topk_result_next], dim=0) diff --git a/python/sglang/srt/layers/attention/nsa/utils.py b/python/sglang/srt/layers/attention/nsa/utils.py index 1fb5a6038..2996c5293 100644 --- a/python/sglang/srt/layers/attention/nsa/utils.py +++ b/python/sglang/srt/layers/attention/nsa/utils.py @@ -180,6 +180,8 @@ class NSAContextParallelMetadata: kv_len_next_tensor: torch.Tensor = None actual_seq_q_prev_tensor: torch.Tensor = None actual_seq_q_next_tensor: torch.Tensor = None + actual_seq_q_prev_cu_tensor: torch.Tensor = None + actual_seq_q_next_cu_tensor: torch.Tensor = None total_seq_lens: torch.Tensor = None page_aligned: bool = False page_size: int = 1 @@ -921,6 +923,12 @@ def prepare_input_dp_with_cp_dsa( actual_seq_q_next_tensor = torch.tensor(actual_seq_q_next).to( device="cuda", dtype=torch.int32 ) + actual_seq_q_prev_cu_tensor = torch.tensor( + [0, actual_seq_q_prev], device="cuda", dtype=torch.int32 + ) + actual_seq_q_next_cu_tensor = torch.tensor( + [0, actual_seq_q_next], device="cuda", dtype=torch.int32 + ) nsa_cp_metadata = NSAContextParallelMetadata( split_list=split_list, @@ -937,6 +945,8 @@ def prepare_input_dp_with_cp_dsa( kv_len_next_tensor=kv_len_next_tensor, actual_seq_q_prev_tensor=actual_seq_q_prev_tensor, actual_seq_q_next_tensor=actual_seq_q_next_tensor, + actual_seq_q_prev_cu_tensor=actual_seq_q_prev_cu_tensor, + actual_seq_q_next_cu_tensor=actual_seq_q_next_cu_tensor, total_seq_lens=comm_total_seq_lens, page_aligned=page_split_info.page_aligned, page_size=page_split_info.page_size, diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index 41c272d07..bfd6d026f 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/python/sglang/srt/layers/attention/nsa_backend.py @@ -357,6 +357,7 @@ class NSAIndexerMetadata(BaseIndexerMetadata): ke_offset: torch.Tensor = None, batch_idx_list: List[int] = None, topk_indices_offset_override: Optional[torch.Tensor] = None, + cu_seqlens_q_topk_override: Optional[torch.Tensor] = None, ) -> torch.Tensor: from sgl_kernel import ( fast_topk_transform_fused, @@ -364,19 +365,6 @@ class NSAIndexerMetadata(BaseIndexerMetadata): fast_topk_v2, ) - if topk_indices_offset_override is not None: - cu_topk_indices_offset = topk_indices_offset_override - cu_seqlens_q_topk = None - elif cu_seqlens_q is not None: - cu_seqlens_q = cu_seqlens_q.to(torch.int32) - cu_seqlens_q_topk = compute_cu_seqlens(cu_seqlens_q) - cu_topk_indices_offset = torch.repeat_interleave( - cu_seqlens_q_topk[:-1], - cu_seqlens_q, - ) - else: - cu_seqlens_q_topk = self.attn_metadata.cu_seqlens_q - cu_topk_indices_offset = self.attn_metadata.topk_indices_offset if ke_offset is not None: seq_lens_topk = ke_offset else: @@ -389,6 +377,14 @@ class NSAIndexerMetadata(BaseIndexerMetadata): if not envs.SGLANG_NSA_FUSE_TOPK.get() or self.force_unfused_topk: return fast_topk_v2(logits, seq_lens_topk, topk, row_starts=ks) elif self.topk_transform_method == TopkTransformMethod.PAGED: + if cu_seqlens_q_topk_override is not None: + cu_seqlens_q_topk = cu_seqlens_q_topk_override + elif cu_seqlens_q is not None: + if cu_seqlens_q.dtype != torch.int32: + cu_seqlens_q = cu_seqlens_q.to(torch.int32) + cu_seqlens_q_topk = compute_cu_seqlens(cu_seqlens_q) + else: + cu_seqlens_q_topk = self.attn_metadata.cu_seqlens_q # NOTE(dark): if fused, we return a transformed page table directly validate_paged_topk = ( self.validate_paged_topk @@ -415,6 +411,21 @@ class NSAIndexerMetadata(BaseIndexerMetadata): _validate_paged_topk_transform_output(topk_result, page_table_size_1) return topk_result elif self.topk_transform_method == TopkTransformMethod.RAGGED: + if topk_indices_offset_override is not None: + cu_topk_indices_offset = topk_indices_offset_override + elif cu_seqlens_q is not None: + if cu_seqlens_q.dtype != torch.int32: + cu_seqlens_q = cu_seqlens_q.to(torch.int32) + if cu_seqlens_q_topk_override is not None: + cu_seqlens_q_topk = cu_seqlens_q_topk_override + else: + cu_seqlens_q_topk = compute_cu_seqlens(cu_seqlens_q) + cu_topk_indices_offset = torch.repeat_interleave( + cu_seqlens_q_topk[:-1], + cu_seqlens_q, + ) + else: + cu_topk_indices_offset = self.attn_metadata.topk_indices_offset return fast_topk_transform_ragged_fused( score=logits, lengths=seq_lens_topk, diff --git a/test/registered/unit/layers/test_nsa_cp_utils.py b/test/registered/unit/layers/test_nsa_cp_utils.py index 9010adae8..76e6f0487 100644 --- a/test/registered/unit/layers/test_nsa_cp_utils.py +++ b/test/registered/unit/layers/test_nsa_cp_utils.py @@ -1,4 +1,5 @@ import unittest +import sys from types import SimpleNamespace from unittest.mock import patch @@ -454,12 +455,14 @@ class TestNSAInSeqCPUtils(unittest.TestCase): shared_index_buffer=None, shared_block_tables=None, actual_seq_q_tensor=None, + actual_seq_q_cu_tensor=None, ): topk_calls.append( { "kv_len": kv_len, "actual_seq_q": actual_seq_q, "actual_seq_q_tensor": actual_seq_q_tensor, + "actual_seq_q_cu_tensor": actual_seq_q_cu_tensor, "shared_index_buffer": shared_index_buffer, "shared_block_tables": shared_block_tables, "current_index_kv": current_index_kv, @@ -479,6 +482,8 @@ class TestNSAInSeqCPUtils(unittest.TestCase): kv_len_next=9, actual_seq_q_prev=3, actual_seq_q_next=2, + actual_seq_q_prev_cu_tensor=torch.tensor([0, 3], dtype=torch.int32), + actual_seq_q_next_cu_tensor=torch.tensor([0, 2], dtype=torch.int32), ) }, )() @@ -505,6 +510,8 @@ class TestNSAInSeqCPUtils(unittest.TestCase): self.assertIsNone(topk_calls[0]["current_index_kv"]) self.assertEqual(topk_calls[0]["kv_len"], 5) self.assertEqual(topk_calls[1]["kv_len"], 9) + self.assertEqual(topk_calls[0]["actual_seq_q_cu_tensor"].tolist(), [0, 3]) + self.assertEqual(topk_calls[1]["actual_seq_q_cu_tensor"].tolist(), [0, 2]) self.assertEqual(result.tolist(), [[1, 1], [1, 1], [1, 1], [2, 2], [2, 2]]) def test_indexer_in_seq_cp_pair_skips_materialize_when_current_index_reused(self): @@ -538,11 +545,13 @@ class TestNSAInSeqCPUtils(unittest.TestCase): shared_index_buffer=None, shared_block_tables=None, actual_seq_q_tensor=None, + actual_seq_q_cu_tensor=None, ): topk_calls.append( { "current_index_kv": current_index_kv, "actual_seq_q_tensor": actual_seq_q_tensor, + "actual_seq_q_cu_tensor": actual_seq_q_cu_tensor, "shared_index_buffer": shared_index_buffer, "shared_block_tables": shared_block_tables, } @@ -561,6 +570,8 @@ class TestNSAInSeqCPUtils(unittest.TestCase): kv_len_next=9, actual_seq_q_prev=3, actual_seq_q_next=2, + actual_seq_q_prev_cu_tensor=torch.tensor([0, 3], dtype=torch.int32), + actual_seq_q_next_cu_tensor=torch.tensor([0, 2], dtype=torch.int32), ) }, )() @@ -581,8 +592,84 @@ class TestNSAInSeqCPUtils(unittest.TestCase): self.assertIs(topk_calls[1]["current_index_kv"], current_index_kv) self.assertIsNone(topk_calls[0]["shared_index_buffer"]) self.assertIsNone(topk_calls[1]["shared_block_tables"]) + self.assertEqual(topk_calls[0]["actual_seq_q_cu_tensor"].tolist(), [0, 3]) + self.assertEqual(topk_calls[1]["actual_seq_q_cu_tensor"].tolist(), [0, 2]) self.assertEqual(result.tolist(), [[1, 1], [1, 1], [1, 1], [2, 2], [2, 2]]) + def test_paged_topk_transform_uses_cu_override_without_scan_metadata_ops(self): + import torch + + from sglang.srt.layers.attention.nsa_backend import ( + NSAMetadata, + NSAIndexerMetadata, + TopkTransformMethod, + ) + + cu_override = torch.tensor([0, 4], dtype=torch.int32) + attn_metadata = NSAMetadata( + page_size=64, + cache_seqlens_int32=torch.tensor([4], dtype=torch.int32), + max_seq_len_q=4, + max_seq_len_k=8, + cu_seqlens_q=torch.tensor([0, 4], dtype=torch.int32), + cu_seqlens_k=torch.tensor([0, 8], dtype=torch.int32), + page_table_1=torch.arange(8, dtype=torch.int32).view(1, 8), + real_page_table=torch.arange(8, dtype=torch.int32).view(1, 8), + nsa_cache_seqlens_int32=torch.tensor([4], dtype=torch.int32), + nsa_cu_seqlens_q=torch.arange(2, dtype=torch.int32), + nsa_cu_seqlens_k=torch.tensor([0, 4], dtype=torch.int32), + nsa_extend_seq_lens_list=[4], + nsa_seqlens_expanded=torch.arange(1, 5, dtype=torch.int32), + topk_indices_offset=torch.zeros(4, dtype=torch.int32), + ) + metadata = NSAIndexerMetadata( + attn_metadata=attn_metadata, + topk_transform_method=TopkTransformMethod.PAGED, + ) + logits = torch.zeros(4, 8) + lengths = torch.arange(1, 5, dtype=torch.int32) + expected = torch.full((4, 2), 7, dtype=torch.int32) + + def fake_fused(**kwargs): + self.assertIs(kwargs["cu_seqlens_q"], cu_override) + self.assertIs(kwargs["lengths"], lengths) + return expected + + fake_sgl_kernel = SimpleNamespace( + fast_topk_transform_fused=fake_fused, + fast_topk_transform_ragged_fused=lambda **_: (_ for _ in ()).throw( + AssertionError("ragged path should not run") + ), + fast_topk_v2=lambda *_, **__: (_ for _ in ()).throw( + AssertionError("unfused path should not run") + ), + ) + + with ( + patch.dict(sys.modules, {"sgl_kernel": fake_sgl_kernel}), + patch( + "sglang.srt.layers.attention.nsa_backend.envs.SGLANG_NSA_FUSE_TOPK.get", + return_value=True, + ), + patch( + "sglang.srt.layers.attention.nsa_backend.compute_cu_seqlens", + side_effect=AssertionError("paged override should skip cumsum"), + ), + patch( + "torch.repeat_interleave", + side_effect=AssertionError("paged topk should not build ragged offsets"), + ), + ): + actual = metadata.topk_transform( + logits, + topk=2, + cu_seqlens_q=torch.tensor([4], dtype=torch.int32), + ke_offset=lengths, + cu_seqlens_q_topk_override=cu_override, + ) + + self.assertIs(actual, expected) + if __name__ == "__main__": unittest.main()