Make CP shared-KV direct writes enforce batch-local ownership
W3 needs batch-size>1 extends to use packed valid tokens while preserving per-request page boundaries. The local out_cache_loc planner now validates against the batch plan's request lengths instead of the first request's scalar split metadata, then reuses the existing batch split helper to produce this rank's logical/physical cache locs. Direct-write failures inside the CP shared-KV contract now fail fast instead of silently falling back to legacy index/MLA stores. This exposes allocator, owner-lane, page-alignment, and shape-contract bugs early for both bs=1 and bs>1. Constraint: bs>1 batching must not pad short requests to the longest request; only per-request page-boundary padding is allowed. Rejected: Keep bs=1 compatibility fallback | it hides CP shared-KV contract violations and caused repeated slow-path ambiguity. Rejected: Pad batch to max request length | wastes compute and complicates cache validity for short extends. Confidence: high Scope-risk: moderate Directive: CP shared-KV contract errors should stay fail-fast; do not reintroduce silent direct-write fallback without ETE evidence and explicit warning semantics. Tested: Remote g0034 py_compile for utils.py nsa_indexer.py forward_mla.py Tested: Remote g0034 PYTHONPATH=python pytest test/registered/unit/layers/test_nsa_cp_utils.py -> 43 passed Tested: Remote g0034 PYTHONPATH=python pytest test_nsa_cp_utils.py test_cp_shared_kv_layout.py test_cp_shared_kv_runtime.py -> 170 passed, 2 subtests passed Not-tested: Full ETE bs>1 serving run with live traffic
This commit is contained in:
laoyao0822
@@ -709,6 +709,95 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
|
||||
+ list(range(6 * page_size, 7 * page_size)),
|
||||
)
|
||||
|
||||
def test_batch_local_out_cache_loc_keeps_request_boundaries(self):
|
||||
import torch
|
||||
from types import SimpleNamespace
|
||||
|
||||
page_size = 4
|
||||
# cp_size=2/cp_rank=1 selects segment 1 and 2 for each request.
|
||||
# req0 has no rank-1 local rows. req1 contributes segment 1 (page 2)
|
||||
# and segment 2 (tail page 4). The synthetic logical page ids encode
|
||||
# the owner-lane invariant through (page_id - 1) % cp_size.
|
||||
out_cache_loc = torch.cat(
|
||||
[
|
||||
torch.arange(1 * page_size, 2 * page_size), # req0 seg0 owner 0
|
||||
torch.arange(3 * page_size, 4 * page_size), # req1 seg0 owner 0
|
||||
torch.arange(2 * page_size, 3 * page_size), # req1 seg1 owner 1
|
||||
torch.tensor([4 * page_size]), # req1 seg2 owner 1 tail
|
||||
]
|
||||
)
|
||||
forward_batch = SimpleNamespace(
|
||||
uses_cp_shared_kv=True,
|
||||
cp_shared_kv_layout=CpSharedKVLayout(
|
||||
page_size=page_size,
|
||||
cp_size=2,
|
||||
cp_rank=1,
|
||||
),
|
||||
nsa_cp_metadata=NSAContextParallelMetadata(
|
||||
batch_size=2,
|
||||
split_list=[4, 0, 0, 0],
|
||||
zigzag_index=[1, 2],
|
||||
page_aligned=True,
|
||||
page_size=page_size,
|
||||
extend_prefix_len=0,
|
||||
request_extend_lens=[4, 9],
|
||||
request_split_lists=[[4, 0, 0, 0], [4, 4, 1, 0]],
|
||||
request_zigzag_indices=[[1, 2], [1, 2]],
|
||||
),
|
||||
out_cache_loc=out_cache_loc,
|
||||
)
|
||||
|
||||
local_locs = get_cp_shared_kv_local_out_cache_loc(forward_batch)
|
||||
|
||||
self.assertIsNotNone(local_locs)
|
||||
self.assertEqual(
|
||||
local_locs.tolist(),
|
||||
list(range(2 * page_size, 3 * page_size)) + [4 * page_size],
|
||||
)
|
||||
|
||||
def test_batch_local_physical_out_cache_loc_reuses_layer_invariant_plan(self):
|
||||
import torch
|
||||
from types import SimpleNamespace
|
||||
|
||||
page_size = 4
|
||||
out_cache_loc = torch.cat(
|
||||
[
|
||||
torch.arange(1 * page_size, 2 * page_size),
|
||||
torch.arange(3 * page_size, 4 * page_size),
|
||||
torch.arange(2 * page_size, 3 * page_size),
|
||||
torch.tensor([4 * page_size]),
|
||||
]
|
||||
)
|
||||
forward_batch = SimpleNamespace(
|
||||
uses_cp_shared_kv=True,
|
||||
cp_shared_kv_layout=CpSharedKVLayout(
|
||||
page_size=page_size,
|
||||
cp_size=2,
|
||||
cp_rank=1,
|
||||
),
|
||||
nsa_cp_metadata=NSAContextParallelMetadata(
|
||||
batch_size=2,
|
||||
split_list=[4, 0, 0, 0],
|
||||
zigzag_index=[1, 2],
|
||||
page_aligned=True,
|
||||
page_size=page_size,
|
||||
extend_prefix_len=0,
|
||||
request_extend_lens=[4, 9],
|
||||
request_split_lists=[[4, 0, 0, 0], [4, 4, 1, 0]],
|
||||
request_zigzag_indices=[[1, 2], [1, 2]],
|
||||
),
|
||||
out_cache_loc=out_cache_loc,
|
||||
)
|
||||
|
||||
physical_locs = get_cp_shared_kv_local_physical_out_cache_loc(forward_batch)
|
||||
second_read = get_cp_shared_kv_local_physical_out_cache_loc(forward_batch)
|
||||
|
||||
self.assertIs(physical_locs, second_read)
|
||||
self.assertEqual(
|
||||
physical_locs.tolist(),
|
||||
list(range(1 * page_size, 2 * page_size)) + [2 * page_size],
|
||||
)
|
||||
|
||||
def test_local_physical_out_cache_loc_is_cached(self):
|
||||
import torch
|
||||
from types import SimpleNamespace
|
||||
@@ -748,7 +837,7 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
|
||||
+ list(range(2 * page_size, 3 * page_size)),
|
||||
)
|
||||
|
||||
def test_local_out_cache_loc_falls_back_when_owner_mismatch(self):
|
||||
def test_local_out_cache_loc_fails_fast_when_owner_mismatch(self):
|
||||
import torch
|
||||
from types import SimpleNamespace
|
||||
|
||||
@@ -771,14 +860,16 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
|
||||
out_cache_loc=out_cache_loc,
|
||||
)
|
||||
|
||||
self.assertIsNone(get_cp_shared_kv_local_out_cache_loc(forward_batch))
|
||||
with self.assertRaisesRegex(
|
||||
RuntimeError,
|
||||
"CP_SHARED_KV_FAIL_FAST.*local_loc_owner_mismatch",
|
||||
):
|
||||
get_cp_shared_kv_local_out_cache_loc(forward_batch)
|
||||
|
||||
def test_local_out_cache_loc_logs_every_fallback_event(self):
|
||||
def test_local_out_cache_loc_fails_fast_every_invalid_event(self):
|
||||
import torch
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.srt.layers.attention.nsa import utils as nsa_utils
|
||||
|
||||
page_size = 4
|
||||
forward_batch = SimpleNamespace(
|
||||
uses_cp_shared_kv=True,
|
||||
@@ -798,17 +889,81 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
|
||||
)
|
||||
|
||||
with self.assertLogs(
|
||||
"sglang.srt.layers.attention.nsa.utils", level="WARNING"
|
||||
"sglang.srt.layers.attention.nsa.utils", level="ERROR"
|
||||
) as cm:
|
||||
self.assertIsNone(get_cp_shared_kv_local_out_cache_loc(forward_batch))
|
||||
self.assertIsNone(get_cp_shared_kv_local_out_cache_loc(forward_batch))
|
||||
with self.assertRaisesRegex(
|
||||
RuntimeError,
|
||||
"CP_SHARED_KV_FAIL_FAST.*not_page_aligned",
|
||||
):
|
||||
get_cp_shared_kv_local_out_cache_loc(forward_batch)
|
||||
with self.assertRaisesRegex(
|
||||
RuntimeError,
|
||||
"CP_SHARED_KV_FAIL_FAST.*not_page_aligned",
|
||||
):
|
||||
get_cp_shared_kv_local_out_cache_loc(forward_batch)
|
||||
|
||||
self.assertEqual(len(cm.output), 2)
|
||||
self.assertIn("[CP_SHARED_KV_FALLBACK][direct_write]", cm.output[0])
|
||||
self.assertIn("[CP_SHARED_KV_FAIL_FAST][direct_write]", cm.output[0])
|
||||
self.assertIn("metadata is not page-aligned", cm.output[0])
|
||||
self.assertIn("[CP_SHARED_KV_FALLBACK][direct_write]", cm.output[1])
|
||||
self.assertIn("[CP_SHARED_KV_FAIL_FAST][direct_write]", cm.output[1])
|
||||
self.assertIn("metadata is not page-aligned", cm.output[1])
|
||||
|
||||
def test_indexer_direct_write_fails_fast_on_local_shape_mismatch(self):
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.attention.nsa import nsa_indexer
|
||||
from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
|
||||
|
||||
indexer = object.__new__(Indexer)
|
||||
indexer.nsa_enable_prefill_cp = True
|
||||
forward_batch = SimpleNamespace(nsa_cp_metadata=object())
|
||||
|
||||
with (
|
||||
patch.object(nsa_indexer, "nsa_use_prefill_cp", return_value=True),
|
||||
patch.object(
|
||||
nsa_indexer,
|
||||
"get_cp_shared_kv_local_out_cache_loc",
|
||||
return_value=torch.tensor([1, 2], dtype=torch.int64),
|
||||
),
|
||||
):
|
||||
with self.assertRaisesRegex(
|
||||
RuntimeError,
|
||||
"CP_SHARED_KV_FAIL_FAST.*index_local_shape_mismatch",
|
||||
):
|
||||
Indexer._store_cp_shared_local_index_k_cache(
|
||||
indexer,
|
||||
forward_batch,
|
||||
layer_id=0,
|
||||
local_key=torch.empty((1, 8)),
|
||||
act_quant=None,
|
||||
)
|
||||
|
||||
def test_mla_direct_write_fails_fast_on_local_shape_mismatch(self):
|
||||
import torch
|
||||
|
||||
from sglang.srt.models.deepseek_common.attention_forward_methods import (
|
||||
forward_mla,
|
||||
)
|
||||
|
||||
mla = SimpleNamespace(attn_mqa=SimpleNamespace(layer_id=3))
|
||||
forward_batch = SimpleNamespace()
|
||||
|
||||
with patch.object(
|
||||
forward_mla,
|
||||
"get_cp_shared_kv_local_out_cache_loc",
|
||||
return_value=torch.tensor([1, 2], dtype=torch.int64),
|
||||
):
|
||||
with self.assertRaisesRegex(
|
||||
RuntimeError,
|
||||
"CP_SHARED_KV_FAIL_FAST.*mla_local_shape_mismatch",
|
||||
):
|
||||
forward_mla.DeepseekMLAForwardMixin._maybe_write_cp_shared_local_mla_kv(
|
||||
mla,
|
||||
forward_batch,
|
||||
k_nope=torch.empty((1, 8)),
|
||||
k_pe=torch.empty((2, 8)),
|
||||
)
|
||||
|
||||
def test_indexer_direct_write_does_not_log_missing_metadata_for_non_cp_batch(self):
|
||||
import torch
|
||||
from types import SimpleNamespace
|
||||
|
||||
Reference in New Issue
Block a user