diff --git a/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md b/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md index 6ea3e90fa..a87ab76c0 100644 --- a/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md +++ b/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md @@ -115,8 +115,8 @@ slot-dense locs but no valid current KV row. ## Root-cause record: unsafe partial-current prefetch -The current MLA partial-current prefetch path has a narrower contract than the -page-level metadata it consumes. +The original MLA partial-current prefetch path had a narrower contract than the +page-level metadata it consumed. Evidence: @@ -126,8 +126,8 @@ Evidence: - `cp_shared_kv_runtime.py::merge_materialized_and_current_kv()` only remaps entries present in `current_locs`; all non-current entries keep the existing `materialized_dense_locs`. -- `nsa_backend.py` calls `consume_prefix_with_current()` whenever an MLA - prefetcher exists on a mixed prefix/current batch, before falling back to the +- `nsa_backend.py` called `consume_prefix_with_current()` whenever an MLA + prefetcher existed on a mixed prefix/current batch, before falling back to the older partial materialize path. That is safe only when the page-level suffix represented by the page table is @@ -137,13 +137,18 @@ exactly covered by `current_locs`. It is not safe when: ceil(extend_len / page_size) * page_size > extend_len ``` -The older non-prefetch partial path is safer because it materializes all -non-current logical locations and only substitutes rows that match -`current_locs`. +The older non-prefetch partial path looked safer because it materialized all +non-current logical locations and only substituted rows that matched +`current_locs`, but it still re-entered the compact +`merge_materialized_and_current_kv()` contract. Under the page-aligned cache +contract this is no longer an acceptable fallback for mixed prefix/current MLA +reuse. -The fix should not be a permanent fallback. The intended fix is to make the -prefetch compose path create a page-padded current suffix buffer and map suffix -slack to `-1`. +The fix is intentionally not a permanent fallback. The page-slot prefetch +compose path creates a page-padded current suffix buffer and maps suffix slack +to `-1`. If that compose path is unavailable, the runtime now fails fast with +`[CP_SHARED_KV_FAIL_FAST][mla_partial_current_prefetch]` instead of silently +using compact materialize/current merge. ## Existing evidence in the codebase @@ -445,7 +450,7 @@ Tests: - `extend_len=3 pages, cp_size=8` should produce three compute owners, not eight or sixteen. -### C4. MLA partial-current prefetch compose is partially corrected +### C4. MLA partial-current prefetch compose is fail-fast without prefetch Current state: @@ -455,16 +460,19 @@ Current state: their dense page slots and masks non-current locs in current pages to `-1`. - A TAI primitive exists for the CUDA hot path, with a visible warning fallback to the Torch reference when unavailable. +- `nsa_backend.py` no longer uses the compact + `merge_materialized_and_current_kv()` helper when MLA partial-current reuse + has prefix + current suffix but page-slot prefetch compose is unavailable. + That path now raises + `[CP_SHARED_KV_FAIL_FAST][mla_partial_current_prefetch]` with layer/rank, + prefix/extend lengths, table shape, and remap capacity metadata. -Correction still needed: +Design decision: -- Verify that every MLA mixed prefix/current path uses the page-slot compose - helper when a dense slot buffer already exists. -- Audit the non-prefetch fallback in `nsa_backend.py`: it still calls - `merge_materialized_and_current_kv()`, which is the compact-current helper. - That path may be correct only when the materialized loc table has already - masked suffix slack. This must be proven by tests or changed to the same - page-slot helper. +- Mixed prefix/current MLA reuse must use the page-slot compose path. If the + prefetcher is missing or `consume_prefix_with_current()` cannot provide a + composed buffer, failing immediately is safer than re-entering the compact + materialize/current merge path. - Keep `merge_materialized_and_current_kv()` documented as compact-materialize only; do not silently use it for page-slot prefetch compose. @@ -472,8 +480,9 @@ Tests: - Existing MLA slack test covers a tail page where only part of the suffix page is valid. -- Add a non-prefetch partial-current test with the same tail-page shape before - changing the fallback path. +- `test_mla_partial_current_path_fails_fast_instead_of_compact_fallback` + statically guards `nsa_backend.py` against reintroducing + `merge_materialized_and_current_kv()` calls and requires the fail-fast marker. - CUDA execution of the TAI primitive must be verified remotely; do not test CUDA locally. @@ -1014,15 +1023,17 @@ Completed C14: 123 passed, 5 warnings ``` -## Testing requirements +## Testing requirements / status -Add targeted tests before implementation: +Targeted coverage needed for this contract: 1. `extend_len % page_size != 0` MLA partial current reuse: - prefix page-aligned, - suffix has page slack, - dense buffer includes padded suffix rows, - slack locs remap to `-1`. + - status: page-slot compose is covered; non-prefetch compact fallback is + guarded by fail-fast source test until a direct forward-path fixture exists. 2. Equivalent NSA index partial reuse case: - tail page included in page table, diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index 9b77be7eb..40bee5d4b 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/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() diff --git a/test/registered/unit/layers/test_nsa_cp_utils.py b/test/registered/unit/layers/test_nsa_cp_utils.py index 8b303c000..a84795e96 100644 --- a/test/registered/unit/layers/test_nsa_cp_utils.py +++ b/test/registered/unit/layers/test_nsa_cp_utils.py @@ -1,3 +1,5 @@ +import ast +from pathlib import Path import unittest import sys from types import SimpleNamespace @@ -823,6 +825,32 @@ class TestNSAInSeqCPUtils(unittest.TestCase): self.assertIs(actual, expected) + def test_mla_partial_current_path_fails_fast_instead_of_compact_fallback(self): + backend_path = ( + Path(__file__).resolve().parents[4] + / "python/sglang/srt/layers/attention/nsa_backend.py" + ) + source = backend_path.read_text() + tree = ast.parse(source) + compact_merge_calls = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "merge_materialized_and_current_kv" + ] + + self.assertEqual( + compact_merge_calls, + [], + "MLA partial-current reuse must not fall back to compact " + "materialize/current merge when page-slot prefetch compose is unavailable.", + ) + self.assertIn( + "[CP_SHARED_KV_FAIL_FAST][mla_partial_current_prefetch]", + source, + ) + if __name__ == "__main__": unittest.main()