diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index 54e5cb889..7fa4c045e 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -75,6 +75,33 @@ if TYPE_CHECKING: DUAL_STREAM_TOKEN_THRESHOLD = 1024 if _is_cuda else 0 +def _compute_contiguous_valid_cp_query_count( + cp_kv_end: int, + actual_seq_q: int, + logical_kv_limit: int, +) -> int: + """Return the visible prefix length for in-seq CP query rows. + + `_get_topk_ragged_with_cp` constructs monotonically increasing `ke_offset` + values for one CP segment: + + [cp_kv_end - actual_seq_q + 1, ..., cp_kv_end] + + The validity predicate is `ke_offset <= logical_kv_limit`, so valid rows are + always a contiguous prefix. Compute the prefix length with CPU metadata that + already exists instead of building a CUDA mask and synchronizing it back via + `.item()`. + """ + + actual_seq_q = int(actual_seq_q) + if actual_seq_q <= 0: + return 0 + + ke_start = int(cp_kv_end) - actual_seq_q + 1 + valid_count = int(logical_kv_limit) - ke_start + 1 + return max(0, min(actual_seq_q, valid_count)) + + class BaseIndexerMetadata(ABC): @abstractmethod def get_seqlens_int32(self) -> torch.Tensor: @@ -845,6 +872,7 @@ class Indexer(MultiPlatformOp): current_index_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, shared_index_buffer: Optional[torch.Tensor] = None, shared_block_tables: Optional[torch.Tensor] = None, + actual_seq_q_tensor: Optional[torch.Tensor] = None, ) -> torch.Tensor: if TYPE_CHECKING: assert isinstance(forward_batch.token_to_kv_pool, NSATokenToKVPool) @@ -977,26 +1005,29 @@ class Indexer(MultiPlatformOp): int(forward_batch.seq_lens_cpu[0].item()), int(page_table_1.shape[1]), ) - ke_offset = torch.arange( - (cp_kv_end - actual_seq_q) + 1, - cp_kv_end + 1, - dtype=torch.int32, - device="cuda", + valid_q_count = _compute_contiguous_valid_cp_query_count( + cp_kv_end=cp_kv_end, + actual_seq_q=actual_seq_q, + logical_kv_limit=logical_kv_limit, ) - valid_q_mask = ke_offset <= logical_kv_limit - valid_q_count = int(valid_q_mask.sum().item()) - topk_result = torch.full( - (actual_seq_q, self.index_topk), - -1, + if valid_q_count == 0: + return torch.full( + (actual_seq_q, self.index_topk), + -1, + dtype=torch.int32, + device=q_fp8.device, + ) + + ke_start = cp_kv_end - actual_seq_q + 1 + ke_offset = torch.arange( + ke_start, + ke_start + valid_q_count, dtype=torch.int32, device=q_fp8.device, ) - if valid_q_count == 0: - return topk_result - q_fp8 = q_fp8[valid_q_mask].contiguous() - weights = weights[valid_q_mask].contiguous() - ke_offset = ke_offset[valid_q_mask].contiguous() + q_fp8 = q_fp8[:valid_q_count] + weights = weights[:valid_q_count] kv_len = min(cp_kv_end, logical_kv_limit) if current_index_kv is None: assert index_buffer is not None @@ -1020,8 +1051,8 @@ class Indexer(MultiPlatformOp): k_fp8 = k_fp8[:kv_len].contiguous() k_scale = k_scale[:kv_len].view(torch.float32).squeeze(-1).contiguous() kv_fp8 = (k_fp8, k_scale) - ks = torch.full((valid_q_count,), offset, dtype=torch.int32, device="cuda") - ke = ks + ke_offset + ks = torch.zeros_like(ke_offset) + ke = ke_offset with self._with_real_sm_count(): logits = deep_gemm.fp8_mqa_logits( @@ -1032,17 +1063,39 @@ class Indexer(MultiPlatformOp): ke, clean_logits=False, ) - actual_seq_q_tensor = torch.tensor([valid_q_count], dtype=torch.int32).to( - device="cuda", non_blocking=True - ) + topk_indices_offset_override = None + if ( + getattr(getattr(metadata, "topk_transform_method", None), "name", None) + == "RAGGED" + ): + # For this single CP segment the old + # compute_cu_seqlens([valid_q_count]) + repeat_interleave path + # produced a zero offset per query. Reuse `ks` and avoid the + # post-MQA metadata kernels entirely. + topk_indices_offset_override = ks + actual_seq_q_tensor = None + 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) + elif actual_seq_q_tensor.ndim == 0: + actual_seq_q_tensor = actual_seq_q_tensor.reshape(1) valid_topk_result = metadata.topk_transform( logits, self.index_topk, ks=ks, cu_seqlens_q=actual_seq_q_tensor, ke_offset=ke_offset, + topk_indices_offset_override=topk_indices_offset_override, ) - topk_result[valid_q_mask] = valid_topk_result + if valid_q_count == actual_seq_q: + topk_result = valid_topk_result + else: + topk_result = torch.full( + (actual_seq_q, self.index_topk), + -1, + dtype=torch.int32, + device=q_fp8.device, + ) + topk_result[:valid_q_count] = valid_topk_result return topk_result @@ -1100,6 +1153,7 @@ class Indexer(MultiPlatformOp): current_index_kv=current_index_kv, 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, ) topk_result_next = self._get_topk_ragged_with_cp( @@ -1113,6 +1167,7 @@ class Indexer(MultiPlatformOp): current_index_kv=current_index_kv, 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, ) return torch.cat([topk_result_prev, topk_result_next], dim=0) diff --git a/test/registered/unit/layers/test_nsa_cp_utils.py b/test/registered/unit/layers/test_nsa_cp_utils.py index ce8237cfc..9010adae8 100644 --- a/test/registered/unit/layers/test_nsa_cp_utils.py +++ b/test/registered/unit/layers/test_nsa_cp_utils.py @@ -19,6 +19,44 @@ register_cpu_ci(est_time=1, suite="stage-a-test-cpu") class TestNSAInSeqCPUtils(unittest.TestCase): + def test_contiguous_valid_cp_query_count(self): + from sglang.srt.layers.attention.nsa.nsa_indexer import ( + _compute_contiguous_valid_cp_query_count, + ) + + self.assertEqual( + _compute_contiguous_valid_cp_query_count( + cp_kv_end=1024, + actual_seq_q=128, + logical_kv_limit=1024, + ), + 128, + ) + self.assertEqual( + _compute_contiguous_valid_cp_query_count( + cp_kv_end=1100, + actual_seq_q=128, + logical_kv_limit=1024, + ), + 52, + ) + self.assertEqual( + _compute_contiguous_valid_cp_query_count( + cp_kv_end=1100, + actual_seq_q=64, + logical_kv_limit=1000, + ), + 0, + ) + self.assertEqual( + _compute_contiguous_valid_cp_query_count( + cp_kv_end=100, + actual_seq_q=0, + logical_kv_limit=100, + ), + 0, + ) + def assert_page_aligned_boundaries( self, split_list, *, extend_prefix_len, extend_len, page_size ): @@ -415,11 +453,13 @@ class TestNSAInSeqCPUtils(unittest.TestCase): current_index_kv=None, shared_index_buffer=None, shared_block_tables=None, + actual_seq_q_tensor=None, ): topk_calls.append( { "kv_len": kv_len, "actual_seq_q": actual_seq_q, + "actual_seq_q_tensor": actual_seq_q_tensor, "shared_index_buffer": shared_index_buffer, "shared_block_tables": shared_block_tables, "current_index_kv": current_index_kv, @@ -497,10 +537,12 @@ class TestNSAInSeqCPUtils(unittest.TestCase): current_index_kv=None, shared_index_buffer=None, shared_block_tables=None, + actual_seq_q_tensor=None, ): topk_calls.append( { "current_index_kv": current_index_kv, + "actual_seq_q_tensor": actual_seq_q_tensor, "shared_index_buffer": shared_index_buffer, "shared_block_tables": shared_block_tables, }