Gate CP shared-KV prefill batching behind explicit limits
The scheduler can now admit multi-request NSA in-seq CP shared-KV prefill batches only when the shared-KV bs>1 flag is explicitly enabled. The gate is still disabled by default and is scoped to CP shared-KV so ordinary CP is not widened accidentally. Batch admission is bounded by optional request-count and page-aligned extend-token limits while real memory capacity remains allocator-owned. This keeps bf16 and fp8 on the same scheduler path because dtype differences are already reflected in KV pool token/page capacity. Constraint: bs>1 runtime paths remain guarded by existing CP shared-KV fail-fast checks. Constraint: Scheduler must not duplicate bf16/fp8 byte-level capacity estimation. Rejected: Open the old CP gate unconditionally | ordinary CP would inherit an unverified shared-KV-specific batching path. Rejected: Treat the extend-token cap as a hard per-request limit | a single large request could deadlock the scheduler. Confidence: medium Scope-risk: moderate Directive: Keep CP shared-KV batching gated until ETE validates EAGLE accept length, output length, and HiCache load/backup behavior under real traffic. Tested: local py_compile for server_args, schedule_policy, scheduler, prefill_adder tests, and server_args tests. Tested: remote g0034 py_compile for the same files. Tested: remote g0034 pytest target set: 5 passed for parser, parameter validation, default single-request CP gate, enabled bs>1 gate, and page-aligned extend cap. Tested: remote g0034 pytest test_prefill_adder.py => 13 passed. Not-tested: full server_args test file has an unrelated HuggingFace DNS/config-download failure in TestPrepareServerArgs.test_prepare_server_args. Not-tested: ETE production traffic with --enable-cp-shared-kv-prefill-bs-gt1.
This commit is contained in:
@@ -53,6 +53,28 @@ def test_enable_nsa_prefill_cp_shared_kv_parser_flag():
|
||||
assert args.enable_nsa_prefill_cp_shared_kv is True
|
||||
|
||||
|
||||
def test_cp_shared_kv_prefill_bs_gt1_parser_limits():
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
ServerArgs.add_cli_args(parser)
|
||||
raw_args = parser.parse_args(
|
||||
[
|
||||
"--model-path",
|
||||
"dummy",
|
||||
"--enable-cp-shared-kv-prefill-bs-gt1",
|
||||
"--cp-shared-kv-prefill-max-batch-requests",
|
||||
"4",
|
||||
"--cp-shared-kv-prefill-max-total-extend-tokens",
|
||||
"8192",
|
||||
]
|
||||
)
|
||||
args = ServerArgs.from_cli_args(raw_args)
|
||||
assert args.enable_cp_shared_kv_prefill_bs_gt1 is True
|
||||
assert args.cp_shared_kv_prefill_max_batch_requests == 4
|
||||
assert args.cp_shared_kv_prefill_max_total_extend_tokens == 8192
|
||||
|
||||
|
||||
class TestLoadBalanceMethod(unittest.TestCase):
|
||||
def test_non_pd_defaults_to_round_robin(self):
|
||||
server_args = ServerArgs(model_path="dummy", disaggregation_mode="null")
|
||||
@@ -438,6 +460,23 @@ class TestHiCacheArgs(CustomTestCase):
|
||||
enable_hisparse=False,
|
||||
)
|
||||
|
||||
def test_cp_shared_kv_prefill_batch_limits_must_be_positive(self):
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, "cp_shared_kv_prefill_max_batch_requests.*positive"
|
||||
):
|
||||
ServerArgs(
|
||||
model_path="dummy",
|
||||
cp_shared_kv_prefill_max_batch_requests=0,
|
||||
)
|
||||
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, "cp_shared_kv_prefill_max_total_extend_tokens.*positive"
|
||||
):
|
||||
ServerArgs(
|
||||
model_path="dummy",
|
||||
cp_shared_kv_prefill_max_total_extend_tokens=0,
|
||||
)
|
||||
|
||||
def test_hicache_io_backend_and_mem_layout_compatibility(self):
|
||||
cases = [
|
||||
{
|
||||
|
||||
Reference in New Issue
Block a user