Protect CP shared-KV cache-hit correctness under batched FP8 reuse

Cache-hit GSM8K regressions only appeared after the second pass reused request-specific suffix pages, so this change adds fail-fast transfer validation, masks stale rectangular page-table tails, and extends CUDA/unit coverage across FP8 CP shared-KV write, load, top-k, and materialization paths. The temporary ledger records eliminated hypotheses to prevent re-debugging the same L2 and persistent-cache paths.\n\nConstraint: CP shared KV stores physical pages but scheduler-visible semantics must remain valid-token/page-bounded.\nConstraint: bs>1 FP8 prefill must preserve existing CP shared-KV fast paths without silent fallback.\nRejected: Blame raw HiCache L2 load without tests | L2 KV and index backup/load/materialize roundtrips pass on remote CUDA.\nRejected: Disable current/partial reuse broadly | hides the cache-hit contract regression and costs performance.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not weaken CP shared-KV fail-fast or rectangular-tail masking without rerunning second-pass cache-hit accuracy tests.\nTested: remote CUDA pytest for fused FP8 MLA store, fused persistent index store, L2-loaded FP8 KV materialize, L2-loaded index materialize, ragged top-k offset, TAI batched index MQA prepare.\nTested: local py_compile for touched test files and git diff --check.\nNot-tested: full second-pass GSM8K accuracy after these diagnostic tests; root cause remains under investigation.

python/sglang/srt/disaggregation/prefill.py

@@ -524,6 +524,8 @@ class PrefillBootstrapQueue:
             [req.disagg_kv_sender for req in self.queue],
             self.scheduler.attn_cp_cpu_group,
             self.scheduler.attn_tp_cpu_group,
+            debug_label="bootstrap",
+            debug_ids=[req.rid for req in self.queue],
         )
         _cp_shared_kv_bs_gt1_prefill_timing(
             "bootstrap_poll_done",
@@ -937,6 +939,8 @@ class SchedulerDisaggregationPrefillMixin:
             [req.disagg_kv_sender for req in self.disagg_prefill_inflight_queue],
             self.attn_cp_cpu_group,
             self.attn_tp_cpu_group,
+            debug_label="inflight",
+            debug_ids=[req.rid for req in self.disagg_prefill_inflight_queue],
         )
         _cp_shared_kv_bs_gt1_prefill_timing(
             "inflight_poll_done",
@@ -1050,6 +1054,8 @@ class SchedulerDisaggregationPrefillMixin:
             [req.disagg_kv_sender for req in self.disagg_prefill_inflight_queue],
             self.attn_cp_cpu_group,
             self.attn_tp_cpu_group,
+            debug_label="get_transferred_rids",
+            debug_ids=[req.rid for req in self.disagg_prefill_inflight_queue],
         )
 
         transferred_rids: List[str] = []

python/sglang/srt/disaggregation/utils.py

@@ -63,11 +63,85 @@ def poll_and_all_reduce(pollers, gloo_group: dist.ProcessGroup):
     return tensor_to_reduce.tolist()
 
 
+def _cp_shared_kv_poll_debug_enabled() -> bool:
+    return envs.SGLANG_CP_SHARED_KV_BS_GT1_DEBUG.get()
+
+
+def _poll_queue_debug_hash(debug_ids: Optional[list[str]]) -> int:
+    if not debug_ids:
+        return 0
+
+    # Stable bounded FNV-1a style digest.  Keep it in signed-int64 range because
+    # the debug consensus uses torch.int64 CPU collectives.
+    value = 1469598103934665603
+    mask = (1 << 63) - 1
+    for item in debug_ids:
+        for byte in str(item).encode("utf-8", errors="replace"):
+            value ^= byte
+            value = (value * 1099511628211) & mask
+        value ^= 0xFF
+        value = (value * 1099511628211) & mask
+    return int(value)
+
+
+def _validate_poll_queue_consensus(
+    *,
+    label: str,
+    scope: str,
+    local_len: int,
+    debug_ids: Optional[list[str]],
+    group: dist.ProcessGroup,
+) -> None:
+    local_hash = _poll_queue_debug_hash(debug_ids)
+    signature = torch.tensor([local_len, local_hash], dtype=torch.int64, device="cpu")
+    min_signature = signature.clone()
+    max_signature = signature.clone()
+
+    dist.all_reduce(min_signature, op=dist.ReduceOp.MIN, group=group)
+    dist.all_reduce(max_signature, op=dist.ReduceOp.MAX, group=group)
+
+    if torch.equal(min_signature, max_signature):
+        return
+
+    ids_head = list(debug_ids[:8]) if debug_ids is not None else None
+    message = (
+        "[CP_SHARED_KV_FAIL_FAST][poll_queue] "
+        f"label={label} scope={scope} local_len={local_len} "
+        f"min_len={int(min_signature[0].item())} "
+        f"max_len={int(max_signature[0].item())} "
+        f"local_hash={local_hash} min_hash={int(min_signature[1].item())} "
+        f"max_hash={int(max_signature[1].item())} ids_head={ids_head}"
+    )
+    logger.error(message)
+    raise RuntimeError(message)
+
+
 def poll_and_all_reduce_attn_cp_tp_group(
     pollers,
     attn_cp_cpu_group: dist.ProcessGroup,
     attn_tp_cpu_group: dist.ProcessGroup,
+    *,
+    debug_label: Optional[str] = None,
+    debug_ids: Optional[list[str]] = None,
 ):
+    if _cp_shared_kv_poll_debug_enabled():
+        label = debug_label or "unknown"
+        local_len = len(pollers)
+        _validate_poll_queue_consensus(
+            label=label,
+            scope="attn_tp",
+            local_len=local_len,
+            debug_ids=debug_ids,
+            group=attn_tp_cpu_group,
+        )
+        _validate_poll_queue_consensus(
+            label=label,
+            scope="attn_cp",
+            local_len=local_len,
+            debug_ids=debug_ids,
+            group=attn_cp_cpu_group,
+        )
+
     # First sync across attn-tp ranks so all TP participants for a given (dp, cp)
     # shard observe the same status transitions.
     polls = poll_and_all_reduce(pollers, attn_tp_cpu_group)

python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py

@@ -228,6 +228,128 @@ def _slot_remap_cache_key(
     )
 
 
+def _to_cpu_int_list(values: Any, *, name: str) -> list[int]:
+    if values is None:
+        raise ValueError(f"{name} must not be None")
+    if isinstance(values, torch.Tensor):
+        if values.dim() == 0:
+            result = [int(values.item())]
+        else:
+            result = [int(x) for x in values.detach().cpu().reshape(-1).tolist()]
+    else:
+        result = [int(x) for x in values]
+    for idx, value in enumerate(result):
+        if value < 0:
+            raise ValueError(f"{name} contains negative value: idx={idx} value={value}")
+    return result
+
+
+def _normalize_seq_lens_for_logical_page_rows(
+    *,
+    seq_lens_cpu: Any,
+    rows: int,
+    repeat_lens_cpu: Any | None = None,
+) -> list[int]:
+    seq_lens = _to_cpu_int_list(seq_lens_cpu, name="seq_lens_cpu")
+    if rows < 0:
+        raise ValueError(f"rows must be non-negative, got {rows}")
+    if rows == 0:
+        if seq_lens:
+            raise ValueError(
+                "logical page table has no rows but seq_lens_cpu is non-empty: "
+                f"seq_lens={seq_lens}"
+            )
+        return []
+    if len(seq_lens) == rows:
+        return seq_lens
+
+    repeat_lens = (
+        _to_cpu_int_list(repeat_lens_cpu, name="repeat_lens_cpu")
+        if repeat_lens_cpu is not None
+        else None
+    )
+    if (
+        repeat_lens is not None
+        and len(repeat_lens) == len(seq_lens)
+        and sum(repeat_lens) == rows
+    ):
+        expanded: list[int] = []
+        for seq_len, repeat in zip(seq_lens, repeat_lens, strict=True):
+            expanded.extend([seq_len] * int(repeat))
+        return expanded
+
+    if len(seq_lens) > 0 and rows % len(seq_lens) == 0:
+        repeat = rows // len(seq_lens)
+        expanded = []
+        for seq_len in seq_lens:
+            expanded.extend([seq_len] * repeat)
+        return expanded
+
+    raise ValueError(
+        "Cannot align seq_lens_cpu with logical page rows: "
+        f"rows={rows} seq_lens={seq_lens} repeat_lens={repeat_lens}"
+    )
+
+
+def mask_batch_logical_pages_to_valid_lengths(
+    logical_pages: torch.Tensor,
+    *,
+    seq_lens_cpu: Any,
+    page_size: int,
+    repeat_lens_cpu: Any | None = None,
+) -> torch.Tensor:
+    """Zero rectangular page-table tails beyond each request's valid pages.
+
+    CP shared-KV slot remap gives every flattened page-table slot a deterministic
+    dense page id.  Therefore stale values in the rectangular tail are not just
+    unused padding: duplicate stale pages can overwrite the logical-page inverse
+    and redirect valid top-k/cache-hit locations to the wrong dense slot.  Mask
+    those tails at the metadata boundary while keeping page-level padding inside
+    the final valid page visible.
+    """
+
+    if page_size <= 0:
+        raise ValueError(f"page_size must be positive, got {page_size}")
+    if logical_pages.dim() == 0:
+        raise ValueError("logical_pages must have at least one dimension")
+
+    if logical_pages.dim() == 1:
+        rows = 1
+        pages_per_request = int(logical_pages.numel())
+    else:
+        rows = int(logical_pages.shape[0])
+        pages_per_request = int(logical_pages.reshape(rows, -1).shape[1])
+
+    seq_lens = _normalize_seq_lens_for_logical_page_rows(
+        seq_lens_cpu=seq_lens_cpu,
+        rows=rows,
+        repeat_lens_cpu=repeat_lens_cpu,
+    )
+    if rows == 0 or pages_per_request == 0:
+        return logical_pages.clone()
+
+    valid_pages_cpu = [
+        (int(seq_len) + int(page_size) - 1) // int(page_size)
+        for seq_len in seq_lens
+    ]
+    valid_pages = torch.tensor(
+        valid_pages_cpu,
+        device=logical_pages.device,
+        dtype=torch.long,
+    )
+    columns = torch.arange(
+        pages_per_request,
+        device=logical_pages.device,
+        dtype=torch.long,
+    )
+    keep_mask = columns.unsqueeze(0) < valid_pages.unsqueeze(1)
+
+    masked = logical_pages.clone()
+    masked_2d = masked.reshape(rows, pages_per_request)
+    masked_2d.masked_fill_(~keep_mask, 0)
+    return masked
+
+
 def _log_slot_remap_cache_not_reused(
     *,
     kind: str,

python/sglang/srt/layers/attention/nsa_backend.py

@@ -31,6 +31,7 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
     get_or_build_shared_token_kv_slot_remap,
     is_current_only_extend_batch,
     is_packed_fp8_mla_kv_cache,
+    mask_batch_logical_pages_to_valid_lengths,
     materialize_prefix_and_reuse_current_kv_page_slots,
     materialize_shared_token_kv_buffer,
     pack_current_mla_kv_for_reuse,
@@ -932,6 +933,21 @@ class NativeSparseAttnBackend(
             except (ImportError, ModuleNotFoundError):
                 paged_mqa_schedule_metadata = None
 
+        real_page_table = self._transform_table_1_to_real(page_table)
+        real_page_valid_seq_lens_cpu = indexer_seq_lens_cpu
+        if draft_token_num:
+            real_page_valid_seq_lens_cpu = (
+                indexer_seq_lens_cpu + draft_token_num
+                if isinstance(indexer_seq_lens_cpu, torch.Tensor)
+                else [int(x) + int(draft_token_num) for x in indexer_seq_lens_cpu]
+            )
+        real_page_table = mask_batch_logical_pages_to_valid_lengths(
+            real_page_table,
+            seq_lens_cpu=real_page_valid_seq_lens_cpu,
+            page_size=self.real_page_size,
+            repeat_lens_cpu=extend_seq_lens_cpu,
+        )
+
         metadata = NSAMetadata(
             page_size=self.real_page_size,
             cache_seqlens_int32=cache_seqlens_int32,
@@ -957,7 +973,7 @@ class NativeSparseAttnBackend(
             nsa_cu_seqlens_k=nsa_cu_seqlens_k,
             nsa_seqlens_expanded=seqlens_expanded,
             nsa_extend_seq_lens_list=extend_seq_lens_cpu,
-            real_page_table=self._transform_table_1_to_real(page_table),
+            real_page_table=real_page_table,
             nsa_max_seqlen_q=1,
             topk_indices_offset=topk_indices_offset,
             indexer_k_start_end=indexer_k_start_end,