diff --git a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py index ad98804b0..c368b3e75 100644 --- a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py +++ b/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( diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index e160cc400..633009d52 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/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()