From 69ca7045ea0d894772aeeea94721a6d441bc2ac8 Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Sat, 13 Jun 2026 02:34:55 +0800 Subject: [PATCH] Preserve chunked request affinity state 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. --- python/sglang/srt/managers/schedule_policy.py | 3 ++ python/sglang/srt/managers/scheduler.py | 6 ++-- .../unit/managers/test_prefill_adder.py | 33 +++++++++++++++++++ 3 files changed, 39 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index abae4051a..32c7642ab 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -500,6 +500,7 @@ class PrefillAdder: self.can_run_list = [] self.preempt_list = [] self.new_chunked_req = None + self.chunked_req_in_batch = None self.has_chunked_req_in_batch = False self.log_hit_tokens = 0 # TODO(lsyin): report the real input tokens excluding page alignment @@ -924,6 +925,7 @@ class PrefillAdder: req.set_extend_input_len(min(req.extend_input_len, _rem_tokens)) req.fill_ids = req.fill_ids[: len(req.prefix_indices) + req.extend_input_len] self.can_run_list.append(req) + self.chunked_req_in_batch = req self.has_chunked_req_in_batch = True # A continued chunk carries a real page-floored prefix. Keep that # footprint in the CP cached/buffer admission state even though the @@ -1187,6 +1189,7 @@ class PrefillAdder: self.can_run_list.append(req) self.new_chunked_req = req + self.chunked_req_in_batch = req self.has_chunked_req_in_batch = True self._req_inc_lock_ref(req) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 0d6b4bb9d..0fdcb8d60 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -2584,13 +2584,13 @@ class Scheduler( affinity_window_used = 0 affinity_head_pending = True # first classified candidate = FCFS head affinity_batch_warm_led: Optional[bool] = None - if affinity_on and adder._chunked_req_in_batch is not None: + if affinity_on and adder.chunked_req_in_batch is not None: # A chunked-led batch must carry the CHUNK's class: a 65536 first # chunk is cold-led (further colds are free density and the caps # bound them); a tail chunk over a big cached prefix is warm-led. # Latching from the first loop candidate instead would invert # cold-led-admits-everything and spuriously STOP the scan. - chunked = adder._chunked_req_in_batch + chunked = adder.chunked_req_in_batch affinity_batch_warm_led = ( len(chunked.prefix_indices) + int(getattr(chunked, "host_hit_length", 0) or 0) @@ -2656,7 +2656,7 @@ class Scheduler( is_warm=req_is_warm, batch_empty_for_affinity=( len(adder.can_run_list) - - (1 if adder._chunked_req_in_batch is not None else 0) + - (1 if adder.chunked_req_in_batch is not None else 0) == 0 ), # None = nothing admitted yet; only reachable together diff --git a/test/registered/unit/managers/test_prefill_adder.py b/test/registered/unit/managers/test_prefill_adder.py index d72e40b15..c62b4ba99 100644 --- a/test/registered/unit/managers/test_prefill_adder.py +++ b/test/registered/unit/managers/test_prefill_adder.py @@ -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(