f75ffff8d9
Cache-hit GSM8K regressions only appeared after the second pass reused request-specific suffix pages, so this change adds fail-fast transfer validation, masks stale rectangular page-table tails, and extends CUDA/unit coverage across FP8 CP shared-KV write, load, top-k, and materialization paths. The temporary ledger records eliminated hypotheses to prevent re-debugging the same L2 and persistent-cache paths.\n\nConstraint: CP shared KV stores physical pages but scheduler-visible semantics must remain valid-token/page-bounded.\nConstraint: bs>1 FP8 prefill must preserve existing CP shared-KV fast paths without silent fallback.\nRejected: Blame raw HiCache L2 load without tests | L2 KV and index backup/load/materialize roundtrips pass on remote CUDA.\nRejected: Disable current/partial reuse broadly | hides the cache-hit contract regression and costs performance.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not weaken CP shared-KV fail-fast or rectangular-tail masking without rerunning second-pass cache-hit accuracy tests.\nTested: remote CUDA pytest for fused FP8 MLA store, fused persistent index store, L2-loaded FP8 KV materialize, L2-loaded index materialize, ragged top-k offset, TAI batched index MQA prepare.\nTested: local py_compile for touched test files and git diff --check.\nNot-tested: full second-pass GSM8K accuracy after these diagnostic tests; root cause remains under investigation.
201 lines
7.7 KiB
Python
201 lines
7.7 KiB
Python
"""Unit tests for prefill bootstrap queue compaction."""
|
|
|
|
import unittest
|
|
from types import SimpleNamespace
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
|
|
from sglang.srt.disaggregation.base import KVPoll
|
|
from sglang.srt.disaggregation.prefill import PrefillBootstrapQueue
|
|
from sglang.srt.disaggregation.utils import poll_and_all_reduce_attn_cp_tp_group
|
|
from sglang.srt.disaggregation.utils import ReqToMetadataIdxAllocator
|
|
from sglang.test.ci.ci_register import register_cpu_ci
|
|
from sglang.test.test_utils import CustomTestCase
|
|
|
|
register_cpu_ci(est_time=6, suite="stage-a-test-cpu")
|
|
|
|
|
|
class FakeSender:
|
|
def __init__(self, should_fail=False):
|
|
self.should_fail = should_fail
|
|
self.init_calls = []
|
|
|
|
def init(self, num_pages, metadata_buffer_index):
|
|
self.init_calls.append((num_pages, metadata_buffer_index))
|
|
|
|
def failure_exception(self):
|
|
if self.should_fail:
|
|
raise RuntimeError("boom")
|
|
|
|
def poll(self):
|
|
return KVPoll.Success
|
|
|
|
|
|
class TestPrefillBootstrapQueue(CustomTestCase):
|
|
def _make_req(self, rid, bootstrap_room, origin_input_ids, sender):
|
|
req = MagicMock()
|
|
req.rid = rid
|
|
req.bootstrap_room = bootstrap_room
|
|
req.origin_input_ids = origin_input_ids
|
|
req.disagg_kv_sender = sender
|
|
req.return_logprob = False
|
|
req.metadata_buffer_index = -1
|
|
req.time_stats = SimpleNamespace(
|
|
set_bootstrap_done_time=MagicMock(),
|
|
set_wait_queue_entry_time=MagicMock(),
|
|
trace_ctx=SimpleNamespace(abort=MagicMock()),
|
|
)
|
|
return req
|
|
|
|
def _build_queue(self, allocator_size=4):
|
|
queue = PrefillBootstrapQueue.__new__(PrefillBootstrapQueue)
|
|
queue.req_to_metadata_buffer_idx_allocator = ReqToMetadataIdxAllocator(
|
|
size=allocator_size
|
|
)
|
|
queue.token_to_kv_pool = SimpleNamespace(page_size=4)
|
|
queue.scheduler = SimpleNamespace(
|
|
stream_output=MagicMock(),
|
|
enable_metrics=False,
|
|
enable_hicache_storage=False,
|
|
attn_cp_cpu_group=MagicMock(),
|
|
attn_tp_cpu_group=MagicMock(),
|
|
)
|
|
queue.tp_rank = 0
|
|
queue.queue = []
|
|
return queue
|
|
|
|
def test_pop_bootstrapped_compacts_queue_in_one_pass(self):
|
|
queue = self._build_queue(allocator_size=2)
|
|
|
|
waiting = self._make_req("wait", 1, list(range(8)), FakeSender())
|
|
ready = self._make_req("ready", 2, list(range(10)), FakeSender())
|
|
failed = self._make_req("failed", 3, list(range(6)), FakeSender(True))
|
|
queue.queue = [waiting, ready, failed]
|
|
|
|
with patch(
|
|
"sglang.srt.disaggregation.prefill.poll_and_all_reduce_attn_cp_tp_group",
|
|
return_value=[
|
|
KVPoll.Bootstrapping,
|
|
KVPoll.WaitingForInput,
|
|
KVPoll.Failed,
|
|
],
|
|
):
|
|
bootstrapped, failed_reqs = queue.pop_bootstrapped(return_failed_reqs=True)
|
|
|
|
self.assertEqual(bootstrapped, [ready])
|
|
self.assertEqual(failed_reqs, [failed])
|
|
self.assertEqual(queue.queue, [waiting])
|
|
self.assertEqual(ready.metadata_buffer_index, 0)
|
|
self.assertEqual(ready.disagg_kv_sender.init_calls, [(3, 0)])
|
|
|
|
def test_pop_bootstrapped_stops_on_metadata_exhaustion(self):
|
|
queue = self._build_queue(allocator_size=1)
|
|
|
|
first_ready = self._make_req("ready-1", 1, list(range(8)), FakeSender())
|
|
second_ready = self._make_req("ready-2", 2, list(range(12)), FakeSender())
|
|
trailing = self._make_req("trailing", 3, list(range(4)), FakeSender())
|
|
queue.queue = [first_ready, second_ready, trailing]
|
|
|
|
with patch(
|
|
"sglang.srt.disaggregation.prefill.poll_and_all_reduce_attn_cp_tp_group",
|
|
return_value=[
|
|
KVPoll.WaitingForInput,
|
|
KVPoll.WaitingForInput,
|
|
KVPoll.Bootstrapping,
|
|
],
|
|
):
|
|
bootstrapped = queue.pop_bootstrapped()
|
|
|
|
self.assertEqual(bootstrapped, [first_ready])
|
|
self.assertEqual(queue.queue, [second_ready, trailing])
|
|
self.assertEqual(first_ready.disagg_kv_sender.init_calls, [(2, 0)])
|
|
self.assertEqual(second_ready.disagg_kv_sender.init_calls, [])
|
|
|
|
def test_pop_bootstrapped_preserves_unchecked_entries(self):
|
|
queue = self._build_queue(allocator_size=2)
|
|
|
|
skipped = self._make_req("skip", 1, list(range(8)), FakeSender())
|
|
checked = self._make_req("check", 2, list(range(10)), FakeSender())
|
|
queue.queue = [skipped, checked]
|
|
|
|
with patch(
|
|
"sglang.srt.disaggregation.prefill.poll_and_all_reduce_attn_cp_tp_group",
|
|
return_value=[KVPoll.WaitingForInput, KVPoll.WaitingForInput],
|
|
):
|
|
bootstrapped = queue.pop_bootstrapped(rids_to_check=["check"])
|
|
|
|
self.assertEqual(bootstrapped, [checked])
|
|
self.assertEqual(queue.queue, [skipped])
|
|
self.assertEqual(skipped.disagg_kv_sender.init_calls, [])
|
|
self.assertEqual(checked.disagg_kv_sender.init_calls, [(3, 0)])
|
|
|
|
def test_poll_consensus_debug_fails_before_shape_mismatch_hang(self):
|
|
reduce_calls = []
|
|
|
|
def fake_all_reduce(tensor, op, group):
|
|
reduce_calls.append((tensor.dtype, int(tensor.numel()), op, group))
|
|
# The new debug guard uses an int64 [queue_len, queue_hash] scalar
|
|
# vector before building the uint8 per-request poll tensor. Simulate
|
|
# another rank having a different queue length and assert that we
|
|
# fail fast before reaching the old variable-length uint8 all_reduce.
|
|
if tensor.dtype == torch.int64 and int(tensor.numel()) == 2:
|
|
if op == dist.ReduceOp.MIN:
|
|
tensor[0] = 1
|
|
elif op == dist.ReduceOp.MAX:
|
|
tensor[0] = 2
|
|
return
|
|
raise AssertionError(
|
|
"poll tensor all_reduce should not run after queue mismatch"
|
|
)
|
|
|
|
with (
|
|
patch(
|
|
"sglang.srt.disaggregation.utils.envs.SGLANG_CP_SHARED_KV_BS_GT1_DEBUG.get",
|
|
return_value=True,
|
|
),
|
|
patch("sglang.srt.disaggregation.utils.dist.all_reduce", fake_all_reduce),
|
|
):
|
|
with self.assertRaisesRegex(RuntimeError, "poll_queue.*inflight"):
|
|
poll_and_all_reduce_attn_cp_tp_group(
|
|
[FakeSender(), FakeSender()],
|
|
MagicMock(name="cp_group"),
|
|
MagicMock(name="tp_group"),
|
|
debug_label="inflight",
|
|
debug_ids=["rid-a", "rid-b"],
|
|
)
|
|
|
|
self.assertTrue(reduce_calls)
|
|
|
|
def test_poll_consensus_debug_disabled_preserves_old_collective_shape(self):
|
|
reduce_calls = []
|
|
|
|
def fake_all_reduce(tensor, op, group):
|
|
reduce_calls.append((tensor.dtype, int(tensor.numel()), op, group))
|
|
|
|
with (
|
|
patch(
|
|
"sglang.srt.disaggregation.utils.envs.SGLANG_CP_SHARED_KV_BS_GT1_DEBUG.get",
|
|
return_value=False,
|
|
),
|
|
patch("sglang.srt.disaggregation.utils.dist.all_reduce", fake_all_reduce),
|
|
):
|
|
polls = poll_and_all_reduce_attn_cp_tp_group(
|
|
[FakeSender(), FakeSender()],
|
|
MagicMock(name="cp_group"),
|
|
MagicMock(name="tp_group"),
|
|
debug_label="inflight",
|
|
debug_ids=["rid-a", "rid-b"],
|
|
)
|
|
|
|
self.assertEqual(polls, [KVPoll.Success, KVPoll.Success])
|
|
self.assertEqual(
|
|
[(dtype, size) for dtype, size, _op, _group in reduce_calls],
|
|
[(torch.uint8, 2), (torch.uint8, 2)],
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|