diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 2d81be044..d1dbdab33 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -203,6 +203,7 @@ class Envs: SGLANG_FORCE_SHUTDOWN = EnvBool(False) SGLANG_DEBUG_MEMORY_POOL = EnvBool(False) SGLANG_DEBUG_CP_SHARED_KV = EnvBool(False) + SGLANG_CP_SHARED_KV_CURRENT_REUSE = EnvBool(False) SGLANG_TEST_REQUEST_TIME_STATS = EnvBool(False) SGLANG_DISABLE_TP_MEMORY_INBALANCE_CHECK = EnvBool(False) SGLANG_SIMULATE_ACC_LEN = EnvFloat(-1) 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 3776e7f45..563e68b80 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 @@ -17,6 +17,44 @@ def cp_shared_kv_debug_enabled() -> bool: return envs.SGLANG_DEBUG_CP_SHARED_KV.get() +def cp_shared_kv_current_reuse_enabled() -> bool: + return envs.SGLANG_CP_SHARED_KV_CURRENT_REUSE.get() + + +def is_current_only_extend_batch(forward_batch) -> bool: + """Return whether an extend batch has no cached/history tokens. + + This intentionally uses CPU metadata instead of scanning CUDA page tables in + Python control flow. It is a conservative gate for Phase 3 current reuse: + when it returns true, `seq_lens == extend_seq_lens` and all prefix lengths + are zero, so the logical KV view for the batch should be exactly the + current `out_cache_loc` chunk. + """ + + if forward_batch is None: + return False + forward_mode = getattr(forward_batch, "forward_mode", None) + if forward_mode is None or not forward_mode.is_extend_without_speculative(): + return False + + extend_prefix_lens_cpu = getattr(forward_batch, "extend_prefix_lens_cpu", None) + extend_seq_lens_cpu = getattr(forward_batch, "extend_seq_lens_cpu", None) + seq_lens_cpu = getattr(forward_batch, "seq_lens_cpu", None) + if ( + extend_prefix_lens_cpu is None + or extend_seq_lens_cpu is None + or seq_lens_cpu is None + ): + return False + + if any(int(prefix_len) != 0 for prefix_len in extend_prefix_lens_cpu): + return False + + seq_lens_list = [int(seq_len) for seq_len in seq_lens_cpu.tolist()] + extend_seq_lens_list = [int(seq_len) for seq_len in extend_seq_lens_cpu] + return seq_lens_list == extend_seq_lens_list + + def cp_shared_kv_debug_log( key: str, message: str, @@ -230,6 +268,49 @@ def remap_logical_locs_to_dense_locs( return dense_locs +def build_current_loc_remap( + query_locs: torch.Tensor, + current_locs: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Map logical locs into rows of the gathered current chunk tensor. + + `query_locs` may be any shape and may contain `-1` invalid sentinels. + `current_locs` is `forward_batch.out_cache_loc`; its row order is the row + order of the already CP-all-gathered current KV/index tensor. + + Returns: + - is_current_mask: true where query_locs is a non-negative current loc. + - compact_row_ids: row id into the current compact tensor where valid, + and -1 elsewhere. The dtype/shape match query_locs. + """ + + is_current = torch.zeros_like(query_locs, dtype=torch.bool) + compact_row_ids = torch.full_like(query_locs, -1) + if query_locs.numel() == 0 or current_locs.numel() == 0: + return is_current, compact_row_ids + + query_flat_long = query_locs.reshape(-1).to(torch.long) + current_flat_long = current_locs.reshape(-1).to(torch.long) + sorted_current_locs, sorted_to_current_rows = torch.sort(current_flat_long) + + insert_positions = torch.searchsorted(sorted_current_locs, query_flat_long) + safe_positions = torch.clamp(insert_positions, max=current_flat_long.numel() - 1) + in_range = insert_positions < current_flat_long.numel() + matched = ( + in_range + & (query_flat_long >= 0) + & (sorted_current_locs[safe_positions] == query_flat_long) + ) + + matched_rows = sorted_to_current_rows[safe_positions].to(compact_row_ids.dtype) + compact_flat = torch.where( + matched, + matched_rows, + torch.full_like(matched_rows, -1), + ) + return matched.reshape(query_locs.shape), compact_flat.reshape(query_locs.shape) + + def logical_pages_from_locs(logical_locs: torch.Tensor, page_size: int) -> torch.Tensor: logical_pages = logical_locs.clone() valid_mask = logical_locs >= 0 diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index b2e3e540b..b444ccc66 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -16,7 +16,9 @@ from sglang.srt.layers.attention.nsa import index_buf_accessor from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( cp_shared_kv_debug_enabled, cp_shared_kv_debug_log, + cp_shared_kv_current_reuse_enabled, filter_owned_logical_locs, + is_current_only_extend_batch, materialize_shared_paged_buffer, tensor_debug_checksum, tensor_debug_summary, @@ -319,6 +321,19 @@ class Indexer(MultiPlatformOp): return physical_out_loc, key[owned_mask].contiguous() + def _can_reuse_current_index_kv(self, forward_batch: ForwardBatch) -> bool: + return ( + cp_shared_kv_current_reuse_enabled() + and forward_batch.uses_cp_shared_kv + and self.nsa_enable_prefill_cp + and forward_batch.nsa_cp_metadata is not None + and is_nsa_prefill_cp_in_seq_split() + and is_current_only_extend_batch(forward_batch) + and forward_batch.hisparse_coordinator is None + and _is_cuda + and not _is_fp8_fnuz + ) + @contextlib.contextmanager def _with_real_sm_count(self): # When pipeline parallelism is enabled, each PP rank initiates a recv operation after the _pp_launch_batch @@ -822,6 +837,7 @@ class Indexer(MultiPlatformOp): kv_len: int, actual_seq_q: int, cp_index: List[Tuple[int, int, int]] = None, + current_index_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, ) -> torch.Tensor: if TYPE_CHECKING: assert isinstance(forward_batch.token_to_kv_pool, NSATokenToKVPool) @@ -839,11 +855,29 @@ class Indexer(MultiPlatformOp): batch_idx_list = [] block_tables = metadata.get_page_table_64() - index_buffer, block_tables = self._maybe_materialize_shared_index_buffer( - forward_batch, - layer_id, - block_tables, - ) + if current_index_kv is not None and cp_index is not None: + current_index_kv = None + if current_index_kv is None: + index_buffer, block_tables = self._maybe_materialize_shared_index_buffer( + forward_batch, + layer_id, + block_tables, + ) + else: + index_buffer = None + if cp_shared_kv_debug_enabled(): + cp_shared_kv_debug_log( + "index_current_reuse", + "NSA index current reuse cp_rank=%s layer=%s kv_len=%s actual_seq_q=%s k_ck=%s s_ck=%s", + forward_batch.cp_shared_kv_layout.cp_rank + if forward_batch.cp_shared_kv_layout is not None + else None, + layer_id, + kv_len, + actual_seq_q, + tensor_debug_checksum(current_index_kv[0]), + tensor_debug_checksum(current_index_kv[1]), + ) assert ( forward_batch.seq_lens_cpu is not None @@ -949,21 +983,27 @@ class Indexer(MultiPlatformOp): weights = weights[valid_q_mask].contiguous() ke_offset = ke_offset[valid_q_mask].contiguous() kv_len = min(cp_kv_end, logical_kv_limit) - k_fp8 = index_buf_accessor.GetK.execute( - forward_batch.token_to_kv_pool, - index_buffer, - seq_len=kv_len, - page_indices=block_tables[0], - ) - k_scale = index_buf_accessor.GetS.execute( - forward_batch.token_to_kv_pool, - index_buffer, - seq_len=kv_len, - page_indices=block_tables[0], - ) + if current_index_kv is None: + assert index_buffer is not None + k_fp8 = index_buf_accessor.GetK.execute( + forward_batch.token_to_kv_pool, + index_buffer, + seq_len=kv_len, + page_indices=block_tables[0], + ) + k_scale = index_buf_accessor.GetS.execute( + forward_batch.token_to_kv_pool, + index_buffer, + seq_len=kv_len, + page_indices=block_tables[0], + ) - k_fp8 = k_fp8.view(torch.float8_e4m3fn) - k_scale = k_scale.view(torch.float32).squeeze(-1) + k_fp8 = k_fp8.view(torch.float8_e4m3fn) + k_scale = k_scale.view(torch.float32).squeeze(-1) + else: + k_fp8, k_scale = current_index_kv + 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 @@ -1269,6 +1309,34 @@ class Indexer(MultiPlatformOp): weights = self._get_logits_head_gate(x_for_gate, q_scale) + current_index_kv = None + if self._can_reuse_current_index_kv(forward_batch): + if key.shape[0] == forward_batch.out_cache_loc.numel(): + current_k_fp8, current_k_scale = act_quant( + key.contiguous(), self.block_size, self.scale_fmt + ) + current_index_kv = ( + current_k_fp8.contiguous(), + current_k_scale.contiguous(), + ) + if cp_shared_kv_debug_enabled(): + cp_shared_kv_debug_log( + "index_current_reuse_prepare", + "NSA index current reuse prepared cp_rank=%s layer=%s locs=%s k_ck=%s s_ck=%s", + forward_batch.cp_shared_kv_layout.cp_rank + if forward_batch.cp_shared_kv_layout is not None + else None, + layer_id, + tensor_debug_summary(forward_batch.out_cache_loc), + tensor_debug_checksum(current_index_kv[0]), + tensor_debug_checksum(current_index_kv[1]), + ) + elif cp_shared_kv_debug_enabled(): + raise RuntimeError( + "CP shared KV current index reuse shape mismatch: " + f"key_tokens={key.shape[0]} out_cache_loc={forward_batch.out_cache_loc.numel()}" + ) + if _is_cuda or _is_hip: assert forward_batch.seq_lens_cpu is not None if len(forward_batch.seq_lens_cpu) == 0: @@ -1320,6 +1388,7 @@ class Indexer(MultiPlatformOp): metadata, kv_len_prev, actual_seq_q_prev, + current_index_kv=current_index_kv, ) topk_result_next = self._get_topk_ragged_with_cp( @@ -1330,6 +1399,7 @@ class Indexer(MultiPlatformOp): metadata, kv_len_next, actual_seq_q_next, + current_index_kv=current_index_kv, ) return torch.cat([topk_result_prev, topk_result_next], dim=0) else: diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index 0c981ba0d..dcba05c8f 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/python/sglang/srt/layers/attention/nsa_backend.py @@ -10,9 +10,12 @@ from sglang.srt.configs.model_config import get_nsa_index_topk, is_deepseek_nsa from sglang.srt.environ import envs from sglang.srt.layers.attention.base_attn_backend import AttentionBackend from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( + build_current_loc_remap, cp_shared_kv_debug_enabled, cp_shared_kv_debug_log, + cp_shared_kv_current_reuse_enabled, filter_owned_logical_locs, + is_current_only_extend_batch, materialize_shared_token_kv_buffer, tensor_debug_checksum, tensor_debug_summary, @@ -1591,13 +1594,49 @@ class NativeSparseAttnBackend( and topk_transform_method == TopkTransformMethod.PAGED ): assert forward_batch.cp_shared_kv_layout is not None - kv_cache, page_table_1 = materialize_shared_token_kv_buffer( - kv_cache=kv_cache, - logical_locs=page_table_1, - remap_logical_locs=metadata.page_table_1, - layout=forward_batch.cp_shared_kv_layout, - page_size=forward_batch.token_to_kv_pool.page_size, + can_reuse_current_kv = ( + cp_shared_kv_current_reuse_enabled() + and is_current_only_extend_batch(forward_batch) + and k is not None + and k_rope is not None + and k.shape[0] == forward_batch.out_cache_loc.numel() + and k_rope.shape[0] == forward_batch.out_cache_loc.numel() ) + if can_reuse_current_kv: + logical_page_table_1 = page_table_1 + current_mask, page_table_1 = build_current_loc_remap( + logical_page_table_1, + forward_batch.out_cache_loc, + ) + if cp_shared_kv_debug_enabled(): + missing_current = (logical_page_table_1 >= 0) & (~current_mask) + if torch.any(missing_current): + bad_locs = logical_page_table_1[missing_current] + raise RuntimeError( + "CP shared KV current MLA reuse expected current-only " + "logical locs but found history locs. " + f"bad_min={int(bad_locs.min().item())} " + f"bad_max={int(bad_locs.max().item())}" + ) + cp_shared_kv_debug_log( + "mla_current_reuse", + "MLA current reuse cp_rank=%s layer=%s current_locs=%s remapped=%s kv_ck=%s rope_ck=%s", + forward_batch.cp_shared_kv_layout.cp_rank, + layer.layer_id, + tensor_debug_summary(forward_batch.out_cache_loc), + tensor_debug_summary(page_table_1), + tensor_debug_checksum(k), + tensor_debug_checksum(k_rope), + ) + kv_cache = _cat([k, k_rope], dim=-1) + else: + kv_cache, page_table_1 = materialize_shared_token_kv_buffer( + kv_cache=kv_cache, + logical_locs=page_table_1, + remap_logical_locs=metadata.page_table_1, + layout=forward_batch.cp_shared_kv_layout, + page_size=forward_batch.token_to_kv_pool.page_size, + ) if nsa_impl == "tilelang": if q_rope is not None: 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 a08385837..cb9e7a1f0 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 @@ -104,6 +104,60 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): self.assertEqual(dense_locs.tolist(), [0, 1, 4, 5, -1, 8]) + def test_build_current_loc_remap_supports_non_contiguous_locs_and_sentinel(self): + from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( + build_current_loc_remap, + ) + + current_locs = torch.tensor([100, 64, 256, 128], dtype=torch.int64) + query_locs = torch.tensor( + [[128, -1, 64], [512, 100, 256]], dtype=torch.int32 + ) + + is_current, compact_rows = build_current_loc_remap(query_locs, current_locs) + + self.assertEqual( + is_current.tolist(), + [[True, False, True], [False, True, True]], + ) + self.assertEqual(compact_rows.tolist(), [[3, -1, 1], [-1, 0, 2]]) + self.assertEqual(compact_rows.dtype, query_locs.dtype) + + def test_build_current_loc_remap_returns_all_invalid_for_empty_current_locs(self): + from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( + build_current_loc_remap, + ) + + query_locs = torch.tensor([4, -1, 8], dtype=torch.int64) + is_current, compact_rows = build_current_loc_remap( + query_locs, torch.empty((0,), dtype=torch.int64) + ) + + self.assertEqual(is_current.tolist(), [False, False, False]) + self.assertEqual(compact_rows.tolist(), [-1, -1, -1]) + + def test_is_current_only_extend_batch_uses_cpu_lengths_without_tensor_scans(self): + from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( + is_current_only_extend_batch, + ) + from sglang.srt.model_executor.forward_batch_info import ForwardMode + + forward_batch = SimpleNamespace( + forward_mode=ForwardMode.EXTEND, + extend_prefix_lens_cpu=[0, 0], + extend_seq_lens_cpu=[3, 5], + seq_lens_cpu=torch.tensor([3, 5], dtype=torch.int32), + ) + + self.assertTrue(is_current_only_extend_batch(forward_batch)) + + forward_batch.extend_prefix_lens_cpu = [0, 1] + self.assertFalse(is_current_only_extend_batch(forward_batch)) + + forward_batch.extend_prefix_lens_cpu = [0, 0] + forward_batch.seq_lens_cpu = torch.tensor([4, 5], dtype=torch.int32) + self.assertFalse(is_current_only_extend_batch(forward_batch)) + def test_materialize_local_token_kv_pages(self): from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( build_dense_page_remap,