diff --git a/python/sglang/srt/layers/attention/nsa/index_buf_accessor.py b/python/sglang/srt/layers/attention/nsa/index_buf_accessor.py index da8121fc7..344deed66 100644 --- a/python/sglang/srt/layers/attention/nsa/index_buf_accessor.py +++ b/python/sglang/srt/layers/attention/nsa/index_buf_accessor.py @@ -624,7 +624,7 @@ def _get_k_and_s_triton( :param page_indices: (num_pages,), int32/int64 :param seq_lens: tensor of sequence lens, int64 :param seq_len_sum: sum of all sequence len, int32 - :param seq_len_sum: max of sequence len, int32 + :param max_seq_len: max of sequence len, int32 :param page_size: int, typically 64 :param index_head_dim: int, typically 128 :return: tuple of (k_out, s_out) where diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index 675b6f9ed..bff848c1b 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -97,6 +97,11 @@ class BaseIndexerMetadata(ABC): Return: seq lens for each batch. """ + def get_indexer_seq_len(self) -> torch.Tensor: + """ + Return: seq lens for each batch. + """ + def get_nsa_extend_len_cpu(self) -> List[int]: """ Return: extend seq lens for each batch. @@ -538,11 +543,12 @@ class Indexer(MultiPlatformOp): ks, ke = metadata.get_indexer_kvcache_range() - seq_len_sum = forward_batch.seq_lens_sum - max_seq_len = torch.max(forward_batch.seq_lens_cpu).item() + indexer_seq_lens_cpu = metadata.get_indexer_seq_len_cpu() + seq_len_sum = torch.sum(indexer_seq_lens_cpu).item() + max_seq_len = torch.max(indexer_seq_lens_cpu).item() k_fp8, k_scale = forward_batch.token_to_kv_pool.get_index_k_scale_buffer( layer_id, - forward_batch.seq_lens, + metadata.get_indexer_seq_len(), block_tables, seq_len_sum, max_seq_len, diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index 875bb5803..84e93c0e1 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/python/sglang/srt/layers/attention/nsa_backend.py @@ -138,6 +138,8 @@ class NSAMetadata: indexer_k_start_end: Optional[Tuple[torch.Tensor, torch.Tensor]] = None # seq lens for each batch. indexer_seq_lens_cpu: Optional[torch.Tensor] = None + # seq lens for each batch. + indexer_seq_lens: Optional[torch.Tensor] = None # batch index for each token. token_to_batch_idx: Optional[torch.Tensor] = None @@ -194,6 +196,9 @@ class NSAIndexerMetadata(BaseIndexerMetadata): def get_indexer_kvcache_range(self) -> Tuple[torch.Tensor, torch.Tensor]: return self.attn_metadata.indexer_k_start_end + def get_indexer_seq_len(self) -> torch.Tensor: + return self.attn_metadata.indexer_seq_lens + def get_indexer_seq_len_cpu(self) -> torch.Tensor: return self.attn_metadata.indexer_seq_lens_cpu @@ -404,6 +409,7 @@ class NativeSparseAttnBackend( bs_idx_cpu = None # seq_len_cpu of selected sequences indexer_seq_lens_cpu = forward_batch.seq_lens_cpu + indexer_seq_lens = forward_batch.seq_lens if forward_batch.forward_mode.is_decode_or_idle(): extend_seq_lens_cpu = [1] * batch_size @@ -504,6 +510,7 @@ class NativeSparseAttnBackend( ) ) indexer_seq_lens_cpu = indexer_seq_lens_cpu[bs_idx_cpu] + indexer_seq_lens = indexer_seq_lens[bs_idx] cache_seqlens_int32 = cache_seqlens_int32[bs_idx] cu_seqlens_k = compute_cu_seqlens(cache_seqlens_int32) max_seqlen_k = ( @@ -641,6 +648,7 @@ class NativeSparseAttnBackend( topk_indices_offset=topk_indices_offset, indexer_k_start_end=indexer_k_start_end, indexer_seq_lens_cpu=indexer_seq_lens_cpu, + indexer_seq_lens=indexer_seq_lens, token_to_batch_idx=token_to_batch_idx, ) self.forward_metadata = metadata diff --git a/test/registered/kernels/test_nsa_indexer.py b/test/registered/kernels/test_nsa_indexer.py index 42f5d316b..abee474f5 100644 --- a/test/registered/kernels/test_nsa_indexer.py +++ b/test/registered/kernels/test_nsa_indexer.py @@ -132,6 +132,10 @@ class MockIndexerMetadata(BaseIndexerMetadata): """Return: seq lens for each batch.""" return torch.tensor(self.seq_lens, dtype=torch.int32, device="cpu") + def get_indexer_seq_len(self) -> torch.Tensor: + """Return: seq lens for each batch.""" + return torch.tensor(self.seq_lens, dtype=torch.int32, device=self.device) + def get_nsa_extend_len_cpu(self) -> List[int]: """ Return: extend seq lens for each batch.