Remove CP indexer syncs around MQA logits

The CP in-seq NSA indexer path was constructing a CUDA validity mask around each fp8_mqa_logits call, synchronizing it back with .item(), and then rebuilding small CUDA tensors before topk. The valid rows are always a contiguous prefix because ke_offset is monotonic, so the count can be derived from existing CP metadata and the path can use slices instead of boolean gather/scatter.

RAGGED topk now reuses the zero row-start vector as the per-row offset override for the single-segment case to avoid cumsum/repeat_interleave metadata kernels after MQA.

Constraint: CP shared KV prefill still needs to preserve prev/next causal visibility and topk transform semantics.
Rejected: Add a new Triton range-builder first | the Python-side synchronization and boolean indexing were the larger immediate risk, and kernel work should be profiled after this smaller change.
Confidence: medium
Scope-risk: moderate
Directive: Do not reintroduce GPU mask .item() or CPU-created CUDA tensors in the fp8_mqa_logits hot path without profiler evidence.
Tested: python3 -m py_compile python/sglang/srt/layers/attention/nsa/nsa_indexer.py test/registered/unit/layers/test_nsa_cp_utils.py
Tested: g0034 docker /sgl-workspace/sglang-tai python3 test/registered/unit/layers/test_nsa_cp_utils.py -> Ran 22 tests OK
Not-tested: full GLM5 PD server throughput/profiler run

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)