Align CP shared-KV prefetch with the attention overlap window
Launching CP shared-KV prefetch from MLA prepare or the indexer can make next-layer prefix work overlap current-layer MQA/materialization instead of the attention window. Centralize the launch in the NSA backend after current-layer materialization and before attention, and leave the early indexer hook inert so the call site cannot regress silently. The accompanying notes capture the draft-as-forward-layer follow-up plan and the latest OOM diagnosis: the observed 178945-token failure matches the CP in-seq MQA logits allocation, so the follow-up fix is q-dimension chunking in _get_topk_ragged_with_cp(), not a max-prefetch-size gate. Constraint: CP shared-KV prefetch must avoid overlapping current-layer MQA/materialization and must not add silent fallback behavior. Rejected: Limit maximum prefetch size | hides the CP logits peak and can reduce cache/prefetch effectiveness. Rejected: Keep prepare/indexer launch sites | they place next-layer collectives in the wrong overlap window. Confidence: medium Scope-risk: moderate Directive: Do not reintroduce early CP prefetch launch without checking Nsight overlap and CP MQA memory peaks. Tested: Local git diff --check and py_compile for touched Python files. Tested: Remote container py_compile plus targeted pytest: 3 passed, 5 warnings. Not-tested: Full ETE under production traffic after this commit. Not-tested: CP in-seq MQA logits chunking; documented as follow-up.
This commit is contained in:
@@ -3101,7 +3101,7 @@ class TestCpSharedKVTaiMaterializeIntegration(unittest.TestCase):
|
||||
self.assertIs(fake_prefetcher.calls[0][1], fake_pool.index_buffer)
|
||||
self.assertIs(fake_prefetcher.calls[0][2], logical_pages)
|
||||
|
||||
def test_index_prefetch_start_targets_next_layer(self):
|
||||
def test_indexer_does_not_start_next_layer_prefetch_before_attention_window(self):
|
||||
from sglang.srt.layers.attention.nsa import nsa_indexer
|
||||
|
||||
class FakePrefetcher:
|
||||
@@ -3121,7 +3121,34 @@ class TestCpSharedKVTaiMaterializeIntegration(unittest.TestCase):
|
||||
|
||||
indexer._maybe_start_next_layer_index_prefetch(forward_batch, layer_id=11)
|
||||
|
||||
self.assertEqual(fake_prefetcher.calls, [(12, token_to_kv_pool)])
|
||||
self.assertEqual(fake_prefetcher.calls, [])
|
||||
|
||||
def test_attention_window_prefetch_starts_index_and_mla_next_layer(self):
|
||||
from sglang.srt.layers.attention import nsa_backend
|
||||
|
||||
class FakePrefetcher:
|
||||
def __init__(self):
|
||||
self.calls = []
|
||||
|
||||
def start_next_layer_prefix(self, *, next_layer_id, token_to_kv_pool):
|
||||
self.calls.append((next_layer_id, token_to_kv_pool))
|
||||
|
||||
token_to_kv_pool = object()
|
||||
index_prefetcher = FakePrefetcher()
|
||||
mla_prefetcher = FakePrefetcher()
|
||||
forward_batch = SimpleNamespace(
|
||||
token_to_kv_pool=token_to_kv_pool,
|
||||
cp_shared_kv_index_prefetcher=index_prefetcher,
|
||||
cp_shared_kv_mla_prefetcher=mla_prefetcher,
|
||||
cp_shared_kv_num_model_layers=12,
|
||||
)
|
||||
|
||||
nsa_backend._maybe_start_cp_shared_kv_attention_prefetch(
|
||||
forward_batch, layer_id=5
|
||||
)
|
||||
|
||||
self.assertEqual(index_prefetcher.calls, [(6, token_to_kv_pool)])
|
||||
self.assertEqual(mla_prefetcher.calls, [(6, token_to_kv_pool)])
|
||||
|
||||
def test_index_prefetch_skips_when_current_layer_is_last(self):
|
||||
from sglang.srt.layers.attention.nsa import nsa_indexer
|
||||
|
||||
Reference in New Issue
Block a user