Fail fast when CP MLA partial-current compose is unavailable

Mixed prefix/current MLA reuse depends on the page-slot prefetch compose path to preserve padded tail semantics. The compact materialize/current merge path can re-expose suffix slack as valid dense rows, so the backend now raises with an explicit fail-fast marker instead of silently falling back.

Constraint: Page-aligned CP shared KV contract requires suffix tail slack to remain invalid.

Rejected: Keep compact merge as a fallback | it has a different dense-row contract and can hide correctness bugs.

Confidence: high

Scope-risk: moderate

Directive: Do not reintroduce merge_materialized_and_current_kv on MLA partial-current reuse without a page-slot correctness proof.

Tested: Remote py_compile for nsa_backend.py and test_nsa_cp_utils.py in g0034 container.

Tested: Remote pytest test_nsa_cp_utils.py::TestNSAInSeqCPUtils::test_mla_partial_current_path_fails_fast_instead_of_compact_fallback.

Tested: Remote pytest test_nsa_cp_utils.py test_cp_shared_kv_layout.py test_cp_shared_kv_runtime.py: 124 passed, 5 warnings.

Not-tested: Live ETE traffic and CUDA kernel execution for this fail-fast path.

Co-authored-by: OmX <omx@oh-my-codex.dev>

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

@@ -27,7 +27,6 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
     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,
@@ -1830,35 +1829,43 @@ class NativeSparseAttnBackend(
                     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,
+                        prefix_lens_cpu = getattr(
+                            forward_batch, "extend_prefix_lens_cpu", None
                         )
-                        materialize_locs = torch.where(
-                            current_mask,
-                            torch.full_like(logical_page_table_1, -1),
-                            logical_page_table_1,
+                        extend_lens_cpu = getattr(
+                            forward_batch, "extend_seq_lens_cpu", None
                         )
-                        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,
-                            )
+                        reason = (
+                            "missing_prefetcher"
+                            if mla_prefetcher is None
+                            else "prefetch_consume_returned_none"
                         )
-                        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,
+                        prefix_lens = (
+                            [int(x) for x in prefix_lens_cpu]
+                            if prefix_lens_cpu is not None
+                            else None
+                        )
+                        extend_lens = (
+                            [int(x) for x in extend_lens_cpu]
+                            if extend_lens_cpu is not None
+                            else None
+                        )
+                        raise RuntimeError(
+                            "[CP_SHARED_KV_FAIL_FAST][mla_partial_current_prefetch] "
+                            "CP shared KV MLA partial-current reuse requires "
+                            "page-slot prefetch compose. Compact "
+                            "materialize/current merge fallback is disabled "
+                            "because it can expose padded tail slack. "
+                            f"reason={reason} "
+                            f"cp_rank={forward_batch.cp_shared_kv_layout.cp_rank} "
+                            f"layer_id={layer.layer_id} "
+                            f"prefix_lens={prefix_lens} "
+                            f"extend_lens={extend_lens} "
+                            f"current_rows={int(current_kv_cache.shape[0])} "
+                            f"logical_page_table_shape={tuple(logical_page_table_1.shape)} "
+                            f"current_locs_shape={tuple(forward_batch.out_cache_loc.shape)} "
+                            f"page_size={current_remap_page_size} "
+                            f"logical_page_capacity={current_remap_logical_page_capacity}"
                         )
                     if (
                         cp_shared_kv_mla_prefetch_log_enabled()