diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 23b41ce06..c98fb0350 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/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 - # updates fill_len, never mutates this array's length. + # Kept in sync by _refresh_fill_ids; admission only updates fill_len, + # 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[ - : self._compute_max_prefix_len(input_len) - ] + # Pass the full array with a raw-token cap (limit) instead of slicing, + # 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) diff --git a/python/sglang/srt/mem_cache/mamba_radix_cache.py b/python/sglang/srt/mem_cache/mamba_radix_cache.py index b4c55a6c3..4a8e53834 100644 --- a/python/sglang/srt/mem_cache/mamba_radix_cache.py +++ b/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 ) diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index 3b103582c..b4365b55c 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/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: + t = self.token_ids + n = self._raw_len() if self.is_bigram: - t = self.token_ids - for i in range(len(t) - 1): + 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 diff --git a/python/sglang/srt/mem_cache/radix_cache_cpp.py b/python/sglang/srt/mem_cache/radix_cache_cpp.py index 344b9a307..3729f397e 100644 --- a/python/sglang/srt/mem_cache/radix_cache_cpp.py +++ b/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), diff --git a/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py b/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py index 733ea74d5..84cce04a3 100644 --- a/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py +++ b/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, diff --git a/python/sglang/srt/session/session_controller.py b/python/sglang/srt/session/session_controller.py index aabb32f94..874de4fdb 100644 --- a/python/sglang/srt/session/session_controller.py +++ b/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 - if ( - tokenizer is not None - and req.input_ids - and req.input_ids[0] == tokenizer.bos_token_id - ): - req.input_ids = req.input_ids[1:] - # Adjust mm_item offsets since they were computed on - # the pre-strip sequence (with BOS at position 0) - 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 - ] - - input_ids = ( - last_req.origin_input_ids - + last_req.output_ids[: last_req.sampling_params.max_new_tokens] + self._strip_bos_token(req, tokenizer) + # In-place sharing is only safe for the plain streaming append: + # streaming sessions allow a single inflight request, last_req has + # finished, and the committed_* lengths recorded by finish_req let + # _share_token_arrays trim away tokens appended by an aborted turn. + # offset / drop_previous_output rewrite history and must copy. + can_share_token_arrays = ( + self.streaming + and self.committed_origin_len is not None + and not session_params.drop_previous_output + and not (session_params.offset and session_params.offset != 0) ) - - 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 + if can_share_token_arrays: + input_ids, input_ids_unpadded, carry_fill = self._share_token_arrays( + last_req, 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).""" diff --git a/python/sglang/srt/session/streaming_session.py b/python/sglang/srt/session/streaming_session.py index c3efbf0b0..d2a7dc33d 100644 --- a/python/sglang/srt/session/streaming_session.py +++ b/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). diff --git a/test/registered/unit/mem_cache/test_session_token_share_unit.py b/test/registered/unit/mem_cache/test_session_token_share_unit.py new file mode 100644 index 000000000..9e9eee10e --- /dev/null +++ b/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()