3a43727216
CP shared-KV bs>1 batching was only bounded by request count, extend tokens, and cached tokens. That left temporary GPU buffers such as MLA/index materialization, remap metadata, logits windows, and transfer descriptors implicit, and raw extend-token limits could exceed the active chunked-prefill budget.\n\nThis adds an explicit max-buffer-size admission gate with a CPU-only stream-aware estimator, wires it through PrefillAdder/Scheduler, performs a startup CUDA smoke allocation when configured, and reports the estimate in the scheduler admission benchmark. When chunked prefill is active, the effective CP extend-token limit is capped by the current chunk budget so the CP path does not advertise unreachable batch capacity or lift max-prefill-tokens too far.\n\nConstraint: Admission estimation must stay CPU-only on the scheduler hot path; CUDA allocation is limited to startup smoke checking.\nConstraint: Single oversized requests must still be allowed to run alone to avoid scheduler deadlock.\nRejected: Rely only on --max-prefill-tokens | it does not reliably bound the first oversized request and does not model cache-hit/load-back pressure.\nRejected: Let CP extend limit exceed chunked-prefill size | it creates an unreachable effective capacity and misleading budget lift.\nConfidence: medium\nScope-risk: moderate\nDirective: If bs>1 L1 prefetch is enabled later, update CPSharedKVPrefillBufferEstimatorContext.bs_gt1_l1_prefetch_enabled and include the live prefetch dense buffers in overlap windows.\nTested: local py_compile for touched files\nTested: local PYTHONPATH=python pytest -q test/registered/unit/managers/test_cp_shared_kv_prefill_buffer_estimator.py (4 passed)\nTested: remote cjy-glm5-new targeted pytest for new server_args, PrefillAdder, estimator, and benchmark cases (10 passed)\nTested: remote cjy-glm5-new PYTHONPATH=python pytest -q test/registered/unit/managers/test_cp_shared_kv_prefill_buffer_estimator.py test/registered/unit/managers/test_prefill_adder.py test/registered/unit/managers/test_prefill_scheduler_admission_bench.py (29 passed before chunk cap, then test_prefill_adder.py 21 passed after chunk cap)\nNot-tested: full server_args suite because existing TestPrepareServerArgs tries to reach HuggingFace and fails under container DNS/network\nNot-tested: GLM5 ETE smoke with --cp-shared-kv-prefill-max-buffer-size
175 lines
6.4 KiB
Python
175 lines
6.4 KiB
Python
import importlib.util
|
|
import sys
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
_REPO_ROOT = Path(__file__).resolve().parents[4]
|
|
sys.path.insert(0, str(_REPO_ROOT / "python"))
|
|
|
|
|
|
_BENCH_PATH = (
|
|
_REPO_ROOT
|
|
/ "benchmark"
|
|
/ "hicache"
|
|
/ "bench_prefill_scheduler_admission.py"
|
|
)
|
|
|
|
|
|
def _load_bench_module():
|
|
spec = importlib.util.spec_from_file_location(
|
|
"bench_prefill_scheduler_admission", _BENCH_PATH
|
|
)
|
|
module = importlib.util.module_from_spec(spec)
|
|
assert spec.loader is not None
|
|
sys.modules[spec.name] = module
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
class TestPrefillSchedulerAdmissionBench(unittest.TestCase):
|
|
def test_l1_l2_and_extend_tokens_are_reported_with_real_prefill_semantics(self):
|
|
bench = _load_bench_module()
|
|
cfg = bench.SchedulerBenchConfig(
|
|
page_size=64,
|
|
available_tokens=10000,
|
|
evictable_tokens=0,
|
|
max_prefill_tokens=1024,
|
|
enable_cp_shared_kv_prefill_bs_gt1=True,
|
|
cp_shared_kv_prefill_max_batch_requests=8,
|
|
cp_shared_kv_prefill_max_total_extend_tokens=4096,
|
|
)
|
|
trace = bench.run_scheduler_admission_trace(
|
|
[
|
|
bench.RequestSpec(
|
|
rid="hit-l1-l2",
|
|
l1_cached_tokens=128,
|
|
l2_cached_tokens=128,
|
|
extend_tokens=64,
|
|
max_new_tokens=1,
|
|
)
|
|
],
|
|
cfg,
|
|
)
|
|
|
|
self.assertEqual(len(trace.ticks), 1)
|
|
tick = trace.ticks[0]
|
|
self.assertEqual([req.rid for req in tick.accepted], ["hit-l1-l2"])
|
|
accepted = tick.accepted[0]
|
|
self.assertEqual(accepted.initial_extend_tokens, 192)
|
|
self.assertEqual(accepted.effective_extend_tokens, 64)
|
|
self.assertEqual(accepted.l1_cached_tokens, 128)
|
|
self.assertEqual(accepted.l2_cached_tokens, 128)
|
|
self.assertEqual(accepted.loaded_l2_tokens, 128)
|
|
self.assertEqual(tick.log_hit_tokens, 256)
|
|
self.assertEqual(tick.log_input_tokens, 64)
|
|
|
|
def test_cp_total_extend_limit_controls_batching_not_generic_max_prefill_tokens(self):
|
|
bench = _load_bench_module()
|
|
cfg = bench.SchedulerBenchConfig(
|
|
page_size=64,
|
|
available_tokens=10000,
|
|
evictable_tokens=0,
|
|
max_prefill_tokens=192,
|
|
enable_cp_shared_kv_prefill_bs_gt1=True,
|
|
cp_shared_kv_prefill_max_batch_requests=8,
|
|
cp_shared_kv_prefill_max_total_extend_tokens=256,
|
|
)
|
|
trace = bench.run_scheduler_admission_trace(
|
|
[
|
|
bench.RequestSpec("a", 0, 0, 128, max_new_tokens=1),
|
|
bench.RequestSpec("b", 0, 0, 128, max_new_tokens=1),
|
|
],
|
|
cfg,
|
|
)
|
|
|
|
self.assertEqual(len(trace.ticks), 1)
|
|
self.assertEqual([req.rid for req in trace.ticks[0].accepted], ["a", "b"])
|
|
self.assertEqual(trace.ticks[0].cp_total_extend_tokens, 256)
|
|
|
|
def test_l2_load_back_consumes_l1_capacity_and_can_stop_later_requests(self):
|
|
bench = _load_bench_module()
|
|
cfg = bench.SchedulerBenchConfig(
|
|
page_size=64,
|
|
available_tokens=320,
|
|
evictable_tokens=0,
|
|
max_prefill_tokens=4096,
|
|
enable_cp_shared_kv_prefill_bs_gt1=True,
|
|
cp_shared_kv_prefill_max_batch_requests=8,
|
|
cp_shared_kv_prefill_max_total_extend_tokens=4096,
|
|
)
|
|
trace = bench.run_scheduler_admission_trace(
|
|
[
|
|
bench.RequestSpec("l2-heavy", 0, 128, 64, max_new_tokens=1),
|
|
bench.RequestSpec("next", 0, 0, 64, max_new_tokens=1),
|
|
],
|
|
cfg,
|
|
)
|
|
|
|
self.assertEqual([req.rid for req in trace.ticks[0].accepted], ["l2-heavy"])
|
|
self.assertEqual(trace.ticks[0].stopped_on_rid, "next")
|
|
self.assertEqual(trace.ticks[0].stopped_result, "NO_TOKEN")
|
|
self.assertEqual(trace.ticks[0].allocator_available_after_tick, 192)
|
|
self.assertEqual(trace.ticks[0].load_back_events[0].loaded_tokens, 128)
|
|
|
|
def test_total_cached_limit_is_observable_in_scheduler_trace(self):
|
|
bench = _load_bench_module()
|
|
cfg = bench.SchedulerBenchConfig(
|
|
page_size=64,
|
|
available_tokens=10000,
|
|
evictable_tokens=0,
|
|
max_prefill_tokens=4096,
|
|
enable_cp_shared_kv_prefill_bs_gt1=True,
|
|
cp_shared_kv_prefill_max_batch_requests=8,
|
|
cp_shared_kv_prefill_max_total_extend_tokens=4096,
|
|
cp_shared_kv_prefill_max_total_cached_tokens=4096,
|
|
)
|
|
trace = bench.run_scheduler_admission_trace(
|
|
[
|
|
bench.RequestSpec("a", 4096, 0, 64, max_new_tokens=1),
|
|
bench.RequestSpec("b", 4096, 0, 64, max_new_tokens=1),
|
|
],
|
|
cfg,
|
|
)
|
|
|
|
self.assertEqual([req.rid for req in trace.ticks[0].accepted], ["a"])
|
|
self.assertEqual(trace.ticks[0].stopped_on_rid, "b")
|
|
self.assertEqual(trace.ticks[0].stopped_result, "OTHER")
|
|
self.assertEqual(trace.ticks[0].cp_total_cached_tokens, 4096)
|
|
|
|
def test_max_buffer_size_is_observable_in_scheduler_trace(self):
|
|
bench = _load_bench_module()
|
|
cfg = bench.SchedulerBenchConfig(
|
|
page_size=64,
|
|
available_tokens=10000,
|
|
evictable_tokens=0,
|
|
max_prefill_tokens=4096,
|
|
enable_cp_shared_kv_prefill_bs_gt1=True,
|
|
cp_shared_kv_prefill_max_batch_requests=8,
|
|
cp_shared_kv_prefill_max_total_extend_tokens=4096,
|
|
cp_shared_kv_prefill_max_total_cached_tokens=4096,
|
|
cp_shared_kv_prefill_max_buffer_size=1,
|
|
kv_cache_dim=1,
|
|
kv_dtype_bytes=2,
|
|
index_head_dim=8,
|
|
vocab_size=16,
|
|
tp_size=1,
|
|
)
|
|
trace = bench.run_scheduler_admission_trace(
|
|
[
|
|
bench.RequestSpec("a", 0, 0, 64, max_new_tokens=1),
|
|
bench.RequestSpec("b", 0, 0, 64, max_new_tokens=1),
|
|
],
|
|
cfg,
|
|
)
|
|
|
|
tick = trace.ticks[0]
|
|
self.assertEqual([req.rid for req in tick.accepted], ["a"])
|
|
self.assertEqual(tick.stopped_on_rid, "b")
|
|
self.assertEqual(tick.stopped_result, "OTHER")
|
|
self.assertGreater(tick.cp_estimated_peak_buffer_bytes, 1)
|
|
self.assertIn("layer_forward_peak_bytes", tick.cp_buffer_breakdown)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|