diff --git a/python/sglang/srt/layers/attention/nsa/utils.py b/python/sglang/srt/layers/attention/nsa/utils.py index c468087fb..af0fd36ef 100644 --- a/python/sglang/srt/layers/attention/nsa/utils.py +++ b/python/sglang/srt/layers/attention/nsa/utils.py @@ -1192,6 +1192,60 @@ def select_cp_local_valid_rows_for_cache_write( return selected +def select_cp_current_valid_rows_for_reuse( + forward_batch, + current_tensor: torch.Tensor, +) -> torch.Tensor | None: + """Return CP-local valid current rows for attention current reuse. + + Cache-write call sites pass CP-local compute-padded rows and can directly use + ``select_cp_local_valid_rows_for_cache_write``. Attention current-reuse + call sites can instead see already-local valid rows or unsplit global valid + current rows. This helper accepts those contracts and returns rows aligned + with ``get_cp_shared_kv_local_out_cache_loc``. + """ + + local_out_cache_loc = get_cp_shared_kv_local_out_cache_loc(forward_batch) + local_valid_rows = ( + int(local_out_cache_loc.numel()) if local_out_cache_loc is not None else None + ) + tensor_rows = int(current_tensor.shape[0]) + if local_valid_rows is not None and tensor_rows == local_valid_rows: + return current_tensor + + plan = get_cp_shared_kv_batch_plan(forward_batch) + if plan is None: + return None + + compute_padding_enabled = bool(getattr(plan, "compute_padding_enabled", False)) + if compute_padding_enabled: + _, expected_compute_rows = _get_cp_local_valid_row_indices_cache( + forward_batch, + plan, + current_tensor.device, + ) + if tensor_rows == expected_compute_rows: + return select_cp_local_valid_rows_for_cache_write( + forward_batch, + current_tensor, + ) + + extend_seq_lens_cpu = getattr(forward_batch, "extend_seq_lens_cpu", None) + if extend_seq_lens_cpu is None or len(extend_seq_lens_cpu) == 0: + return None + global_valid_rows = sum(int(x) for x in extend_seq_lens_cpu) + if global_valid_rows <= 0 or tensor_rows < global_valid_rows: + return None + + return split_tensor_by_cp_batch_plan( + current_tensor[:global_valid_rows].contiguous(), + plan, + mode="1d" if current_tensor.dim() == 1 else "data", + split_kind="valid", + static_padded_tokens=_get_forward_batch_static_padded_tokens(forward_batch), + ) + + def _pad_cp_request_tensor_for_split( tensor: torch.Tensor, *, diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index 3fc5873c8..998e44ea9 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/python/sglang/srt/layers/attention/nsa_backend.py @@ -63,7 +63,7 @@ from sglang.srt.layers.attention.nsa.utils import ( nsa_cp_round_robin_split_q_seqs, nsa_use_prefill_cp, pad_nsa_cache_seqlens, - select_cp_local_valid_rows_for_cache_write, + select_cp_current_valid_rows_for_reuse, ) from sglang.srt.layers.attention.utils import ( concat_mla_absorb_q_general, @@ -1853,10 +1853,10 @@ class NativeSparseAttnBackend( current_k_rope = None if compute_padding_current: assert k is not None and k_rope is not None - current_k_nope = select_cp_local_valid_rows_for_cache_write( + current_k_nope = select_cp_current_valid_rows_for_reuse( forward_batch, k ) - current_k_rope = select_cp_local_valid_rows_for_cache_write( + current_k_rope = select_cp_current_valid_rows_for_reuse( forward_batch, k_rope ) current_locs_for_reuse = get_cp_shared_kv_local_out_cache_loc( @@ -1869,6 +1869,8 @@ class NativeSparseAttnBackend( ) can_reuse_current_kv = ( current_kv_rows_for_reuse is not None + and current_k_nope is not None + and current_k_rope is not None and int(current_k_nope.shape[0]) == current_kv_rows_for_reuse and int(current_k_rope.shape[0]) == current_kv_rows_for_reuse ) @@ -2358,6 +2360,7 @@ class NativeSparseAttnBackend( ) if index_prefetcher is not None: index_prefetcher.launch_pending_reduce() + try: if nsa_impl == "tilelang": if q_rope is not None: @@ -2381,14 +2384,165 @@ class NativeSparseAttnBackend( assert page_table_1_flattened is not None if forward_batch.uses_cp_shared_kv: assert forward_batch.cp_shared_kv_layout is not None + if k is None or k_rope is None: + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][mla_ragged_current_sync] " + "CP shared KV RAGGED cache-hit compose requires " + "fresh current MLA KV rows." + ) + batch_plan = get_cp_shared_kv_batch_plan(forward_batch) + compute_padding_current = batch_plan is not None and bool( + getattr(batch_plan, "compute_padding_enabled", False) + ) + current_locs_for_reuse = get_cp_shared_kv_local_out_cache_loc( + forward_batch + ) + if current_locs_for_reuse is None: + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][mla_ragged_current_sync] " + "CP shared KV RAGGED cache-hit compose requires " + "local out_cache_loc for current suffix rows." + ) + if compute_padding_current: + current_k_nope = select_cp_current_valid_rows_for_reuse( + forward_batch, + k, + ) + current_k_rope = select_cp_current_valid_rows_for_reuse( + forward_batch, + k_rope, + ) + if current_k_nope is None or current_k_rope is None: + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][mla_ragged_current_sync] " + "CP shared KV RAGGED cache-hit compose could " + "not derive valid current rows under " + "compute padding. " + f"k_rows={int(k.shape[0])} " + f"k_rope_rows={int(k_rope.shape[0])} " + f"local_locs={int(current_locs_for_reuse.numel())}" + ) + else: + local_current_rows = int(current_locs_for_reuse.numel()) + if int(k.shape[0]) == local_current_rows and int( + k_rope.shape[0] + ) == local_current_rows: + current_k_nope = k + current_k_rope = k_rope + else: + valid_current_rows = current_extend_kv_rows_for_reuse( + forward_batch, + k, + k_rope, + ) + if valid_current_rows is None: + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][mla_ragged_current_sync] " + "CP shared KV RAGGED cache-hit compose could " + "not derive valid current rows. " + f"k_rows={int(k.shape[0])} " + f"k_rope_rows={int(k_rope.shape[0])} " + f"local_locs={local_current_rows}" + ) + current_k_nope = cp_split_and_rebuild_data( + forward_batch, + k[: int(valid_current_rows)].contiguous(), + ) + current_k_rope = cp_split_and_rebuild_data( + forward_batch, + k_rope[: int(valid_current_rows)].contiguous(), + ) + if ( + int(current_k_nope.shape[0]) + != int(current_locs_for_reuse.numel()) + or int(current_k_rope.shape[0]) + != int(current_locs_for_reuse.numel()) + ): + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][mla_ragged_current_sync] " + "CP shared KV RAGGED current rows do not match " + "local out_cache_loc. " + f"k_rows={int(current_k_nope.shape[0])} " + f"k_rope_rows={int(current_k_rope.shape[0])} " + f"local_locs={int(current_locs_for_reuse.numel())}" + ) + if is_packed_fp8_mla_kv_cache(kv_cache): + current_kv_cache = pack_current_mla_kv_for_reuse( + current_k_nope, + current_k_rope, + kv_cache=kv_cache, + ) + else: + current_kv_cache = _cat( + [current_k_nope, current_k_rope], + dim=-1, + ) + page_size = int(forward_batch.token_to_kv_pool.page_size) + prefix_lens_cpu = getattr( + forward_batch, + "extend_prefix_lens_cpu", + None, + ) + extend_lens_cpu = getattr( + forward_batch, + "extend_seq_lens_cpu", + None, + ) + prefix_lens_valid = ( + prefix_lens_cpu is not None + and len(prefix_lens_cpu) > 0 + and all( + int(prefix_len) >= 0 + and int(prefix_len) % page_size == 0 + for prefix_len in prefix_lens_cpu + ) + ) + if not prefix_lens_valid: + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][mla_ragged_current_sync] " + "CP shared KV RAGGED cache-hit compose requires " + "page-aligned prefix metadata. " + f"prefix_lens={prefix_lens_cpu} " + f"extend_lens={extend_lens_cpu} " + f"page_size={page_size}" + ) + if len(prefix_lens_cpu) == 1: + prefix_pages = int(prefix_lens_cpu[0]) // page_size + prefix_slot_spans = None + else: + prefix_pages = 0 + prefix_slot_spans = build_batch_prefix_slot_spans( + logical_pages=metadata.real_page_table, + prefix_lens_cpu=prefix_lens_cpu, + page_size=page_size, + ) + current_slot_spans = build_batch_current_slot_spans( + logical_pages=metadata.real_page_table, + prefix_lens_cpu=prefix_lens_cpu, + extend_lens_cpu=extend_lens_cpu, + page_size=page_size, + ) + slot_remap = get_or_build_shared_token_kv_slot_remap( + forward_batch, + kv_cache=kv_cache, + remap_logical_pages=metadata.real_page_table, + layout=forward_batch.cp_shared_kv_layout, + page_size=page_size, + ) kv_cache, page_table_1_flattened = ( - materialize_shared_token_kv_buffer( + materialize_prefix_and_reuse_current_kv_page_slots( kv_cache=kv_cache, logical_locs=page_table_1_flattened, + current_kv_cache=current_kv_cache, + current_locs=current_locs_for_reuse, + slot_remap=slot_remap, layout=forward_batch.cp_shared_kv_layout, - page_size=forward_batch.token_to_kv_pool.page_size, - nvtx_source="mla.ragged_full_materialize", - nvtx_layer_id=layer.layer_id, + page_size=page_size, + prefix_pages=prefix_pages, + prefix_slot_spans=prefix_slot_spans, + current_slot_spans=current_slot_spans, + layer_id=layer.layer_id, + nvtx_source="mla.ragged_partial_current_sync", ) ) kv_cache = dequantize_k_cache_paged( diff --git a/test/registered/unit/layers/test_nsa_cp_utils.py b/test/registered/unit/layers/test_nsa_cp_utils.py index 55fbf5f25..c1ff1d18f 100644 --- a/test/registered/unit/layers/test_nsa_cp_utils.py +++ b/test/registered/unit/layers/test_nsa_cp_utils.py @@ -32,6 +32,7 @@ from sglang.srt.layers.attention.nsa.utils import ( nsa_use_prefill_cp, pad_cp_local_input_ids_for_embedding, prepare_input_dp_with_cp_dsa, + select_cp_current_valid_rows_for_reuse, select_cp_local_valid_rows_for_cache_write, split_tensor_by_cp_batch_plan, split_in_seq_cp_local_pair, @@ -1703,6 +1704,43 @@ class TestNSAInSeqCPUtils(unittest.TestCase): self.assertEqual(selected.tolist(), [[50.0, 51.0]]) + def test_select_cp_current_valid_rows_accepts_global_rows_under_compute_padding( + self, + ): + import torch + + plan = build_batch_page_aligned_in_seq_split_plan( + extend_lens=[65], + prefix_lens=[0], + page_size=64, + cp_size=8, + cp_rank=1, + ) + self.assertTrue(plan.compute_padding_enabled) + global_current = torch.arange(65 * 2, dtype=torch.float32).view(65, 2) + expected = split_tensor_by_cp_batch_plan( + global_current, + plan, + mode="data", + split_kind="valid", + ) + forward_batch = SimpleNamespace( + extend_seq_lens_cpu=[65], + cp_local_out_cache_loc=torch.arange(expected.shape[0], dtype=torch.long), + nsa_cp_metadata=NSAContextParallelMetadata( + batch_size=1, + batch_plan=plan, + ), + ) + + selected = select_cp_current_valid_rows_for_reuse( + forward_batch, + global_current, + ) + + self.assertIsNotNone(selected) + self.assertTrue(torch.equal(selected, expected)) + def test_cp_split_and_rebuild_position_is_batch_aware_and_compute_padded(self): import torch diff --git a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py b/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py index b7b485cdd..aa641898d 100644 --- a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py +++ b/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py @@ -1330,7 +1330,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): self.assertIsNone(result) - def test_fp8_ragged_mla_defers_cp_materialize_to_flattened_path(self): + def test_fp8_ragged_mla_uses_page_slot_current_compose_for_cache_hit(self): from pathlib import Path source = ( @@ -1353,7 +1353,13 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): ragged_end = source.index(" attn_output = self._forward_flashmla_sparse", ragged_start) ragged_source = source[ragged_start:ragged_end] self.assertIn("page_table_1_flattened", ragged_source) - self.assertIn("materialize_shared_token_kv_buffer", ragged_source) + self.assertIn( + "materialize_prefix_and_reuse_current_kv_page_slots", + ragged_source, + ) + self.assertIn("select_cp_current_valid_rows_for_reuse", ragged_source) + self.assertIn("prefix_slot_spans=", ragged_source) + self.assertIn("current_slot_spans=", ragged_source) @unittest.skipIf(not torch.cuda.is_available(), "CUDA is required") def test_tai_current_slot_fill_sparse_page_self_test_passes_on_installed_kernel(