Make CP shared-KV direct writes enforce batch-local ownership

W3 needs batch-size>1 extends to use packed valid tokens while preserving per-request page boundaries. The local out_cache_loc planner now validates against the batch plan's request lengths instead of the first request's scalar split metadata, then reuses the existing batch split helper to produce this rank's logical/physical cache locs.

Direct-write failures inside the CP shared-KV contract now fail fast instead of silently falling back to legacy index/MLA stores. This exposes allocator, owner-lane, page-alignment, and shape-contract bugs early for both bs=1 and bs>1.

Constraint: bs>1 batching must not pad short requests to the longest request; only per-request page-boundary padding is allowed.

Rejected: Keep bs=1 compatibility fallback | it hides CP shared-KV contract violations and caused repeated slow-path ambiguity.

Rejected: Pad batch to max request length | wastes compute and complicates cache validity for short extends.

Confidence: high

Scope-risk: moderate

Directive: CP shared-KV contract errors should stay fail-fast; do not reintroduce silent direct-write fallback without ETE evidence and explicit warning semantics.

Tested: Remote g0034 py_compile for utils.py nsa_indexer.py forward_mla.py

Tested: Remote g0034 PYTHONPATH=python pytest test/registered/unit/layers/test_nsa_cp_utils.py -> 43 passed

Tested: Remote g0034 PYTHONPATH=python pytest test_nsa_cp_utils.py test_cp_shared_kv_layout.py test_cp_shared_kv_runtime.py -> 170 passed, 2 subtests passed

Not-tested: Full ETE bs>1 serving run with live traffic

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

@@ -67,8 +67,8 @@ from sglang.srt.layers.attention.nsa.utils import (
     get_cp_shared_kv_local_physical_out_cache_loc,
     is_nsa_enable_prefill_cp,
     is_nsa_prefill_cp_in_seq_split,
-    log_cp_shared_kv_direct_write_fallback,
     nsa_use_prefill_cp,
+    raise_cp_shared_kv_direct_write_error,
     split_in_seq_cp_local_pair,
 )
 from sglang.srt.layers.communicator import ScatterMode
@@ -1618,7 +1618,7 @@ class Indexer(MultiPlatformOp):
         if local_out_loc is None:
             return False
         if local_key.shape[0] != local_out_loc.numel():
-            log_cp_shared_kv_direct_write_fallback(
+            raise_cp_shared_kv_direct_write_error(
                 "index_local_shape_mismatch",
                 "NSA index local key token count does not match local out_cache_loc: "
                 "local_key=%s local_out_cache_loc=%s layer_id=%s",
@@ -1626,7 +1626,6 @@ class Indexer(MultiPlatformOp):
                 local_out_loc.numel(),
                 layer_id,
             )
-            return False
         if local_out_loc.numel() == 0:
             return True
 

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

@@ -65,6 +65,21 @@ def log_cp_shared_kv_direct_write_fallback(
     )
 
 
+def raise_cp_shared_kv_direct_write_error(
+    reason: str,
+    message: str,
+    *args,
+) -> None:
+    """Fail fast when CP shared-KV direct-write invariants are violated."""
+
+    formatted = message % args if args else message
+    error_msg = (
+        f"[CP_SHARED_KV_FAIL_FAST][direct_write] reason={reason} {formatted}"
+    )
+    logger.error(error_msg)
+    raise RuntimeError(error_msg)
+
+
 def compute_nsa_seqlens(original_seq_lens, nsa_index_topk: int):
     return original_seq_lens.clamp(max=nsa_index_topk)
 
@@ -1055,51 +1070,51 @@ def get_cp_shared_kv_local_out_cache_loc(forward_batch: "ForwardBatch"):
     layout = getattr(forward_batch, "cp_shared_kv_layout", None)
     out_cache_loc = getattr(forward_batch, "out_cache_loc", None)
     if metadata is None:
-        log_cp_shared_kv_direct_write_fallback(
+        raise_cp_shared_kv_direct_write_error(
             "missing_metadata",
             "nsa_cp_metadata is missing",
         )
-        return None
     if layout is None:
-        log_cp_shared_kv_direct_write_fallback(
+        raise_cp_shared_kv_direct_write_error(
             "missing_layout",
             "cp_shared_kv_layout is missing",
         )
-        return None
     if out_cache_loc is None:
-        log_cp_shared_kv_direct_write_fallback(
+        raise_cp_shared_kv_direct_write_error(
             "missing_out_cache_loc",
             "out_cache_loc is missing",
         )
-        return None
     if not getattr(metadata, "page_aligned", False):
-        log_cp_shared_kv_direct_write_fallback(
+        raise_cp_shared_kv_direct_write_error(
             "not_page_aligned",
             "metadata is not page-aligned: page_size=%s extend_prefix_len=%s",
             getattr(metadata, "page_size", None),
             getattr(metadata, "extend_prefix_len", None),
         )
-        return None
     try:
         in_seq_split = is_nsa_prefill_cp_in_seq_split()
     except ValueError:
         in_seq_split = True
     if not in_seq_split:
-        log_cp_shared_kv_direct_write_fallback(
+        raise_cp_shared_kv_direct_write_error(
             "not_in_seq_split",
             "nsa_prefill_cp_mode is not in-seq-split",
         )
-        return None
-    split_tokens = sum(int(x) for x in metadata.split_list)
+    batch_plan = get_cp_shared_kv_batch_plan(forward_batch)
+    if batch_plan is not None and int(getattr(batch_plan, "batch_size", 1) or 1) > 1:
+        split_tokens = sum(int(x) for x in getattr(batch_plan, "request_extend_lens", []))
+        mismatch_reason = "batch_split_out_cache_len_mismatch"
+    else:
+        split_tokens = sum(int(x) for x in metadata.split_list)
+        mismatch_reason = "split_out_cache_len_mismatch"
     out_cache_tokens = int(out_cache_loc.numel())
     if split_tokens != out_cache_tokens:
-        log_cp_shared_kv_direct_write_fallback(
-            "split_out_cache_len_mismatch",
+        raise_cp_shared_kv_direct_write_error(
+            mismatch_reason,
             "split_list tokens=%s out_cache_loc tokens=%s",
             split_tokens,
             out_cache_tokens,
         )
-        return None
 
     local_out_cache_loc = cp_split_and_rebuild_1d(
         forward_batch,
@@ -1111,14 +1126,13 @@ def get_cp_shared_kv_local_out_cache_loc(forward_batch: "ForwardBatch"):
 
     valid_locs = local_out_cache_loc[local_out_cache_loc > 0]
     if valid_locs.numel() > 0 and not torch.all(layout.owned_by_this_rank(valid_locs)):
-        log_cp_shared_kv_direct_write_fallback(
+        raise_cp_shared_kv_direct_write_error(
             "local_loc_owner_mismatch",
             "local out_cache_loc contains pages not owned by this rank: cp_rank=%s cp_size=%s page_size=%s",
             layout.cp_rank,
             layout.cp_size,
             layout.page_size,
         )
-        return None
 
     forward_batch.cp_local_out_cache_loc = local_out_cache_loc
     return local_out_cache_loc

python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py

@@ -10,8 +10,8 @@ from sglang.srt.layers.attention.nsa.utils import (
     get_cp_shared_kv_local_out_cache_loc,
     get_cp_shared_kv_local_physical_out_cache_loc,
     log_cp_draft_shared_kv_debug,
-    log_cp_shared_kv_direct_write_fallback,
     nsa_use_prefill_cp,
+    raise_cp_shared_kv_direct_write_error,
 )
 from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
     should_reuse_current_extend_kv,
@@ -561,7 +561,7 @@ class DeepseekMLAForwardMixin:
             k_nope.shape[0] != local_out_cache_loc.numel()
             or k_pe.shape[0] != local_out_cache_loc.numel()
         ):
-            log_cp_shared_kv_direct_write_fallback(
+            raise_cp_shared_kv_direct_write_error(
                 "mla_local_shape_mismatch",
                 "MLA local KV token count does not match local out_cache_loc: "
                 "k_nope=%s k_pe=%s local_out_cache_loc=%s layer_id=%s",
@@ -570,7 +570,6 @@ class DeepseekMLAForwardMixin:
                 local_out_cache_loc.numel(),
                 self.attn_mqa.layer_id,
             )
-            return False
         if local_out_cache_loc.numel() == 0:
             return True