Fix the _chunked_req_scheduled_last_iter flag with a content-based stash gate (#26938)

test/registered/unit/managers/test_scheduler_chunked_req_gate.py

@@ -25,13 +25,14 @@ def _make_req(
     fill_ids: list,
     prefix_indices: torch.Tensor,
     extend_input_len: int,
+    fill_len: int,
 ) -> Req:
     req = Req.__new__(Req)
     req.rid = "test-req"
     req.origin_input_ids = array("q", fill_ids)
     req.output_ids = array("q")
     req.full_untruncated_fill_ids = array("q", fill_ids)
-    req.fill_len = len(req.full_untruncated_fill_ids)
+    req.fill_len = fill_len
     req.prefix_indices = prefix_indices
     req.req_pool_idx = req_pool_idx
     req.extend_input_len = extend_input_len
@@ -104,15 +105,17 @@ def _scheduler_for_get_next_batch(*, tree_cache, chunked_req) -> Scheduler:
 class TestStashGatePreservesPrefixIndices(CustomTestCase):
     """Consumer side: real ChunkCache.cache_unfinished_req mutates
     req.prefix_indices iff stash actually runs, so prefix_indices content
-    is the bug-detection signal."""
+    is the bug-detection signal. The stash gate is content-based:
+    `fill_len > len(prefix_indices)` means there is freshly computed KV to
+    cache; otherwise the chunk was parked and stashing must be skipped."""
 
     POOL_IDX = 4
     INITIAL_PREFIX_LEN = 8  # what was really cached last iter
-    POST_RESET_FILL_LEN = 32  # length after init_next_round_input
+    POST_RESET_FILL_LEN = 32  # length after init_next_round_input rebuilds
     NUM_SLOTS = 8
     MAX_CONTEXT = 64
 
-    def _build(self, flag: bool):
+    def _build(self, *, fill_len: int):
         pool = _make_req_to_token_pool(self.NUM_SLOTS, self.MAX_CONTEXT)
         cache = _make_chunk_cache(pool)
         initial_prefix = pool.req_to_token[self.POOL_IDX, : self.INITIAL_PREFIX_LEN].to(
@@ -122,16 +125,16 @@ class TestStashGatePreservesPrefixIndices(CustomTestCase):
             req_pool_idx=self.POOL_IDX,
             fill_ids=list(range(self.POST_RESET_FILL_LEN)),
             prefix_indices=initial_prefix,
-            extend_input_len=0,
+            extend_input_len=fill_len - self.INITIAL_PREFIX_LEN,
+            fill_len=fill_len,
         )
         s = _scheduler_for_get_next_batch(tree_cache=cache, chunked_req=req)
-        s._chunked_req_scheduled_last_iter = flag
         return s, req, initial_prefix, pool
 
-    def test_deferred_chunked_req_keeps_real_prefix_indices(self):
-        # The bug case: a spurious stash on a deferred chunked_req
-        # would extend prefix_indices to len(fill_ids).
-        s, req, initial_prefix, _ = self._build(flag=False)
+    def test_parked_chunked_req_keeps_real_prefix_indices(self):
+        # A parked chunk has fill_len == len(prefix_indices): no new KV was
+        # computed, so the gate must skip stash and leave prefix_indices intact.
+        s, req, initial_prefix, _ = self._build(fill_len=self.INITIAL_PREFIX_LEN)
 
         Scheduler.get_next_batch_to_run(s)
 
@@ -139,9 +142,9 @@ class TestStashGatePreservesPrefixIndices(CustomTestCase):
         self.assertTrue(torch.equal(req.prefix_indices, initial_prefix))
 
     def test_scheduled_chunked_req_advances_prefix_indices_via_real_stash(self):
-        # Symmetric guard against over-gating: when the chunked_req was
-        # actually scheduled, stash must run and advance prefix_indices.
-        s, req, _, pool = self._build(flag=True)
+        # Symmetric guard against over-gating: when fill_len has advanced past
+        # the cached prefix, stash must run and advance prefix_indices.
+        s, req, _, pool = self._build(fill_len=self.POST_RESET_FILL_LEN)
 
         Scheduler.get_next_batch_to_run(s)
 
@@ -151,13 +154,12 @@ class TestStashGatePreservesPrefixIndices(CustomTestCase):
         self.assertEqual(req.prefix_indices.shape[0], self.POST_RESET_FILL_LEN)
         self.assertTrue(torch.equal(req.prefix_indices, expected))
 
-    def test_no_chunked_req_never_mutates_state_even_with_stale_flag(self):
-        # Retract path clears chunked_req without resetting the flag;
-        # the outer `if chunked_req is not None` guard must hold.
+    def test_no_chunked_req_never_mutates_state(self):
+        # The outer `if chunked_req is not None` guard must hold on the retract
+        # path that clears chunked_req.
         pool = _make_req_to_token_pool(self.NUM_SLOTS, self.MAX_CONTEXT)
         cache = _make_chunk_cache(pool)
         s = _scheduler_for_get_next_batch(tree_cache=cache, chunked_req=None)
-        s._chunked_req_scheduled_last_iter = True
 
         Scheduler.get_next_batch_to_run(s)
         self.assertIsNone(s.chunked_req)