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

@@ -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)

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

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(