[perf] reduce overhead of fill_ids list reconstruction and decref (#27965)

python/sglang/srt/managers/schedule_batch.py

@@ -695,11 +695,14 @@ class Req(ReqDllmMixin):
             if origin_input_ids_unpadded
             else self.origin_input_ids
         )  # Before image padding
-        # Each decode stage's output ids
+        # Each decode stage's output ids. Append-only by contract:
+        # _refresh_fill_ids infers how many output tokens are already in
+        # full_untruncated_fill_ids from lengths alone, so in-place rewrites
+        # that preserve length would silently corrupt fill_ids.
         self.output_ids = array("q")
         # Full untruncated sequence: origin + output (+ DLLM mask block).
-        # Rebuilt at the top of each init_next_round_input; admission only
+        # Kept in sync by _refresh_fill_ids; admission only updates fill_len,
-        # updates fill_len, never mutates this array's length.
+        # never mutates this array's length.
         self.full_untruncated_fill_ids = array("q")
         self.fill_len: int = 0
 
@@ -1069,6 +1072,27 @@ class Req(ReqDllmMixin):
     def get_fill_ids(self) -> array:
         return self.full_untruncated_fill_ids[: self.fill_len]
 
+    def _refresh_fill_ids(self) -> None:
+        """Keep full_untruncated_fill_ids == origin_input_ids + output_ids by
+        appending only the new output tokens.
+
+        Falls back to a full rebuild when the in-place append is invalid:
+        - aliasing: scheduler_pp_mixin assigns full_untruncated_fill_ids =
+          origin_input_ids directly, so extending in place would write output
+          tokens into the origin;
+        - lengths disagree: fresh req (array still empty), retraction
+          (output_ids reset to empty), or set_finish_with_abort (origin
+          replaced by a 1-token stub).
+        """
+        n_have_output = len(self.full_untruncated_fill_ids) - len(self.origin_input_ids)
+        if (
+            self.full_untruncated_fill_ids is not self.origin_input_ids
+            and 0 <= n_have_output <= len(self.output_ids)
+        ):
+            self.full_untruncated_fill_ids.extend(self.output_ids[n_have_output:])
+        else:
+            self.full_untruncated_fill_ids = self.origin_input_ids + self.output_ids
+
     def init_next_round_input(
         self,
         tree_cache: Optional[BasePrefixCache] = None,
@@ -1078,7 +1102,7 @@ class Req(ReqDllmMixin):
             self._init_fill_ids_for_dllm()
             self.determine_dllm_phase()
         else:
-            self.full_untruncated_fill_ids = self.origin_input_ids + self.output_ids
+            self._refresh_fill_ids()
 
         input_len = len(self.full_untruncated_fill_ids)
 
@@ -1098,14 +1122,16 @@ class Req(ReqDllmMixin):
             )
             self.logprob_start_len = -1
 
-        token_ids_to_match = self.full_untruncated_fill_ids[
+        # Pass the full array with a raw-token cap (limit) instead of slicing,
-            : self._compute_max_prefix_len(input_len)
+        # avoiding an O(context) copy per prefill-batch build.
-        ]
+        token_ids_to_match = self.full_untruncated_fill_ids
+        key_limit: Optional[int] = self._compute_max_prefix_len(input_len)
 
         # Disable prefix caching when embed overrides are present: same token IDs
         # with different override vectors must not share cached KV values.
         if self.positional_embed_overrides is not None:
             token_ids_to_match = array("q")
+            key_limit = None
 
         if tree_cache is not None:
             if cow_mamba is None:
@@ -1113,7 +1139,9 @@ class Req(ReqDllmMixin):
             match_result = tree_cache.match_prefix(
                 MatchPrefixParams(
                     key=RadixKey(
-                        token_ids=token_ids_to_match, extra_key=self.extra_key
+                        token_ids=token_ids_to_match,
+                        extra_key=self.extra_key,
+                        limit=key_limit,
                     ),
                     req=self,
                     cow_mamba=cow_mamba,
@@ -2314,7 +2342,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
         running_bs = running_batch.batch_size()
 
         for req in running_batch.reqs:
-            req.full_untruncated_fill_ids = req.origin_input_ids + req.output_ids
+            req._refresh_fill_ids()
             req.fill_len = len(req.full_untruncated_fill_ids)
             req.set_extend_input_len(1)
 

python/sglang/srt/mem_cache/mamba_radix_cache.py

@@ -507,7 +507,7 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache):
         prev_prefix_len = params.prev_prefix_len
 
         if value is None:
-            value = torch.tensor([x for x in key.token_ids], dtype=torch.int64)
+            value = torch.tensor([x for x in key.raw_token_ids()], dtype=torch.int64)
         prefix_len, mamba_exist = self._insert_helper(
             self.root_node, key, value, mamba_value, params.chunked, prev_prefix_len
         )

python/sglang/srt/mem_cache/radix_cache.py

@@ -56,13 +56,14 @@ if TYPE_CHECKING:
 class RadixKey:
     """is_bigram=True: token_ids holds raw tokens (N+1 for N bigrams); slices share one boundary token."""
 
-    __slots__ = ("token_ids", "extra_key", "is_bigram")
+    __slots__ = ("token_ids", "extra_key", "is_bigram", "limit")
 
     def __init__(
         self,
         token_ids: array[int],
         extra_key: Optional[str] = None,
         is_bigram: bool = False,
+        limit: Optional[int] = None,
     ):
         # token ids sequence (raw ints in both modes)
         self.token_ids = token_ids
@@ -70,21 +71,40 @@ class RadixKey:
         self.extra_key = extra_key
         # bigram view over token_ids: length = max(0, len(token_ids) - 1)
         self.is_bigram = is_bigram
+        # Optional cap on raw tokens: behave as if token_ids were sliced to
+        # token_ids[:limit], without the O(n) copy. None = use all tokens.
+        self.limit = limit
+
+    def _raw_len(self) -> int:
+        n = len(self.token_ids)
+        if self.limit is not None and self.limit < n:
+            return self.limit
+        return n
+
+    def raw_token_ids(self) -> array:
+        """token_ids honoring `limit` (copies only when capped)."""
+        n = self._raw_len()
+        t = self.token_ids
+        return t if n == len(t) else t[:n]
 
     def __len__(self) -> int:
+        n = self._raw_len()
         if self.is_bigram:
-            n = len(self.token_ids)
             return n - 1 if n > 0 else 0
-        return len(self.token_ids)
+        return n
 
     # TODO(Jialin): vectorize with numpy without PyLong boxing
     def __iter__(self) -> Iterator:
-        if self.is_bigram:
         t = self.token_ids
-            for i in range(len(t) - 1):
+        n = self._raw_len()
+        if self.is_bigram:
+            for i in range(n - 1 if n > 0 else 0):
                 yield (t[i], t[i + 1])
+        elif n == len(t):
+            yield from t
         else:
-            yield from self.token_ids
+            for i in range(n):
+                yield t[i]
 
     def __getitem__(self, idx: Union[int, slice]) -> RadixKey:
         # Normalize int -> 1-element slice so the rest handles one shape.
@@ -166,6 +186,7 @@ class RadixKey:
             matched = max(0, min(matched_tokens - 1, len(self), len(other)))
             return (matched // page_size) * page_size if page_size > 1 else matched
 
+        matched_tokens = min(matched_tokens, len(self), len(other))
         if page_size == 1:
             return matched_tokens
         return (matched_tokens // page_size) * page_size

python/sglang/srt/mem_cache/radix_cache_cpp.py

@@ -101,7 +101,7 @@ class RadixCacheCpp(BasePrefixCache):
     def match_prefix(self, params: MatchPrefixParams) -> MatchResult:
         key = params.key
         device_indices_vec, host_indices_length, node_gpu, node_cpu = (
-            self.tree.match_prefix(key.token_ids)
+            self.tree.match_prefix(key.raw_token_ids())
         )
         return MatchResult(
             device_indices=self._merge_tensor(device_indices_vec),

python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py

@@ -46,7 +46,8 @@ class _LMCacheLoadBackMarker:
     ``match_prefix`` call in MP mode.
     """
 
-    key: RadixKey  # page-aligned key the scheduler matched on
+    key: RadixKey  # detached snapshot of the matched key (the live query key
+    # aliases the req's growing fill_ids and must not be retained)
     value_numel: int  # number of tokens already in radix at match time
 
 
@@ -216,14 +217,17 @@ class LMCRadixCache(RadixCache):
         LMCache has tokens beyond radix. Otherwise releases
         the held read locks and returns the radix-only result.
         """
-        matched = self.lmcache_connector.lookup_kv(key.token_ids, req.rid)
+        token_ids = key.raw_token_ids()
+        matched = self.lmcache_connector.lookup_kv(token_ids, req.rid)
         if matched <= value.numel():
             # Release the read locks; keep the pending session for end_session.
             self.lmcache_connector.release_pending(req.rid)
             return base_res
 
+        if token_ids is key.token_ids:
+            token_ids = token_ids[:]
         self._mp_load_back_markers[req.rid] = _LMCacheLoadBackMarker(
-            key=key,
+            key=RadixKey(token_ids, key.extra_key, key.is_bigram),
             value_numel=int(value.numel()),
         )
         return MatchResult(
@@ -254,13 +258,14 @@ class LMCRadixCache(RadixCache):
         if uncached_len == 0:
             return base_res
 
+        token_ids = key.raw_token_ids()
         result = self._load_back(
             key=key,
             value_numel=int(value.numel()),
             uncached_len=uncached_len,
             last_node=last_node,
             load_fn=lambda sm, pp: self._ip_load_back(
-                token_ids=key.token_ids,
+                token_ids=token_ids,
                 value_numel=int(value.numel()),
                 slot_mapping=sm,
                 prefix_pad=pp,

python/sglang/srt/session/session_controller.py

@@ -15,6 +15,7 @@ from __future__ import annotations
 import logging
 import time
 import uuid
+from array import array
 from typing import TYPE_CHECKING, Dict, Optional
 
 from sglang.srt.managers.io_struct import (
@@ -94,12 +95,111 @@ class Session:
         self.req_nodes: Dict[str, SessionReqNode] = {}
         self.close_on_finish: bool = False
         self._inflight: bool = False
+        # Token-array lengths of last_req as of its finish_req. The share path
+        # appends speculatively beyond these; only finish_req confirms them, so
+        # _share_token_arrays trims back first (heals aborted turns).
+        self.committed_origin_len: Optional[int] = None
+        self.committed_unpadded_len: Optional[int] = None
+        self.committed_fill_len: Optional[int] = None
 
     def is_timed_out(self) -> bool:
         if self.timeout is None:
             return False
         return time.monotonic() - self.last_active_time > self.timeout
 
+    @staticmethod
+    def _strip_bos_token(req: TokenizedGenerateReqInput, tokenizer) -> None:
+        """Trim a leading BOS on an appended turn; shift mm offsets to match."""
+        if not (
+            tokenizer is not None
+            and req.input_ids
+            and req.input_ids[0] == tokenizer.bos_token_id
+        ):
+            return
+        req.input_ids = req.input_ids[1:]
+        if req.mm_inputs:
+            for item in req.mm_inputs.mm_items:
+                if item.offsets:
+                    if any(s == 0 for s, _ in item.offsets):
+                        logging.warning(
+                            "mm_item offset starts at 0 (BOS position), "
+                            "clamping to 0 after BOS strip"
+                        )
+                    item.offsets = [
+                        (max(0, s - 1), max(0, e - 1)) for s, e in item.offsets
+                    ]
+
+    def _share_token_arrays(self, last_req: Req, new_input_ids):
+        """Plain streaming append: reuse last_req's token arrays in place.
+
+        Trims each array back to its committed length first — an earlier turn
+        may have appended its tokens and then aborted before finish_req, and
+        req_nodes still points at last_req, so anything beyond the committed
+        lengths is unconfirmed. Then extends with last turn's output and the
+        new input. Returns (input_ids, input_ids_unpadded, carry_fill);
+        carry_fill (== the new origin) spares the first fill_ids rebuild.
+        """
+        out_tail = last_req.output_ids[: last_req.sampling_params.max_new_tokens]
+
+        input_ids = last_req.origin_input_ids
+        del input_ids[self.committed_origin_len :]
+        if last_req.origin_input_ids_unpadded is input_ids:
+            input_ids_unpadded = input_ids
+        else:
+            input_ids_unpadded = last_req.origin_input_ids_unpadded
+            del input_ids_unpadded[self.committed_unpadded_len :]
+
+        carry_fill = last_req.full_untruncated_fill_ids
+        if (
+            not isinstance(carry_fill, array)
+            or carry_fill is input_ids
+            or carry_fill is input_ids_unpadded
+        ):
+            # Unexpected type or aliased with an origin array (extending it
+            # below would double-append): let _refresh_fill_ids rebuild.
+            carry_fill = None
+        else:
+            del carry_fill[self.committed_fill_len :]
+            baked = len(carry_fill) - len(input_ids)
+            if 0 <= baked <= len(out_tail):
+                carry_fill.extend(out_tail[baked:])
+                carry_fill.extend(new_input_ids)
+            else:
+                carry_fill = None
+
+        input_ids.extend(out_tail)
+        input_ids.extend(new_input_ids)
+        if input_ids_unpadded is not input_ids:
+            input_ids_unpadded.extend(out_tail)
+            input_ids_unpadded.extend(new_input_ids)
+        return input_ids, input_ids_unpadded, carry_fill
+
+    @staticmethod
+    def _concat_token_arrays(
+        last_req: Req, req: TokenizedGenerateReqInput, session_params
+    ):
+        """Copy-based assembly for replace/offset/drop_previous_output turns."""
+        out_tail = last_req.output_ids[: last_req.sampling_params.max_new_tokens]
+
+        input_ids = last_req.origin_input_ids + out_tail
+        if session_params.drop_previous_output:
+            input_ids = last_req.origin_input_ids[:]
+        if session_params.offset and session_params.offset != 0:
+            input_ids = input_ids[: session_params.offset] + req.input_ids
+        else:
+            input_ids += req.input_ids
+
+        input_ids_unpadded = last_req.origin_input_ids_unpadded + out_tail
+        if session_params.drop_previous_output:
+            input_ids_unpadded = last_req.origin_input_ids_unpadded[:]
+        if session_params.offset and session_params.offset != 0:
+            input_ids_unpadded = (
+                input_ids_unpadded[: session_params.offset] + req.input_ids
+            )
+        else:
+            input_ids_unpadded += req.input_ids
+        return input_ids, input_ids_unpadded
+
     def create_req(
         self,
         req: TokenizedGenerateReqInput,
@@ -163,54 +263,28 @@ class Session:
                         abort_message = "Session request is appending to a request that hasn't finished."
                         logging.warning(abort_message)
 
+        carry_fill = None
         if last_req is not None:
-            # trim bos token if it is an append
+            self._strip_bos_token(req, tokenizer)
-            if (
+            # In-place sharing is only safe for the plain streaming append:
-                tokenizer is not None
+            # streaming sessions allow a single inflight request, last_req has
-                and req.input_ids
+            # finished, and the committed_* lengths recorded by finish_req let
-                and req.input_ids[0] == tokenizer.bos_token_id
+            # _share_token_arrays trim away tokens appended by an aborted turn.
-            ):
+            # offset / drop_previous_output rewrite history and must copy.
-                req.input_ids = req.input_ids[1:]
+            can_share_token_arrays = (
-                # Adjust mm_item offsets since they were computed on
+                self.streaming
-                # the pre-strip sequence (with BOS at position 0)
+                and self.committed_origin_len is not None
-                if req.mm_inputs:
+                and not session_params.drop_previous_output
-                    for item in req.mm_inputs.mm_items:
+                and not (session_params.offset and session_params.offset != 0)
-                        if item.offsets:
-                            if any(s == 0 for s, _ in item.offsets):
-                                logging.warning(
-                                    "mm_item offset starts at 0 (BOS position), "
-                                    "clamping to 0 after BOS strip"
             )
-                            item.offsets = [
+            if can_share_token_arrays:
-                                (max(0, s - 1), max(0, e - 1)) for s, e in item.offsets
+                input_ids, input_ids_unpadded, carry_fill = self._share_token_arrays(
-                            ]
+                    last_req, req.input_ids
-
-            input_ids = (
-                last_req.origin_input_ids
-                + last_req.output_ids[: last_req.sampling_params.max_new_tokens]
-            )
-
-            if session_params.drop_previous_output:
-                input_ids = last_req.origin_input_ids[:]
-
-            if session_params.offset and session_params.offset != 0:
-                input_ids = input_ids[: session_params.offset] + req.input_ids
-            else:
-                input_ids += req.input_ids
-
-            input_ids_unpadded = (
-                last_req.origin_input_ids_unpadded
-                + last_req.output_ids[: last_req.sampling_params.max_new_tokens]
-            )
-            if session_params.drop_previous_output:
-                input_ids_unpadded = last_req.origin_input_ids_unpadded[:]
-
-            if session_params.offset and session_params.offset != 0:
-                input_ids_unpadded = (
-                    input_ids_unpadded[: session_params.offset] + req.input_ids
                 )
             else:
-                input_ids_unpadded += req.input_ids
+                input_ids, input_ids_unpadded = self._concat_token_arrays(
+                    last_req, req, session_params
+                )
         else:
             input_ids = req.input_ids
             input_ids_unpadded = req.input_ids
@@ -243,6 +317,8 @@ class Session:
         if last_req is not None:
             new_req.multimodal_inputs = last_req.multimodal_inputs
         new_req.tokenizer = tokenizer
+        if carry_fill is not None:
+            new_req.full_untruncated_fill_ids = carry_fill
 
         if abort:
             new_req.set_finish_with_abort(abort_message)
@@ -263,6 +339,10 @@ class Session:
             prev_node.req.session = None
             self.req_nodes.clear()
         self.req_nodes[req.rid] = SessionReqNode(req)
+        # Confirm this req's token arrays as the session's rollback point.
+        self.committed_origin_len = len(req.origin_input_ids)
+        self.committed_unpadded_len = len(req.origin_input_ids_unpadded)
+        self.committed_fill_len = len(req.full_untruncated_fill_ids)
 
     def abort_req(self):
         """Clear inflight flag on abort (req_nodes stays unchanged)."""

python/sglang/srt/session/streaming_session.py

@@ -247,7 +247,7 @@ class StreamingSession(BasePrefixCache):
         # token_ids = get_fill_ids()[:input_len-1] (1-token logit reserve
         # already applied). min handles retract retry where committed_len
         # can exceed len(token_ids) by 1.
-        prefix_len = min(req.kv_committed_len, len(params.key.token_ids))
+        prefix_len = min(req.kv_committed_len, len(params.key))
 
         # Streaming sessions are append-only (session_controller rollback
         # ensures req_nodes always points to the last successful req).

test/registered/unit/mem_cache/test_session_token_share_unit.py

@@ -0,0 +1,163 @@
+"""Unit tests for the streaming-session in-place token-array share protocol
+(`Session.create_req` / `finish_req` / `abort_req`):
+
+- token arrays are extended in place and shared across turns (no per-turn copy);
+- committed_* lengths recorded at finish_req trim away tokens appended by a
+  turn that aborted before finishing (mid-turn and first-turn aborts);
+- max_new_tokens overshoot falls back to a fill_ids rebuild instead of
+  carrying an inconsistent array.
+"""
+
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=15, suite="base-a-test-cpu")
+
+import unittest
+from array import array
+from types import SimpleNamespace
+
+from sglang.srt.sampling.sampling_params import SamplingParams
+from sglang.srt.session.session_controller import Session
+from sglang.test.test_utils import CustomTestCase
+
+VOCAB = 1 << 20
+
+
+def _recv(rid, input_ids, max_new_tokens=8):
+    return SimpleNamespace(
+        rid=rid,
+        input_ids=array("q", input_ids),
+        mm_inputs=None,
+        session_params=SimpleNamespace(
+            id="s", rid=None, offset=None, replace=False, drop_previous_output=False
+        ),
+        sampling_params=SamplingParams(max_new_tokens=max_new_tokens),
+        lora_id=None,
+        custom_logit_processor=None,
+        stream=False,
+        return_logprob=False,
+        top_logprobs_num=0,
+        token_ids_logprob=None,
+        require_reasoning=False,
+        return_hidden_states=False,
+        return_routed_experts=False,
+        routed_experts_start_len=0,
+        priority=None,
+        routing_key=None,
+        extra_key=None,
+        http_worker_ipc=None,
+        time_stats=None,
+    )
+
+
+class TestSessionTokenShare(CustomTestCase):
+
+    def setUp(self):
+        self.session = Session(capacity_of_str_len=0, session_id="s", streaming=True)
+
+    def _create(self, rid, input_ids, max_new_tokens=8):
+        return self.session.create_req(
+            _recv(rid, input_ids, max_new_tokens=max_new_tokens),
+            tokenizer=None,
+            vocab_size=VOCAB,
+        )
+
+    def _decode_and_finish(self, req, output, baked=None):
+        """Simulate decode then a successful finish.
+
+        `baked` output tokens are folded into the fill array before the rest
+        arrive (mix_with_running refreshes mid-decode, so the bake is often
+        partial).
+        """
+        if baked is None:
+            baked = len(output)
+        req.output_ids.extend(output[:baked])
+        req._refresh_fill_ids()
+        req.output_ids.extend(output[baked:])
+        self.session.finish_req(req)
+
+    def test_normal_multi_turn_share_and_carry(self):
+        in1, out1 = list(range(100, 110)), [1, 2, 3]
+        r1 = self._create("r1", in1)
+        self.assertEqual(list(r1.origin_input_ids), in1)
+        self._decode_and_finish(r1, out1, baked=2)  # partial bake
+        self.assertEqual(self.session.committed_origin_len, len(in1))
+        self.assertEqual(self.session.committed_fill_len, len(in1) + 2)
+
+        in2, out2 = [7, 8], [4, 5]
+        r2 = self._create("r2", in2)
+        # In-place share: same objects, extended to the new prompt.
+        self.assertIs(r2.origin_input_ids, r1.origin_input_ids)
+        self.assertEqual(list(r2.origin_input_ids), in1 + out1 + in2)
+        # Carry: the fill array handed over and equal to the new origin.
+        self.assertIs(r2.full_untruncated_fill_ids, r1.full_untruncated_fill_ids)
+        self.assertEqual(list(r2.full_untruncated_fill_ids), list(r2.origin_input_ids))
+        self._decode_and_finish(r2, out2)
+
+        r3 = self._create("r3", [9])
+        self.assertEqual(list(r3.origin_input_ids), in1 + out1 + in2 + out2 + [9])
+        self.assertEqual(list(r3.full_untruncated_fill_ids), list(r3.origin_input_ids))
+
+    def test_mid_turn_abort_then_continue(self):
+        in1, out1 = list(range(200, 210)), [1, 2, 3]
+        r1 = self._create("r1", in1)
+        self._decode_and_finish(r1, out1)
+
+        # Turn 2 extends the shared arrays, decodes a bit, then aborts:
+        # finish_req never runs, req_nodes still points at r1.
+        r2 = self._create("r2", [50, 51])
+        self.assertEqual(list(r2.origin_input_ids), in1 + out1 + [50, 51])
+        r2.output_ids.extend([6, 7])
+        r2._refresh_fill_ids()
+        self.session.abort_req()
+        self.assertEqual(self.session.committed_origin_len, len(in1))
+
+        # Turn 3 must see exactly r1's history — no [50, 51], no doubled out1.
+        r3 = self._create("r3", [60])
+        self.assertEqual(list(r3.origin_input_ids), in1 + out1 + [60])
+        self.assertEqual(list(r3.full_untruncated_fill_ids), list(r3.origin_input_ids))
+
+        # Two aborted attempts in a row heal idempotently.
+        self.session.abort_req()
+        r4 = self._create("r4", [70])
+        self.assertEqual(list(r4.origin_input_ids), in1 + out1 + [70])
+        self.assertEqual(list(r4.full_untruncated_fill_ids), list(r4.origin_input_ids))
+
+    def test_first_turn_abort(self):
+        self._create("r1", [1, 2, 3])
+        self.assertTrue(self.session._inflight)
+        self.session.abort_req()
+        self.assertFalse(self.session._inflight)
+        # No finish_req ran: nothing committed, next turn starts from scratch.
+        self.assertIsNone(self.session.committed_origin_len)
+        r2 = self._create("r2", [4, 5])
+        self.assertEqual(list(r2.origin_input_ids), [4, 5])
+        self._decode_and_finish(r2, [9])
+        r3 = self._create("r3", [6])
+        self.assertEqual(list(r3.origin_input_ids), [4, 5, 9, 6])
+
+    def test_max_new_tokens_overshoot_falls_back(self):
+        in1 = list(range(300, 310))
+        r1 = self._create("r1", in1, max_new_tokens=4)
+        # Spec-decode overshoot: 6 tokens decoded and baked into the fill
+        # array, then output trimmed to finished_len (like _trim_overshoot)
+        # before finish.
+        r1.output_ids.extend([1, 2, 3, 4, 5, 6])
+        r1._refresh_fill_ids()
+        del r1.output_ids[4:]
+        self.session.finish_req(r1)
+        self.assertEqual(
+            self.session.committed_fill_len, len(in1) + 6
+        )  # fill kept the overshoot
+
+        # Next turn: out_tail is output[:max_new]; the carried fill has more
+        # baked than out_tail, so the carry is dropped and the fill rebuilds.
+        r2 = self._create("r2", [50])
+        self.assertEqual(list(r2.origin_input_ids), in1 + [1, 2, 3, 4] + [50])
+        self.assertEqual(len(r2.full_untruncated_fill_ids), 0)  # carry skipped
+        r2._refresh_fill_ids()
+        self.assertEqual(list(r2.full_untruncated_fill_ids), list(r2.origin_input_ids))
+
+
+if __name__ == "__main__":
+    unittest.main()