chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Preserve chunked request affinity state

Browse Source 
The affinity scheduler needs to know whether the current batch is led by a chunked request. The rebase carried call sites that referenced an old private field name, while PrefillAdder only retained a boolean flag, causing startup failure before scheduling could run.

Constraint: CP bs>1 affinity must classify a chunked-led batch without reopening chunked-tail mixing behavior.

Rejected: Recreate the old private _chunked_req_in_batch attribute | keeps a stale name and hides the public state transition in PrefillAdder.

Confidence: high

Scope-risk: narrow

Directive: Keep chunked_req_in_batch and has_chunked_req_in_batch updated together when adding new chunked admission paths.

Tested: Remote cjy-glm5-new: PYTHONPATH=python python -m pytest -q test/registered/unit/managers/test_prefill_adder.py -> 26 passed; combined run with scheduler load accounting tests -> 27 passed.

Not-tested: Full ETE restart after this commit alone.
This commit is contained in:
laoyao0822
2026-06-13 02:34:55 +08:00
parent ee843a946b
commit 69ca7045ea
3 changed files with 39 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

33
test/registered/unit/managers/test_prefill_adder.py
View File

@@ -835,6 +835,8 @@ class TestPrefillAdder(CustomTestCase):
adder.new_chunked_req = adder.add_chunked_req(chunked)
self.assertEqual([req.rid for req in adder.can_run_list], ["chunked"])
self.assertIsNone(adder.new_chunked_req)
self.assertIs(adder.chunked_req_in_batch, chunked)
self.assertTrue(adder.has_chunked_req_in_batch)
self.assertEqual(chunked.extend_input_len, 128)
self.assertEqual(
@@ -982,9 +984,40 @@ class TestPrefillAdder(CustomTestCase):
chunked = self.create_prefill_req("chunked", extend_input_len=128)
chunked.prefix_indices = torch.zeros((256,), dtype=torch.int64)
adder.add_chunked_req(chunked)
self.assertIs(adder.chunked_req_in_batch, chunked)
self.assertEqual(adder.cp_shared_kv_prefill_total_cached_tokens, 256)
self.assertEqual(adder.cp_shared_kv_prefill_total_extend_tokens, 128)
def test_new_chunked_req_records_chunked_req_in_batch_for_affinity(self):
set_global_server_args_for_scheduler(
ServerArgs(
model_path="dummy",
enable_nsa_prefill_context_parallel=True,
nsa_prefill_cp_mode="in-seq-split",
)
)
self.mock_token_allocator.available_size.return_value = 10000
adder = self.create_adder(
self.create_running_batch(),
page_size=64,
rem_input_tokens=4096,
rem_chunk_tokens=128,
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,
)
req = self.create_prefill_req("chunked", extend_input_len=512)
self.assertEqual(
adder.add_one_req(req, has_chunked_req=False, truncation_align_size=None),
AddReqResult.OTHER,
)
self.assertIs(adder.new_chunked_req, req)
self.assertIs(adder.chunked_req_in_batch, req)
self.assertTrue(adder.has_chunked_req_in_batch)
self.assertEqual(req.extend_input_len, 128)
def test_cp_prefill_total_cached_limit_stops_second_cached_request(self):
set_global_server_args_for_scheduler(
ServerArgs(