From c3fc3ff75269af0e416b662cd6274b8307b2b08f Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Fri, 29 May 2026 00:33:41 +0800 Subject: [PATCH] Stabilize CP shared-KV prefetch around draft cache hits Cache-hit EAGLE/NextN draft extends can enter the draft DeepEP MoE immediately after CP shared-KV attention. The partial current-reuse path is kept for target layers, but draft cache-hit suffixes now use full materialization until draft has an explicit same-layer reuse contract. Next-layer MLA/index prefetch is also gated by the actual model depth, so the single-layer draft model does not enqueue unused next-layer async work. The temporary stage traces used to isolate the hang are removed. The retained draft current-reuse fallback is a bounded warning because it changes the runtime path intentionally. Constraint: EAGLE/NextN has one executable draft layer and mirrors target KV state. Rejected: Keep partial current reuse for draft cache-hit suffixes | reproduced hangs at draft layer0 before DeepEP MoE completion. Rejected: Keep temporary stage traces | useful for diagnosis but too noisy for normal runs. Confidence: medium Scope-risk: moderate Directive: Do not re-enable draft cache-hit partial current reuse without an explicit draft same-layer reuse contract and ETE validation with CP shared KV + HiCache + EAGLE. Tested: py_compile on edited Python files; git diff --check; temp trace grep returned no matches. Not-tested: Local targeted pytest is blocked by missing pybase64 in this environment; full ETE after log cleanup not run. --- .../attention/nsa/cp_shared_kv_prefetch.py | 103 +++++ .../attention/nsa/cp_shared_kv_runtime.py | 176 ++++++++- .../srt/layers/attention/nsa/nsa_indexer.py | 3 + .../srt/layers/attention/nsa_backend.py | 193 +++++++--- .../attention_forward_methods/forward_mla.py | 14 +- python/sglang/srt/models/deepseek_v2.py | 28 +- .../mem_cache/test_cp_shared_kv_runtime.py | 362 +++++++++++++++++- 7 files changed, 808 insertions(+), 71 deletions(-) diff --git a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py index 71c6a0f79..76567bfe7 100644 --- a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py +++ b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py @@ -22,6 +22,7 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( get_or_build_shared_token_kv_slot_remap, materialize_local_paged_buffer_page_slots_into, materialize_local_token_kv_page_slots_into, + merge_materialized_and_current_kv, remap_logical_pages_to_slot_dense_pages, remap_logical_locs_to_slot_dense_locs_optimized, slot_range_to_page_slice, @@ -647,6 +648,108 @@ class CpSharedKVMlaPrefetcher: ) return dense_kv_cache, dense_locs + def consume_prefix_with_current( + self, + *, + layer_id: int, + kv_cache: torch.Tensor, + logical_locs: torch.Tensor, + current_kv_cache: torch.Tensor, + current_locs: torch.Tensor, + current_remap_page_size: int | None = None, + current_remap_logical_page_capacity: int | None = None, + ) -> Optional[tuple[torch.Tensor, torch.Tensor]]: + """Consume the prefetched prefix and append current-layer KV rows. + + This is the partial-current-reuse variant of :meth:`consume`: prefix + pages are already materialized/reduced by the prefetch stream, while + current/suffix pages are not copied from the shared pool. Current locs + in ``logical_locs`` are remapped to the appended current KV rows. + """ + + if self.disabled: + self._log_layer( + layer_id, + "consume_prefix_current_skip reason=disabled layer=%s", + layer_id, + ) + return None + self._log_layer( + layer_id, + "consume_prefix_current_enter layer=%s prefix_pages=%s total_slots=%s handles=%s", + layer_id, + self.prefix_pages, + self.total_slots, + _debug_handle_keys(layer_id, self.handles), + ) + handle = self.handles.get(layer_id) + if handle is None: + self._log_layer(layer_id, "consume_prefix_current_miss layer=%s", layer_id) + return None + if handle.event is None: + self.launch_pending_reduce() + handle = self.handles.get(layer_id) + if handle is None or handle.event is None: + self._log_layer( + layer_id, + "consume_prefix_current_miss reason=prefix_reduce_not_ready layer=%s", + layer_id, + ) + return None + handle = self.handles.pop(layer_id) + if self.pending_attention_handle is handle: + self.pending_attention_handle = None + if handle.layer_id != layer_id: + self.disabled = True + self._log_layer( + layer_id, + "consume_prefix_current_skip reason=layer_mismatch expected=%s actual=%s", + layer_id, + handle.layer_id, + ) + return None + + consume_cpu = _cpu_timing_start() + wait_cpu = _cpu_timing_start() + torch.cuda.current_stream().wait_event(handle.event) + wait_ms = _cpu_timing_ms(wait_cpu) + dense_kv_cache = handle.dense_kv_cache + + remap_cpu = _cpu_timing_start() + logical_locs = filter_locs_mappable_to_physical_pool( + logical_locs=logical_locs, + layout=self.layout, + physical_token_capacity=kv_cache.shape[0], + ) + dense_locs = remap_logical_locs_to_slot_dense_locs_optimized( + logical_locs, + page_inverse=self.page_inverse, + page_size=self.page_size, + ) + mixed_kv_cache, mixed_locs, _ = merge_materialized_and_current_kv( + materialized_kv_cache=dense_kv_cache, + materialized_dense_locs=dense_locs, + current_kv_cache=current_kv_cache, + logical_locs=logical_locs, + current_locs=current_locs, + page_size=current_remap_page_size, + logical_page_capacity=current_remap_logical_page_capacity, + ) + remap_ms = _cpu_timing_ms(remap_cpu) + total_ms = _cpu_timing_ms(consume_cpu) + self._log_layer( + layer_id, + "consume_prefix_current_hit layer=%s prefix_pages=%s dense_rows=%s current_rows=%s total_ms=%.3f wait_ms=%.3f remap_ms=%.3f", + layer_id, + self.prefix_pages, + int(dense_kv_cache.shape[0]), + int(current_kv_cache.shape[0]), + total_ms, + wait_ms, + remap_ms, + ) + return mixed_kv_cache, mixed_locs + def start_next_layer_prefix( self, *, 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 6d151318b..56f74ff77 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 @@ -18,6 +18,7 @@ logger = logging.getLogger(__name__) _DEBUG_LOG_COUNTS: dict[str, int] = {} _TAI_MATERIALIZE_FALLBACK_LOG_COUNTS: dict[str, int] = {} _TAI_FUSED_MLA_STORE_FALLBACK_LOG_COUNTS: dict[str, int] = {} +_CURRENT_REUSE_FALLBACK_LOG_COUNTS: dict[str, int] = {} _SLOT_REMAP_CACHE_LOG_COUNTS: dict[str, int] = {} _MLA_PREFETCH_LOG_PROBE_LAYER = 2 _MLA_PREFETCH_DEFAULT_MIN_PREFIX_TOKENS = max( @@ -70,7 +71,8 @@ def cp_shared_kv_mla_prefetch_min_prefix_pages( """Minimum prefix pages required to enable Phase8 prefetch. Negative env values mean "use the dynamic default": at least one page per CP - lane and, when the runtime page size is known, at least 1K prefix tokens. + lane and, when the runtime page size is known, at least the configured + prefix-token threshold. This keeps tiny cache-hit prefixes on the simpler synchronous materialize path where prefix prefetch launch/collective overhead can dominate. Set the env to 0 to disable the gate, or to a positive absolute page count for @@ -103,6 +105,47 @@ def cp_shared_kv_mla_prefetch_should_log_layer(layer_id: int) -> bool: return int(layer_id) == _MLA_PREFETCH_LOG_PROBE_LAYER +def cp_shared_kv_is_draft_input(forward_batch: Any) -> bool: + spec_info = getattr(forward_batch, "spec_info", None) + is_draft_input = getattr(spec_info, "is_draft_input", None) + return callable(is_draft_input) and is_draft_input() + + +def cp_shared_kv_should_prefetch_next_layer( + forward_batch: Any, + layer_id: int, +) -> bool: + """Return whether layer ``layer_id`` has a real next layer to prefetch. + + CP shared-KV prefetch is a next-layer optimization. Draft/NextN models reuse + decoder layer id 0 but have only one executable layer, so blindly prefetching + layer 1 creates unused async work before the draft MoE/DeepEP collective. + The decoder layer publishes the current model depth on the ForwardBatch for + this check. + """ + + if cp_shared_kv_is_draft_input(forward_batch): + return False + + num_model_layers = getattr(forward_batch, "cp_shared_kv_num_model_layers", None) + if num_model_layers is None: + return True + return int(layer_id) + 1 < int(num_model_layers) + + +def _log_current_reuse_fallback( + key: str, + message: str, + *args, + limit: int = 8, +) -> None: + count = _CURRENT_REUSE_FALLBACK_LOG_COUNTS.get(key, 0) + if count >= limit: + return + _CURRENT_REUSE_FALLBACK_LOG_COUNTS[key] = count + 1 + logger.warning(message, *args) + + @dataclass(frozen=True) class SharedTokenKVSlotRemap: slot_logical_pages: torch.Tensor @@ -723,6 +766,137 @@ def is_current_only_extend_batch(forward_batch) -> bool: return seq_lens_list == extend_seq_lens_list +def can_reuse_current_extend_kv(forward_batch) -> bool: + """Return whether the current extend chunk can be used as dense KV rows. + + Unlike :func:`is_current_only_extend_batch`, this allows a cached/history + prefix. The contract is intentionally narrow for now: single-batch extend, + CPU length metadata present, and ``out_cache_loc`` exactly covers the current + extend chunk. The caller still owns model/backend gates such as CP shared KV + enabled, in-seq-split mode, and tensor shape compatibility. + """ + + 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 + + if int(getattr(forward_batch, "batch_size", 0)) != 1: + return False + + extend_seq_lens_cpu = getattr(forward_batch, "extend_seq_lens_cpu", None) + seq_lens_cpu = getattr(forward_batch, "seq_lens_cpu", None) + out_cache_loc = getattr(forward_batch, "out_cache_loc", None) + if extend_seq_lens_cpu is None or seq_lens_cpu is None or out_cache_loc is None: + return False + if len(extend_seq_lens_cpu) != 1 or int(seq_lens_cpu.numel()) != 1: + return False + + extend_len = int(extend_seq_lens_cpu[0]) + seq_len = int(seq_lens_cpu[0].item()) + if extend_len <= 0 or seq_len < extend_len: + return False + return int(out_cache_loc.numel()) == extend_len + + +def should_reuse_current_extend_kv(forward_batch) -> bool: + """Return whether MLA should splice current extend KV into materialized KV. + + Partial current reuse appends the freshly computed suffix KV to the + materialized prefix buffer. That path is safe for the target model, but the + EAGLE/NextN draft layer has a different lifetime contract: it mirrors target + KV state and immediately enters DeepEP MoE after attention. A cache-hit + draft suffix observed in production can leave all ranks stuck in that MoE + collective after the partial-reuse attention path returns. Keep draft + cache-hit suffixes on the older full-materialize path until draft gets an + explicit same-layer reuse contract. + """ + + if not cp_shared_kv_current_reuse_enabled(): + return False + + current_only = is_current_only_extend_batch(forward_batch) + partial_current = can_reuse_current_extend_kv(forward_batch) + if not (current_only or partial_current): + return False + + if cp_shared_kv_is_draft_input(forward_batch) and not current_only: + _log_current_reuse_fallback( + "draft_partial_current_reuse_disabled", + "CP shared KV current-reuse fallback (draft_partial_current_reuse): " + "cache-hit EAGLE/NextN draft uses full materialize instead of " + "partial current reuse. prefix_lens=%s extend_lens=%s", + getattr(forward_batch, "extend_prefix_lens_cpu", None), + getattr(forward_batch, "extend_seq_lens_cpu", None), + ) + return False + + return True + + +def current_loc_remap_fast_path_args( + forward_batch, +) -> tuple[int | None, int | None]: + """Return page-inverse remap args when the current chunk is page aligned. + + The page-inverse path in :func:`build_current_loc_remap` assumes row 0 of + ``current_locs`` is the first token of a logical page. That is true for the + existing current-only extend path, but not guaranteed for cache-hit partial + extend where the cached prefix may end mid-page. For partial extend, return + ``(None, None)`` so callers take the general sort/search remap path. + """ + + if not is_current_only_extend_batch(forward_batch): + return None, None + if len(getattr(forward_batch, "extend_seq_lens_cpu", []) or []) != 1: + return None, None + token_to_kv_pool = getattr(forward_batch, "token_to_kv_pool", None) + layout = getattr(forward_batch, "cp_shared_kv_layout", None) + if token_to_kv_pool is None or layout is None: + return None, None + page_size = int(getattr(token_to_kv_pool, "page_size", 0)) + if page_size <= 0: + return None, None + logical_page_capacity = ( + max(int(getattr(token_to_kv_pool, "size", 0)) // page_size - 1, 0) + * int(layout.cp_size) + + 1 + ) + return page_size, logical_page_capacity + + +def merge_materialized_and_current_kv( + *, + materialized_kv_cache: torch.Tensor, + materialized_dense_locs: torch.Tensor, + current_kv_cache: torch.Tensor, + logical_locs: torch.Tensor, + current_locs: torch.Tensor, + page_size: int | None = None, + logical_page_capacity: int | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Append current KV rows to a materialized prefix buffer and remap locs. + + ``materialized_dense_locs`` is the remap returned by the prefix/history + materialization path. Entries corresponding to current extend tokens are + replaced with offsets into the appended ``current_kv_cache``. Non-current + entries remain untouched, including ``-1`` invalid sentinels. + """ + + current_mask, current_rows = build_current_loc_remap( + logical_locs, + current_locs, + page_size=page_size, + logical_page_capacity=logical_page_capacity, + ) + current_offset = int(materialized_kv_cache.shape[0]) + current_dense_locs = current_rows + current_offset + mixed_locs = torch.where(current_mask, current_dense_locs, materialized_dense_locs) + mixed_kv_cache = torch.cat([materialized_kv_cache, current_kv_cache], dim=0) + return mixed_kv_cache, mixed_locs, current_mask + + def cp_shared_kv_debug_log( key: str, message: str, diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index 36d506d5e..c36bea641 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -21,6 +21,7 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( cp_shared_kv_mla_prefetch_log, cp_shared_kv_mla_prefetch_log_enabled, cp_shared_kv_mla_prefetch_should_log_layer, + cp_shared_kv_should_prefetch_next_layer, filter_owned_logical_locs, get_or_build_shared_paged_buffer_slot_remap, is_current_only_extend_batch, @@ -405,6 +406,8 @@ class Indexer(MultiPlatformOp): ) if index_prefetcher is None: return + if not cp_shared_kv_should_prefetch_next_layer(forward_batch, layer_id): + return index_prefetcher.start_next_layer_prefix( next_layer_id=next_layer_id, token_to_kv_pool=forward_batch.token_to_kv_pool, diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index 76be1bbda..9b77be7eb 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/python/sglang/srt/layers/attention/nsa_backend.py @@ -17,14 +17,18 @@ 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, cp_shared_kv_mla_prefetch_log, cp_shared_kv_mla_prefetch_log_enabled, cp_shared_kv_mla_prefetch_should_log_layer, + cp_shared_kv_is_draft_input, + cp_shared_kv_should_prefetch_next_layer, + current_loc_remap_fast_path_args, filter_owned_logical_locs, get_or_build_shared_token_kv_slot_remap, is_current_only_extend_batch, materialize_shared_token_kv_buffer, + merge_materialized_and_current_kv, + should_reuse_current_extend_kv, tensor_debug_checksum, tensor_debug_summary, ) @@ -878,12 +882,21 @@ class NativeSparseAttnBackend( token_to_batch_idx=token_to_batch_idx, ) self.forward_metadata = metadata - mla_prefetcher = CpSharedKVMlaPrefetcher.maybe_create( - forward_batch=forward_batch, - metadata=metadata, - topk_transform_is_paged=( - topk_transform_method == TopkTransformMethod.PAGED - ), + # EAGLE/NextN draft has a single executable layer. The current + # CP-shared prefetch pipeline is a target-model next-layer optimization; + # keep it enabled for target extend/verify, but do not create draft + # prefetchers until EAGLE gets an explicit same-layer prefetch contract. + disable_draft_prefetch = cp_shared_kv_is_draft_input(forward_batch) + mla_prefetcher = ( + None + if disable_draft_prefetch + else CpSharedKVMlaPrefetcher.maybe_create( + forward_batch=forward_batch, + metadata=metadata, + topk_transform_is_paged=( + topk_transform_method == TopkTransformMethod.PAGED + ), + ) ) forward_batch.cp_shared_kv_mla_prefetcher = mla_prefetcher # Use one FIFO stream for index and MLA prefix prefetch. Both paths @@ -896,7 +909,9 @@ class NativeSparseAttnBackend( else None ) forward_batch.cp_shared_kv_index_prefetcher = ( - CpSharedKVIndexPrefetcher.maybe_create( + None + if disable_draft_prefetch + else CpSharedKVIndexPrefetcher.maybe_create( forward_batch=forward_batch, metadata=metadata, topk_transform_is_paged=( @@ -1726,8 +1741,7 @@ class NativeSparseAttnBackend( forward_batch, "cp_shared_kv_mla_prefetcher", None ) can_reuse_current_kv = ( - cp_shared_kv_current_reuse_enabled() - and is_current_only_extend_batch(forward_batch) + should_reuse_current_extend_kv(forward_batch) and k is not None and k_rope is not None and k.shape[0] == forward_batch.out_cache_loc.numel() @@ -1767,48 +1781,122 @@ class NativeSparseAttnBackend( else None, ) if can_reuse_current_kv: + current_kv_cache = _cat([k, k_rope], dim=-1) logical_page_table_1 = page_table_1 - current_remap_page_size = None - current_remap_logical_page_capacity = None - if len(getattr(forward_batch, "extend_seq_lens_cpu", []) or []) == 1: - current_remap_page_size = forward_batch.token_to_kv_pool.page_size - current_remap_logical_page_capacity = ( - max( - forward_batch.token_to_kv_pool.size - // current_remap_page_size - - 1, - 0, - ) - * forward_batch.cp_shared_kv_layout.cp_size - + 1 - ) - current_mask, page_table_1 = build_current_loc_remap( - logical_page_table_1, - forward_batch.out_cache_loc, - page_size=current_remap_page_size, - logical_page_capacity=current_remap_logical_page_capacity, + current_remap_page_size, current_remap_logical_page_capacity = ( + current_loc_remap_fast_path_args(forward_batch) ) - 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), + + if is_current_only_extend_batch(forward_batch): + current_mask, page_table_1 = build_current_loc_remap( + logical_page_table_1, + forward_batch.out_cache_loc, + page_size=current_remap_page_size, + logical_page_capacity=current_remap_logical_page_capacity, ) - kv_cache = _cat([k, k_rope], dim=-1) + 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 = current_kv_cache + else: + prefetched_kv = None + if mla_prefetcher is not None: + prefetched_kv = mla_prefetcher.consume_prefix_with_current( + layer_id=layer.layer_id, + kv_cache=kv_cache, + logical_locs=logical_page_table_1, + current_kv_cache=current_kv_cache, + current_locs=forward_batch.out_cache_loc, + current_remap_page_size=current_remap_page_size, + current_remap_logical_page_capacity=current_remap_logical_page_capacity, + ) + if prefetched_kv is not None: + kv_cache, page_table_1 = prefetched_kv + else: + current_mask, _ = build_current_loc_remap( + logical_page_table_1, + forward_batch.out_cache_loc, + page_size=current_remap_page_size, + logical_page_capacity=current_remap_logical_page_capacity, + ) + materialize_locs = torch.where( + current_mask, + torch.full_like(logical_page_table_1, -1), + logical_page_table_1, + ) + prefix_kv_cache, prefix_dense_locs = ( + materialize_shared_token_kv_buffer( + kv_cache=kv_cache, + logical_locs=materialize_locs, + layout=forward_batch.cp_shared_kv_layout, + page_size=forward_batch.token_to_kv_pool.page_size, + nvtx_source="mla.partial_current_materialize", + nvtx_layer_id=layer.layer_id, + ) + ) + kv_cache, page_table_1, _ = merge_materialized_and_current_kv( + materialized_kv_cache=prefix_kv_cache, + materialized_dense_locs=prefix_dense_locs, + current_kv_cache=current_kv_cache, + logical_locs=logical_page_table_1, + current_locs=forward_batch.out_cache_loc, + page_size=current_remap_page_size, + logical_page_capacity=current_remap_logical_page_capacity, + ) + if ( + cp_shared_kv_mla_prefetch_log_enabled() + and cp_shared_kv_mla_prefetch_should_log_layer(layer.layer_id) + ): + prefix_lens_cpu = getattr( + forward_batch, "extend_prefix_lens_cpu", None + ) + extend_lens_cpu = getattr( + forward_batch, "extend_seq_lens_cpu", None + ) + cp_shared_kv_mla_prefetch_log( + "forward_partial_current_reuse cp_rank=%s layer=%s used_prefetch=%s " + "prefix_lens=%s extend_lens=%s current_rows=%s kv_rows=%s page_table_shape=%s", + forward_batch.cp_shared_kv_layout.cp_rank, + layer.layer_id, + prefetched_kv is not None, + [int(x) for x in prefix_lens_cpu] + if prefix_lens_cpu is not None + else None, + [int(x) for x in extend_lens_cpu] + if extend_lens_cpu is not None + else None, + int(current_kv_cache.shape[0]), + int(kv_cache.shape[0]), + tuple(page_table_1.shape), + ) + if cp_shared_kv_debug_enabled(): + cp_shared_kv_debug_log( + "mla_partial_current_reuse", + "MLA partial 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), + ) else: prefetched_kv = None if mla_prefetcher is not None: @@ -1838,8 +1926,9 @@ class NativeSparseAttnBackend( nvtx_source="mla.full_materialize", nvtx_layer_id=layer.layer_id, ) - - if mla_prefetcher is not None: + if mla_prefetcher is not None and cp_shared_kv_should_prefetch_next_layer( + forward_batch, layer.layer_id + ): mla_prefetcher.start_next_layer_prefix( next_layer_id=layer.layer_id + 1, token_to_kv_pool=forward_batch.token_to_kv_pool, @@ -1852,7 +1941,6 @@ class NativeSparseAttnBackend( ) if index_prefetcher is not None: index_prefetcher.launch_pending_reduce() - try: if nsa_impl == "tilelang": if q_rope is not None: @@ -1950,7 +2038,6 @@ class NativeSparseAttnBackend( ) if index_prefetcher is not None: index_prefetcher.wait_attention_window() - return attn_output def forward_decode( diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py index b57b2ba3a..c0df53737 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py @@ -14,8 +14,8 @@ from sglang.srt.layers.attention.nsa.utils import ( nsa_use_prefill_cp, ) from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( - cp_shared_kv_current_reuse_enabled, - is_current_only_extend_batch, + cp_shared_kv_should_prefetch_next_layer, + should_reuse_current_extend_kv, try_tai_fused_mla_store, ) from sglang.srt.layers.communicator import get_attn_tp_context @@ -114,6 +114,11 @@ class DeepseekMLAForwardMixin: if token_to_kv_pool is None: return + if not cp_shared_kv_should_prefetch_next_layer( + forward_batch, self.layer_id + ): + return + next_layer_id = int(self.layer_id) + 1 index_prefetcher = getattr( forward_batch, "cp_shared_kv_index_prefetcher", None @@ -362,9 +367,8 @@ class DeepseekMLAForwardMixin: shared_mla_direct_write_done and getattr(forward_batch, "uses_cp_shared_kv", False) ) - current_reuse_needs_full_current_kv = ( - cp_shared_kv_current_reuse_enabled() - and is_current_only_extend_batch(forward_batch) + current_reuse_needs_full_current_kv = should_reuse_current_extend_kv( + forward_batch ) if ( not shared_kv_materialize_will_read_pool diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 07b30ae98..1a6e3ef87 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -1680,14 +1680,28 @@ class DeepseekV2DecoderLayer(nn.Module): quant_format, ) - hidden_states = self.self_attn( - positions=positions, - hidden_states=hidden_states, - forward_batch=forward_batch, - zero_allocator=zero_allocator, - llama_4_scaling=llama_4_scaling, - layer_scatter_modes=self.layer_scatter_modes, + previous_cp_shared_kv_num_model_layers = getattr( + forward_batch, "cp_shared_kv_num_model_layers", None ) + forward_batch.cp_shared_kv_num_model_layers = ( + 1 if self.is_nextn else self.config.num_hidden_layers + ) + try: + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + zero_allocator=zero_allocator, + llama_4_scaling=llama_4_scaling, + layer_scatter_modes=self.layer_scatter_modes, + ) + finally: + if previous_cp_shared_kv_num_model_layers is None: + delattr(forward_batch, "cp_shared_kv_num_model_layers") + else: + forward_batch.cp_shared_kv_num_model_layers = ( + previous_cp_shared_kv_num_model_layers + ) hidden_states, residual = self.layer_communicator.prepare_mlp( hidden_states, residual, forward_batch 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 b50db01f5..80b6e94ce 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 @@ -25,6 +25,14 @@ for _name in ("flash_attn_varlen_func", "flash_attn_with_kvcache"): if not hasattr(flash_attn_stub, _name): setattr(flash_attn_stub, _name, lambda *args, **kwargs: None) +sgl_kernel_stub = sys.modules.setdefault( + "sgl_kernel", types.ModuleType("sgl_kernel") +) +if not hasattr(sgl_kernel_stub, "__path__"): + sgl_kernel_stub.__path__ = [] +if not hasattr(sgl_kernel_stub, "flash_attn"): + sgl_kernel_stub.flash_attn = flash_attn_stub + from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=1, suite="stage-a-test-cpu") @@ -34,6 +42,11 @@ def _identity_all_reduce(buffer, *args, **kwargs): return buffer +class _FakeExtendForwardMode: + def is_extend_without_speculative(self): + return True + + class TestCpSharedKVRuntimeHelpers(unittest.TestCase): def test_mla_prefetch_materializes_and_reduces_on_prefetch_stream( self, @@ -396,10 +409,9 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): 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, + forward_mode=_FakeExtendForwardMode(), extend_prefix_lens_cpu=[0, 0], extend_seq_lens_cpu=[3, 5], seq_lens_cpu=torch.tensor([3, 5], dtype=torch.int32), @@ -414,6 +426,299 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): forward_batch.seq_lens_cpu = torch.tensor([4, 5], dtype=torch.int32) self.assertFalse(is_current_only_extend_batch(forward_batch)) + def test_can_reuse_current_extend_kv_allows_partial_cache_hit_single_batch(self): + from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( + can_reuse_current_extend_kv, + ) + + forward_batch = SimpleNamespace( + forward_mode=_FakeExtendForwardMode(), + batch_size=1, + extend_seq_lens_cpu=[128], + seq_lens_cpu=torch.tensor([40384 + 128], dtype=torch.int32), + out_cache_loc=torch.arange(128, dtype=torch.int64), + ) + + self.assertTrue(can_reuse_current_extend_kv(forward_batch)) + + forward_batch.batch_size = 2 + self.assertFalse(can_reuse_current_extend_kv(forward_batch)) + + forward_batch.batch_size = 1 + forward_batch.out_cache_loc = torch.arange(127, dtype=torch.int64) + self.assertFalse(can_reuse_current_extend_kv(forward_batch)) + + def test_should_reuse_current_extend_kv_disables_draft_cache_hit_suffix(self): + from sglang.srt.environ import envs + from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime + + class DraftSpecInfo: + def is_draft_input(self): + return True + + class TargetSpecInfo: + def is_draft_input(self): + return False + + runtime._CURRENT_REUSE_FALLBACK_LOG_COUNTS.clear() + forward_batch = SimpleNamespace( + forward_mode=_FakeExtendForwardMode(), + batch_size=1, + extend_prefix_lens_cpu=[40384], + extend_seq_lens_cpu=[56], + seq_lens_cpu=torch.tensor([40384 + 56], dtype=torch.int32), + out_cache_loc=torch.arange(56, dtype=torch.int64), + spec_info=DraftSpecInfo(), + ) + + with envs.SGLANG_CP_SHARED_KV_CURRENT_REUSE.override(True): + with self.assertLogs(runtime.logger.name, level="WARNING") as logs: + self.assertFalse(runtime.should_reuse_current_extend_kv(forward_batch)) + self.assertTrue( + any( + "draft_partial_current_reuse" in message + for message in logs.output + ) + ) + + forward_batch.spec_info = TargetSpecInfo() + self.assertTrue(runtime.should_reuse_current_extend_kv(forward_batch)) + + forward_batch.spec_info = DraftSpecInfo() + forward_batch.extend_prefix_lens_cpu = [0] + forward_batch.seq_lens_cpu = torch.tensor([56], dtype=torch.int32) + self.assertTrue(runtime.should_reuse_current_extend_kv(forward_batch)) + + def test_current_loc_remap_fast_path_args_only_for_current_only_extend(self): + from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( + current_loc_remap_fast_path_args, + ) + from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout + + forward_batch = SimpleNamespace( + forward_mode=_FakeExtendForwardMode(), + batch_size=1, + extend_prefix_lens_cpu=[0], + extend_seq_lens_cpu=[128], + seq_lens_cpu=torch.tensor([128], dtype=torch.int32), + out_cache_loc=torch.arange(128, dtype=torch.int64), + token_to_kv_pool=SimpleNamespace(page_size=64, size=4096), + cp_shared_kv_layout=CpSharedKVLayout(page_size=64, cp_size=8, cp_rank=0), + ) + + self.assertEqual(current_loc_remap_fast_path_args(forward_batch), (64, 505)) + + forward_batch.extend_prefix_lens_cpu = [40389] + forward_batch.seq_lens_cpu = torch.tensor([40389 + 128], dtype=torch.int32) + self.assertEqual(current_loc_remap_fast_path_args(forward_batch), (None, None)) + + def test_merge_materialized_and_current_kv_remaps_only_current_locs(self): + from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( + merge_materialized_and_current_kv, + ) + + materialized_kv = torch.arange(0, 8, dtype=torch.float32).view(8, 1, 1) + current_kv = torch.arange(100, 103, dtype=torch.float32).view(3, 1, 1) + logical_locs = torch.tensor([[4, 20, -1], [21, 7, 99]], dtype=torch.int32) + materialized_locs = torch.tensor([[4, -1, -1], [-1, 7, -1]], dtype=torch.int32) + current_locs = torch.tensor([20, 21, 22], dtype=torch.int64) + + mixed_kv, mixed_locs, current_mask = merge_materialized_and_current_kv( + materialized_kv_cache=materialized_kv, + materialized_dense_locs=materialized_locs, + current_kv_cache=current_kv, + logical_locs=logical_locs, + current_locs=current_locs, + ) + + self.assertTrue(torch.equal(mixed_kv[:8], materialized_kv)) + self.assertTrue(torch.equal(mixed_kv[8:], current_kv)) + self.assertEqual( + current_mask.tolist(), + [[False, True, False], [True, False, False]], + ) + self.assertEqual(mixed_locs.tolist(), [[4, 8, -1], [9, 7, -1]]) + + def test_mla_prefetch_consume_prefix_with_current_skips_suffix_materialize(self): + from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch + from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout + + class FakeCurrentStream: + def __init__(self): + self.events = [] + + def wait_event(self, event): + self.events.append(event) + + current_stream = FakeCurrentStream() + fake_event = object() + dense_kv = torch.arange(0, 16, dtype=torch.float32).view(16, 1, 1) + current_kv = torch.arange(100, 102, dtype=torch.float32).view(2, 1, 1) + page_inverse = torch.tensor([0, 1, 2, -1, -1, 3], dtype=torch.int64) + prefetcher = prefetch.CpSharedKVMlaPrefetcher( + layout=CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=0), + page_size=4, + prefix_pages=2, + slot_logical_pages=torch.tensor([1, 2, 5], dtype=torch.int64), + page_inverse=page_inverse, + dense_num_pages=4, + stream=object(), + ) + handle = prefetch.CpSharedKVMlaPrefetchHandle( + layer_id=1, + dense_kv_cache=dense_kv, + prefix_rows=slice(4, 12), + event=fake_event, + ) + prefetcher.handles[1] = handle + prefetcher.pending_attention_handle = handle + + with patch.object( + prefetch.torch.cuda, "current_stream", return_value=current_stream + ), patch.object( + prefetch, + "materialize_local_token_kv_page_slots_into", + side_effect=AssertionError("suffix materialize must not run"), + ): + mixed_kv, mixed_locs = prefetcher.consume_prefix_with_current( + layer_id=1, + kv_cache=torch.zeros((64, 1, 1), dtype=torch.float32), + logical_locs=torch.tensor([[4, 20], [21, 7]], dtype=torch.int32), + current_kv_cache=current_kv, + current_locs=torch.tensor([20, 21], dtype=torch.int64), + ) + + self.assertEqual(current_stream.events, [fake_event]) + self.assertEqual(prefetcher.handles, {}) + self.assertIsNone(prefetcher.pending_attention_handle) + self.assertTrue(torch.equal(mixed_kv[:16], dense_kv)) + self.assertTrue(torch.equal(mixed_kv[16:], current_kv)) + self.assertEqual(mixed_locs.tolist(), [[4, 16], [17, 7]]) + + def test_mla_prefetch_attention_window_waits_on_pending_event(self): + from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch + from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout + + class FakeCurrentStream: + def __init__(self): + self.events = [] + + def wait_event(self, event): + self.events.append(event) + + current_stream = FakeCurrentStream() + fake_event = object() + prefetcher = prefetch.CpSharedKVMlaPrefetcher( + layout=CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=0), + page_size=4, + prefix_pages=2, + slot_logical_pages=torch.tensor([1, 2, 5], dtype=torch.int64), + page_inverse=torch.tensor([0, 1, 2, -1, -1, 3], dtype=torch.int64), + dense_num_pages=4, + stream=object(), + ) + handle = prefetch.CpSharedKVMlaPrefetchHandle( + layer_id=1, + dense_kv_cache=torch.zeros((16, 1, 1), dtype=torch.float32), + prefix_rows=slice(4, 12), + event=fake_event, + ) + prefetcher.handles[1] = handle + prefetcher.pending_attention_handle = handle + + with patch.object( + prefetch.torch.cuda, "current_stream", return_value=current_stream + ): + prefetcher.wait_attention_window() + + self.assertEqual(current_stream.events, [fake_event]) + self.assertIsNone(prefetcher.pending_attention_handle) + self.assertIs(prefetcher.handles[1], handle) + + def test_mla_prefetch_attention_window_launches_pending_reduce_before_wait(self): + from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch + from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout + + class FakeCurrentStream: + def __init__(self): + self.events = [] + + def wait_event(self, event): + self.events.append(event) + + current_stream = FakeCurrentStream() + fake_event = object() + prefetcher = prefetch.CpSharedKVMlaPrefetcher( + layout=CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=0), + page_size=4, + prefix_pages=2, + slot_logical_pages=torch.tensor([1, 2, 5], dtype=torch.int64), + page_inverse=torch.tensor([0, 1, 2, -1, -1, 3], dtype=torch.int64), + dense_num_pages=4, + stream=object(), + ) + handle = prefetch.CpSharedKVMlaPrefetchHandle( + layer_id=1, + dense_kv_cache=torch.zeros((16, 1, 1), dtype=torch.float32), + prefix_rows=slice(4, 12), + event=None, + ) + prefetcher.handles[1] = handle + prefetcher.pending_attention_handle = handle + + def finish_reduce(): + handle.event = fake_event + + with patch.object( + prefetch.torch.cuda, "current_stream", return_value=current_stream + ), patch.object( + prefetcher, "launch_pending_reduce", side_effect=finish_reduce + ) as launch_pending_reduce: + prefetcher.wait_attention_window() + + launch_pending_reduce.assert_called_once_with() + self.assertEqual(current_stream.events, [fake_event]) + self.assertIsNone(prefetcher.pending_attention_handle) + + def test_index_prefetch_attention_window_waits_on_pending_event(self): + from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch + from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout + + class FakeCurrentStream: + def __init__(self): + self.events = [] + + def wait_event(self, event): + self.events.append(event) + + current_stream = FakeCurrentStream() + fake_event = object() + prefetcher = prefetch.CpSharedKVIndexPrefetcher( + layout=CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=0), + prefix_pages=2, + slot_logical_pages=torch.tensor([1, 2, 5], dtype=torch.int64), + page_inverse=torch.tensor([0, 1, 2, -1, -1, 3], dtype=torch.int64), + dense_num_pages=4, + stream=object(), + ) + handle = prefetch.CpSharedKVIndexPrefetchHandle( + layer_id=1, + dense_page_buffer=torch.zeros((4, 3), dtype=torch.uint8), + prefix_rows=slice(1, 3), + event=fake_event, + ) + prefetcher.handles[1] = handle + prefetcher.pending_attention_handle = handle + + with patch.object( + prefetch.torch.cuda, "current_stream", return_value=current_stream + ): + prefetcher.wait_attention_window() + + self.assertEqual(current_stream.events, [fake_event]) + self.assertIsNone(prefetcher.pending_attention_handle) + self.assertIs(prefetcher.handles[1], handle) + def test_materialize_local_token_kv_pages(self): from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( build_dense_page_remap, @@ -718,14 +1023,17 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): with envs.SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH.override(True): self.assertTrue(cp_shared_kv_mla_prefetch_log_enabled()) - def test_mla_prefetch_min_prefix_pages_defaults_to_1k_tokens_and_can_override(self): + def test_mla_prefetch_min_prefix_pages_uses_cached_token_default_and_can_override(self): from sglang.srt.environ import envs from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime envs.SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_PAGES.clear() - self.assertEqual(runtime._MLA_PREFETCH_DEFAULT_MIN_PREFIX_TOKENS, 1024) + default_tokens = envs.SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_TOKENS.get() + self.assertEqual(runtime._MLA_PREFETCH_DEFAULT_MIN_PREFIX_TOKENS, default_tokens) + expected_pages = (default_tokens + 63) // 64 self.assertEqual( - runtime.cp_shared_kv_mla_prefetch_min_prefix_pages(8, page_size=64), 16 + runtime.cp_shared_kv_mla_prefetch_min_prefix_pages(8, page_size=64), + max(8, expected_pages), ) self.assertEqual( runtime.cp_shared_kv_mla_prefetch_min_prefix_pages(32, page_size=64), 32 @@ -1823,6 +2131,50 @@ class TestCpSharedKVTaiMaterializeIntegration(unittest.TestCase): self.assertEqual(fake_prefetcher.calls, [(12, token_to_kv_pool)]) + def test_index_prefetch_skips_when_current_layer_is_last(self): + from sglang.srt.layers.attention.nsa import nsa_indexer + from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( + cp_shared_kv_should_prefetch_next_layer, + ) + + class FakePrefetcher: + def __init__(self): + self.calls = [] + + def start_next_layer_prefix(self, *, next_layer_id, token_to_kv_pool): + self.calls.append((next_layer_id, token_to_kv_pool)) + + token_to_kv_pool = object() + fake_prefetcher = FakePrefetcher() + forward_batch = SimpleNamespace( + token_to_kv_pool=token_to_kv_pool, + cp_shared_kv_index_prefetcher=fake_prefetcher, + cp_shared_kv_num_model_layers=12, + ) + indexer = object.__new__(nsa_indexer.Indexer) + + self.assertFalse(cp_shared_kv_should_prefetch_next_layer(forward_batch, 11)) + indexer._maybe_start_next_layer_index_prefetch(forward_batch, layer_id=11) + + self.assertEqual(fake_prefetcher.calls, []) + + def test_index_prefetch_skips_eagle_draft_next_layer(self): + from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( + cp_shared_kv_is_draft_input, + cp_shared_kv_should_prefetch_next_layer, + ) + + class FakeSpecInfo: + def is_draft_input(self): + return True + + forward_batch = SimpleNamespace( + spec_info=FakeSpecInfo(), + ) + + self.assertTrue(cp_shared_kv_is_draft_input(forward_batch)) + self.assertFalse(cp_shared_kv_should_prefetch_next_layer(forward_batch, 0)) + def test_index_prefetch_consume_miss_logs_fallback_after_first_layer(self): from sglang.srt.layers.attention.nsa import nsa_indexer