Cache CP shared-KV flattened request row ids per forward

build_flattened_request_row_ids ran once per layer in the per-layer
prefill attention loop and rebuilt the indexer-seq-len-derived row ids
on the GPU via repeat_interleave with a *device* repeats tensor, which
forces a device sync (it reads sum(repeats) to size the output) and
serializes the launch thread every layer.

The row ids depend only on indexer_seq_lens_cpu, which is identical
across every layer of a forward pass.  Add
get_cp_shared_kv_flattened_request_row_ids, which builds them once and
caches on the ForwardBatch (keyed by expected flattened length + device
so a different batch shape or the draft-vs-target pass never reuses a
stale tensor), and build on CPU then move once to drop the sync.

In a TP-5 warm cachebench trace this collapsed 1817 wrapper calls to 46
actual builds (97.5% fewer) and dropped ipc_materialize visibly.
GSM8K 200q x2 = 0.955 / 0.950, 0 invalid.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py

@@ -4396,9 +4396,41 @@ def build_flattened_request_row_ids(
     seq_lens = _to_cpu_int_list(seq_lens_cpu, name="seq_lens_cpu")
     if len(seq_lens) == 0:
         return torch.empty((0,), device=device, dtype=torch.long)
-    lengths = torch.tensor(seq_lens, device=device, dtype=torch.long)
-    rows = torch.arange(len(seq_lens), device=device, dtype=torch.long)
-    return torch.repeat_interleave(rows, lengths)
+    # Build on CPU then move once.  `repeat_interleave` with a *device* repeats
+    # tensor forces a device sync (it reads sum(repeats) to size the output),
+    # which on the per-layer prefill path serializes the launch thread; the row
+    # ids are tiny so a single H2D is strictly cheaper.
+    rows = torch.arange(len(seq_lens), dtype=torch.long)
+    lengths = torch.tensor(seq_lens, dtype=torch.long)
+    return torch.repeat_interleave(rows, lengths).to(device, non_blocking=True)
+
+
+def get_cp_shared_kv_flattened_request_row_ids(
+    forward_batch: Any,
+    seq_lens_cpu: Any,
+    *,
+    device: torch.device,
+) -> torch.Tensor:
+    """Per-forward-cached flattened request-row-ids.
+
+    The row ids depend only on ``seq_lens_cpu`` (the indexer seq lengths), which
+    is identical across every layer of a forward pass, so build them once and
+    reuse across layers instead of rebuilding (and re-syncing) in the per-layer
+    attention loop -- mirrors ``get_cp_shared_kv_token_loc_req_id``'s caching.
+    Keyed on the expected flattened length + device so a different batch shape
+    (or the draft vs target pass on its own forward batch) never reuses a stale
+    tensor.
+    """
+    total = sum(int(x) for x in _to_cpu_int_list(seq_lens_cpu, name="seq_lens_cpu"))
+    key = (int(total), str(device))
+    cached = getattr(forward_batch, "cp_flattened_row_ids", None)
+    cached_key = getattr(forward_batch, "cp_flattened_row_ids_key", None)
+    if cached is not None and cached_key == key and cached.device == device:
+        return cached
+    row_ids = build_flattened_request_row_ids(seq_lens_cpu, device=device)
+    forward_batch.cp_flattened_row_ids = row_ids
+    forward_batch.cp_flattened_row_ids_key = key
+    return row_ids
 
 
 def remap_logical_locs_to_slot_dense_locs(

python/sglang/srt/layers/attention/nsa_backend.py

@@ -16,8 +16,8 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_prefetch import (
 )
 from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
     build_batch_current_slot_spans,
-    build_flattened_request_row_ids,
     build_batch_prefix_slot_spans,
+    get_cp_shared_kv_flattened_request_row_ids,
     build_current_loc_remap,
     cp_shared_kv_debug_enabled,
     cp_shared_kv_debug_log,
@@ -2599,9 +2599,12 @@ class NativeSparseAttnBackend(
                                     want_prefix=len(prefix_lens_cpu) > 1,
                                 )
                             )
-                            logical_locs_row_ids = build_flattened_request_row_ids(
-                                metadata.indexer_seq_lens_cpu,
-                                device=page_table_1_flattened.device,
+                            logical_locs_row_ids = (
+                                get_cp_shared_kv_flattened_request_row_ids(
+                                    forward_batch,
+                                    metadata.indexer_seq_lens_cpu,
+                                    device=page_table_1_flattened.device,
+                                )
                             )
                             if int(logical_locs_row_ids.numel()) != int(
                                 page_table_1_flattened.numel()