Preserve request boundaries in CP shared-KV index top-k
W4-1 needs target index/top-k sync correctness before current/partial-current reuse can be restored. Batch-size>1 in-seq CP produces local q/weights in request-segment order, so top-k must consume req0_prev, req0_next, req1_prev, req1_next rather than treating the flattened batch as one scalar prev/next pair. The implementation adds a batch dispatch for _get_topk_in_seq_cp_pair, reuses one synchronous shared-index materialization per layer, and calls _get_topk_ragged_with_cp per request segment with an explicit batch_idx. The scalar bs=1 path remains unchanged. Constraint: This is W4-1 target index/top-k sync correctness; original W4 current/partial-current reuse remains a separate follow-up. Constraint: Phase W4-1 must not enable bs>1 index prefetch, current reuse, partial-current reuse, or the cp_index multi-batch branch. Rejected: Use cp_index branch for multi-batch | source marks that path as having accuracy issues. Rejected: Pad batch requests to max length | wastes compute and violates packed/ragged batch contract. Confidence: high Scope-risk: moderate Directive: Keep bs>1 target top-k ordered by request segment unless a later fused descriptor proves identical ordering and correctness. Tested: Remote g0034 py_compile for nsa_indexer.py Tested: Remote g0034 PYTHONPATH=python pytest test/registered/unit/layers/test_nsa_cp_utils.py -> 45 passed Tested: Remote g0034 PYTHONPATH=python pytest test_nsa_cp_utils.py test_cp_shared_kv_layout.py test_cp_shared_kv_runtime.py -> 172 passed, 2 subtests passed Not-tested: Full ETE bs>1 serving run with live traffic
This commit is contained in:
laoyao0822
@@ -1090,6 +1090,156 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
|
||||
self.assertEqual(topk_calls[1]["actual_seq_q_cu_tensor"].tolist(), [0, 2])
|
||||
self.assertEqual(result.tolist(), [[1, 1], [1, 1], [1, 1], [2, 2], [2, 2]])
|
||||
|
||||
def test_indexer_in_seq_cp_pair_batch_preserves_request_segment_order(self):
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
|
||||
|
||||
indexer = object.__new__(Indexer)
|
||||
indexer.index_topk = 2
|
||||
logical_pages = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.int32)
|
||||
materialized_index = torch.tensor([11], dtype=torch.int32)
|
||||
dense_pages = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.int32)
|
||||
materialize_calls = []
|
||||
topk_calls = []
|
||||
|
||||
class Metadata:
|
||||
def get_page_table_64(self):
|
||||
return logical_pages
|
||||
|
||||
def fake_materialize(forward_batch, layer_id, logical_page_table):
|
||||
materialize_calls.append((layer_id, logical_page_table))
|
||||
return materialized_index, dense_pages
|
||||
|
||||
def fake_get_topk(
|
||||
forward_batch,
|
||||
layer_id,
|
||||
q_fp8,
|
||||
weights,
|
||||
metadata,
|
||||
kv_len,
|
||||
actual_seq_q,
|
||||
cp_index=None,
|
||||
current_index_kv=None,
|
||||
shared_index_buffer=None,
|
||||
shared_block_tables=None,
|
||||
actual_seq_q_tensor=None,
|
||||
actual_seq_q_cu_tensor=None,
|
||||
batch_idx=0,
|
||||
):
|
||||
topk_calls.append(
|
||||
{
|
||||
"batch_idx": batch_idx,
|
||||
"kv_len": kv_len,
|
||||
"actual_seq_q": actual_seq_q,
|
||||
"q": q_fp8.flatten().tolist(),
|
||||
"weights": weights.flatten().tolist(),
|
||||
"actual_seq_q_tensor": actual_seq_q_tensor,
|
||||
"actual_seq_q_cu_tensor": actual_seq_q_cu_tensor,
|
||||
"shared_index_buffer": shared_index_buffer,
|
||||
"shared_block_tables": shared_block_tables,
|
||||
"current_index_kv": current_index_kv,
|
||||
}
|
||||
)
|
||||
return torch.full(
|
||||
(actual_seq_q, 2),
|
||||
len(topk_calls),
|
||||
dtype=torch.int32,
|
||||
)
|
||||
|
||||
indexer._maybe_materialize_shared_index_buffer = fake_materialize
|
||||
indexer._get_topk_ragged_with_cp = fake_get_topk
|
||||
|
||||
forward_batch = SimpleNamespace(
|
||||
batch_size=2,
|
||||
nsa_cp_metadata=NSAContextParallelMetadata(
|
||||
batch_size=2,
|
||||
kv_len_prev=100,
|
||||
kv_len_next=200,
|
||||
actual_seq_q_prev=2,
|
||||
actual_seq_q_next=1,
|
||||
actual_seq_q_prev_cu_tensor=torch.tensor([0, 2], dtype=torch.int32),
|
||||
actual_seq_q_next_cu_tensor=torch.tensor([0, 1], dtype=torch.int32),
|
||||
request_kv_len_prev=[100, 300],
|
||||
request_kv_len_next=[200, 400],
|
||||
request_actual_seq_q_prev=[2, 1],
|
||||
request_actual_seq_q_next=[1, 3],
|
||||
),
|
||||
)
|
||||
q_fp8 = torch.arange(7, dtype=torch.float32).view(7, 1)
|
||||
weights = (torch.arange(7, dtype=torch.float32) + 100).view(7, 1)
|
||||
|
||||
result = Indexer._get_topk_in_seq_cp_pair(
|
||||
indexer,
|
||||
forward_batch,
|
||||
layer_id=7,
|
||||
q_fp8=q_fp8,
|
||||
weights=weights,
|
||||
metadata=Metadata(),
|
||||
current_index_kv=None,
|
||||
)
|
||||
|
||||
self.assertEqual(len(materialize_calls), 1)
|
||||
self.assertIs(materialize_calls[0][1], logical_pages)
|
||||
self.assertEqual(
|
||||
[
|
||||
(
|
||||
call["batch_idx"],
|
||||
call["kv_len"],
|
||||
call["actual_seq_q"],
|
||||
call["q"],
|
||||
call["weights"],
|
||||
call["actual_seq_q_tensor"].tolist(),
|
||||
call["actual_seq_q_cu_tensor"].tolist(),
|
||||
)
|
||||
for call in topk_calls
|
||||
],
|
||||
[
|
||||
(0, 100, 2, [0.0, 1.0], [100.0, 101.0], [2], [0, 2]),
|
||||
(0, 200, 1, [2.0], [102.0], [1], [0, 1]),
|
||||
(1, 300, 1, [3.0], [103.0], [1], [0, 1]),
|
||||
(1, 400, 3, [4.0, 5.0, 6.0], [104.0, 105.0, 106.0], [3], [0, 3]),
|
||||
],
|
||||
)
|
||||
self.assertTrue(all(call["shared_index_buffer"] is materialized_index for call in topk_calls))
|
||||
self.assertTrue(all(call["shared_block_tables"] is dense_pages for call in topk_calls))
|
||||
self.assertTrue(all(call["current_index_kv"] is None for call in topk_calls))
|
||||
self.assertEqual(
|
||||
result.tolist(),
|
||||
[[1, 1], [1, 1], [2, 2], [3, 3], [4, 4], [4, 4], [4, 4]],
|
||||
)
|
||||
|
||||
def test_indexer_in_seq_cp_pair_batch_rejects_current_index_reuse(self):
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
|
||||
|
||||
indexer = object.__new__(Indexer)
|
||||
forward_batch = SimpleNamespace(
|
||||
batch_size=2,
|
||||
nsa_cp_metadata=NSAContextParallelMetadata(
|
||||
batch_size=2,
|
||||
request_kv_len_prev=[100, 300],
|
||||
request_kv_len_next=[200, 400],
|
||||
request_actual_seq_q_prev=[2, 1],
|
||||
request_actual_seq_q_next=[1, 3],
|
||||
),
|
||||
)
|
||||
|
||||
with self.assertRaisesRegex(
|
||||
RuntimeError,
|
||||
"CP_SHARED_KV_FAIL_FAST.*batch_gt1_index_current_reuse_unsupported",
|
||||
):
|
||||
Indexer._get_topk_in_seq_cp_pair(
|
||||
indexer,
|
||||
forward_batch,
|
||||
layer_id=7,
|
||||
q_fp8=torch.empty(7, 1),
|
||||
weights=torch.empty(7, 1),
|
||||
metadata=SimpleNamespace(),
|
||||
current_index_kv=(torch.empty(1), torch.empty(1)),
|
||||
)
|
||||
|
||||
def test_indexer_in_seq_cp_pair_skips_materialize_when_current_index_reused(self):
|
||||
import torch
|
||||
|
||||
|
||||
Reference in New Issue
Block a user