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
2026-05-29 06:46:39 +08:00
co-authored by OmX
parent 4f9bb7ce30
commit c9f790cde9
3 changed files with 96 additions and 50 deletions
@@ -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()