Fail fast when CP compose would hide dense fallback

CP shared-KV bs>1 compose must not silently fall back to dense full-buffer collectives when CUDA TAI materialize is expected. The fallback masks both correctness contract drift and severe synchronization/communication regressions, especially while comparing the symm path with the older IPC path.\n\nThis keeps CPU/unit-test fallback available, but makes production CUDA+TAI runs raise an explicit compose_v2 fail-fast for token-KV and index dense fallback. It also records the symm-vs-IPC comparison contract so barrier and collective counts are evaluated alongside elapsed time.\n\nConstraint: Production cache-hit-heavy bs>1 paths must expose unexpected dense collectives instead of silently taking them.\nRejected: Cherry-pick the old IPC branch wholesale | it conflicts with the symm compose design and would mix two transport protocols before benchmarking.\nRejected: Allow dense fallback with warning only | warning can be missed and still corrupts performance conclusions.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not re-enable dense full fallback in CUDA+TAI compose paths without a benchmark proving it is intentional and a correctness test covering cache-hit bs>1.\nTested: python -m py_compile for cp_shared_kv_runtime.py and test_cp_shared_kv_runtime.py; git diff --check.\nNot-tested: Remote container pytest/ETE; local pytest is not reliable in this workspace because dependencies such as orjson are missing.

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

@@ -832,6 +832,27 @@ def _log_tai_index_mqa_prepare_fallback(
     )
 
 
+def _should_fail_fast_compose_v2_dense_fallback(dense_tensor: torch.Tensor) -> bool:
+    """Whether compose_v2 may fall back to dense full-buffer collectives.
+
+    CPU/unit-test paths can still use the simple fallback.  In production CUDA
+    runs with TAI materialize enabled, a prefix-IPC miss would silently turn a
+    bs>1 cache-hit compose into a dense all_reduce over the whole buffer.  That
+    hides both correctness-contract drift and severe performance regressions,
+    so fail fast instead.
+    """
+
+    return bool(dense_tensor.is_cuda and cp_shared_kv_tai_materialize_enabled())
+
+
+def _raise_compose_v2_dense_fallback_required(reason: str, **details: Any) -> None:
+    detail_str = " ".join(f"{key}={value}" for key, value in details.items())
+    message = f"[CP_SHARED_KV_FAIL_FAST][compose_v2] reason={reason}"
+    if detail_str:
+        message = f"{message} {detail_str}"
+    raise RuntimeError(message)
+
+
 def _tai_materialize_runtime_disabled_reason() -> str | None:
     if not cp_shared_kv_tai_materialize_enabled():
         return "env_disabled"
@@ -4838,6 +4859,19 @@ def _compose_token_kv_partial_current_v2(
             )
     else:
         dense_kv_cache = kv_cache.new_zeros((dense_rows, *kv_cache.shape[1:]))
+        if prefix_spans and _should_fail_fast_compose_v2_dense_fallback(
+            dense_kv_cache
+        ):
+            _raise_compose_v2_dense_fallback_required(
+                "token_kv_dense_fallback",
+                cp_rank=layout.cp_rank,
+                cp_size=layout.cp_size,
+                layer_id=layer_id,
+                prefix_spans=prefix_spans,
+                current_spans=current_spans,
+                dense_shape=tuple(dense_kv_cache.shape),
+                kv_dtype=kv_cache.dtype,
+            )
         for prefix_start_slot, prefix_end_slot in prefix_spans:
             materialize_local_token_kv_page_slots_into(
                 kv_cache=kv_cache,
@@ -5294,6 +5328,19 @@ def _compose_index_partial_current_v2(
         dense_page_buffer = page_buffer.new_zeros(
             (dense_num_pages, *page_buffer.shape[1:])
         )
+        if prefix_spans and _should_fail_fast_compose_v2_dense_fallback(
+            dense_page_buffer
+        ):
+            _raise_compose_v2_dense_fallback_required(
+                "index_dense_fallback",
+                cp_rank=layout.cp_rank,
+                cp_size=layout.cp_size,
+                layer_id=layer_id,
+                prefix_spans=prefix_spans,
+                current_spans=current_spans,
+                dense_shape=tuple(dense_page_buffer.shape),
+                index_dtype=page_buffer.dtype,
+            )
         for prefix_start_slot, prefix_end_slot in prefix_spans:
             materialize_local_paged_buffer_page_slots_into(
                 page_buffer=page_buffer,