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>
This commit is contained in:
laoyao0822
@@ -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,
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user