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.

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,
         *,

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,

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,

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(
+        # EAGLE/NextN draft has a single executable layer.  The current
-            forward_batch=forward_batch,
+        # CP-shared prefetch pipeline is a target-model next-layer optimization;
-            metadata=metadata,
+        # keep it enabled for target extend/verify, but do not create draft
-            topk_transform_is_paged=(
+        # prefetchers until EAGLE gets an explicit same-layer prefetch contract.
-                topk_transform_method == TopkTransformMethod.PAGED
+        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()
+                should_reuse_current_extend_kv(forward_batch)
-                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()
@@ -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_page_size, current_remap_logical_page_capacity = (
-                current_remap_logical_page_capacity = None
+                    current_loc_remap_fast_path_args(forward_batch)
-                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,
                 )
-                if cp_shared_kv_debug_enabled():
+
-                    missing_current = (logical_page_table_1 >= 0) & (~current_mask)
+                if is_current_only_extend_batch(forward_batch):
-                    if torch.any(missing_current):
+                    current_mask, page_table_1 = build_current_loc_remap(
-                        bad_locs = logical_page_table_1[missing_current]
+                        logical_page_table_1,
-                        raise RuntimeError(
+                        forward_batch.out_cache_loc,
-                            "CP shared KV current MLA reuse expected current-only "
+                        page_size=current_remap_page_size,
-                            "logical locs but found history locs. "
+                        logical_page_capacity=current_remap_logical_page_capacity,
-                            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)
+                    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 and cp_shared_kv_should_prefetch_next_layer(
-            if mla_prefetcher is not None:
+                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(

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,
+    cp_shared_kv_should_prefetch_next_layer,
-    is_current_only_extend_batch,
+    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 = (
+            current_reuse_needs_full_current_kv = should_reuse_current_extend_kv(
-                cp_shared_kv_current_reuse_enabled()
+                forward_batch
-                and is_current_only_extend_batch(forward_batch)
             )
             if (
                 not shared_kv_materialize_will_read_pool

python/sglang/srt/models/deepseek_v2.py

@@ -1680,14 +1680,28 @@ class DeepseekV2DecoderLayer(nn.Module):
             quant_format,
         )
 
-        hidden_states = self.self_attn(
+        previous_cp_shared_kv_num_model_layers = getattr(
-            positions=positions,
+            forward_batch, "cp_shared_kv_num_model_layers", None
-            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,
         )
+        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

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