From 259a2da3e0153cc9b76abcef2b10145e22308e77 Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Mon, 8 Jun 2026 14:52:18 +0800 Subject: [PATCH] Refactor Req.fill_ids into full_untruncated_fill_ids + fill_len with equivalence (#26637) --- python/sglang/bench_one_batch.py | 17 ++++--- python/sglang/srt/disaggregation/decode.py | 24 +++++----- .../decode_schedule_batch_mixin.py | 2 +- python/sglang/srt/disaggregation/prefill.py | 6 +-- python/sglang/srt/dllm/mixin/req.py | 9 ++-- python/sglang/srt/dllm/mixin/scheduler.py | 4 +- .../srt/hardware_backend/mlx/tp_worker.py | 4 +- .../srt/managers/hisparse_coordinator.py | 6 +-- python/sglang/srt/managers/schedule_batch.py | 46 +++++++++++-------- python/sglang/srt/managers/schedule_policy.py | 12 ++--- python/sglang/srt/managers/scheduler.py | 8 ++-- .../logprob_result_processor.py | 2 +- .../sglang/srt/managers/scheduler_pp_mixin.py | 9 ++-- python/sglang/srt/mem_cache/chunk_cache.py | 2 +- .../sglang/srt/mem_cache/mamba_radix_cache.py | 4 +- python/sglang/srt/mem_cache/radix_cache.py | 2 +- .../sglang/srt/mem_cache/radix_cache_cpp.py | 2 +- .../sglang/srt/mem_cache/swa_radix_cache.py | 4 +- .../srt/mem_cache/unified_radix_cache.py | 2 +- .../sglang/srt/session/streaming_session.py | 8 ++-- test/manual/test_forward_split_prefill.py | 5 +- .../test_input_embeds_chunked.py | 12 ++--- .../unit/managers/test_hisparse_unit.py | 4 +- .../unit/managers/test_prefill_adder.py | 12 +++-- .../test_scheduler_chunked_req_gate.py | 3 +- .../mem_cache/test_decode_radix_lock_ref.py | 6 ++- .../test_unified_radix_cache_bench.py | 3 +- .../test_unified_radix_cache_unittest.py | 17 ++++--- test/registered/unit/utils/test_common.py | 4 +- 29 files changed, 136 insertions(+), 103 deletions(-) diff --git a/python/sglang/bench_one_batch.py b/python/sglang/bench_one_batch.py index 1e7b90609..da26065ec 100644 --- a/python/sglang/bench_one_batch.py +++ b/python/sglang/bench_one_batch.py @@ -371,9 +371,10 @@ def prepare_inputs_for_correctness_test(bench_args, tokenizer, custom_prompts): origin_input_ids=array("q", tmp_input_ids), sampling_params=sampling_params, ) - req.fill_ids = req.origin_input_ids + req.full_untruncated_fill_ids = req.origin_input_ids + req.fill_len = len(req.full_untruncated_fill_ids) req.logprob_start_len = -1 - req.set_extend_input_len(len(req.fill_ids) - len(req.prefix_indices)) + req.set_extend_input_len(req.fill_len - len(req.prefix_indices)) reqs.append(req) return input_ids, reqs @@ -384,14 +385,15 @@ def prepare_extend_inputs_for_correctness_test( ): for i in range(len(reqs)): req: Req = reqs[i] - req.fill_ids.extend(input_ids[i][bench_args.cut_len :]) + req.full_untruncated_fill_ids += input_ids[i][bench_args.cut_len :] + req.fill_len = len(req.full_untruncated_fill_ids) if model_runner is not None: # Use req.req_pool_idx instead of i to handle slot 0 padding correctly req.prefix_indices = model_runner.req_to_token_pool.req_to_token[ req.req_pool_idx, : bench_args.cut_len ].to(req.prefix_indices.dtype) req.logprob_start_len = -1 - req.set_extend_input_len(len(req.fill_ids) - len(req.prefix_indices)) + req.set_extend_input_len(req.fill_len - len(req.prefix_indices)) return reqs @@ -416,9 +418,10 @@ def prepare_synthetic_inputs_for_latency_test( origin_input_ids=array("q", input_ids[i]), sampling_params=sampling_params, ) - req.fill_ids = req.origin_input_ids + req.full_untruncated_fill_ids = req.origin_input_ids + req.fill_len = len(req.full_untruncated_fill_ids) req.logprob_start_len = -1 - req.set_extend_input_len(len(req.fill_ids) - len(req.prefix_indices)) + req.set_extend_input_len(req.fill_len - len(req.prefix_indices)) reqs.append(req) return reqs @@ -557,7 +560,7 @@ class _MlxBenchRunner: req_ids = [str(req.rid) for req in reqs] results = [] for rid, req in zip(req_ids, reqs): - token_ids = [int(t) for t in req.fill_ids] + token_ids = [int(t) for t in req.get_fill_ids()] next_token = self.mlx_runner.prefill( req_id=rid, new_token_ids=token_ids, diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 34ba4e0df..b53302cb6 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -1048,9 +1048,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): @property def num_tokens_pre_allocated(self): - return sum( - len(decode_req.req.fill_ids) for decode_req in self.transfer_queue.queue - ) + return sum(decode_req.req.fill_len for decode_req in self.transfer_queue.queue) def _need_space_for_single_req( self, retractable_tokens: Optional[int] = None @@ -1380,10 +1378,11 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): kv_loc, ) - # Truncate fill_ids to kv_committed_len so cache_unfinished_req only + # Truncate fill_len to kv_committed_len so cache_unfinished_req only # inserts committed KV into the radix tree. The last output token - # hasn't had KV committed yet (fill_ids is 1 ahead). - req.fill_ids = (req.origin_input_ids + req.output_ids)[: req.kv_committed_len] + # hasn't had KV committed yet (output_ids is 1 ahead). + req.full_untruncated_fill_ids = req.origin_input_ids + req.output_ids + req.fill_len = req.kv_committed_len # Set prefix_indices so downstream consumers (init_next_round_input, # prepare_for_extend) see the correct prefix length. In the agg path # this is done inside init_next_round_input, but decode-disagg needs @@ -1391,7 +1390,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): req.prefix_indices = ( prefix_indices if prefix_len > 0 else torch.empty((0,), dtype=torch.int64) ) - req.set_extend_input_len(len(req.fill_ids) - total_prefix_len) + req.set_extend_input_len(req.fill_len - total_prefix_len) # Return the transfer destination indices: if self.scheduler.enable_hisparse: @@ -1839,13 +1838,12 @@ class SchedulerDisaggregationDecodeMixin: else: tree_cache = self.tree_cache req.init_next_round_input(tree_cache) - # Truncate fill_ids to kv_committed_len so cache_unfinished_req - # only sees committed KV (fill_ids includes one uncommitted token). + # Truncate fill_len to kv_committed_len so cache_unfinished_req + # only sees committed KV (full array includes one uncommitted + # token because init_next_round_input rebuilt it as full). if req.kv_committed_len is not None: - req.fill_ids = req.fill_ids[: req.kv_committed_len] - req.set_extend_input_len( - len(req.fill_ids) - len(req.prefix_indices) - ) + req.fill_len = req.kv_committed_len + req.set_extend_input_len(req.fill_len - len(req.prefix_indices)) else: waiting_queue.append(req) diff --git a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py index 95e663128..68a06ce8d 100644 --- a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py +++ b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py @@ -29,7 +29,7 @@ class ScheduleBatchDisaggregationDecodeMixin: self.forward_mode = ForwardMode.PREBUILT reqs = self.reqs - input_ids = [r.fill_ids[len(r.prefix_indices) :] for r in reqs] + input_ids = [r.get_fill_ids()[len(r.prefix_indices) :] for r in reqs] extend_num_tokens = sum(len(ids) for ids in input_ids) seq_lens = [] pre_lens = [] diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 172c30adb..21695efee 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -870,7 +870,7 @@ class SchedulerDisaggregationPrefillMixin: elif self.enable_overlap: # Delay KV transfer to process_batch_result_disagg_prefill when overlap is enabled to ensure results are resolved self.chunked_req.tmp_end_idx = min( - len(self.chunked_req.fill_ids), + self.chunked_req.fill_len, len(self.chunked_req.origin_input_ids), ) else: @@ -906,7 +906,7 @@ class SchedulerDisaggregationPrefillMixin: end_idx = ( end_idx if end_idx is not None - else min(len(req.fill_ids), len(req.origin_input_ids)) + else min(req.fill_len, len(req.origin_input_ids)) ) if not last_chunk: @@ -937,7 +937,7 @@ class SchedulerDisaggregationPrefillMixin: # length here avoids emitting an extra state page when the sampled # token crosses a page boundary, which mismatched src/dst lengths in # group_concurrent_contiguous. - seq_len = min(len(req.fill_ids), len(req.origin_input_ids)) + seq_len = min(req.fill_len, len(req.origin_input_ids)) def _mamba_payload(): return [ diff --git a/python/sglang/srt/dllm/mixin/req.py b/python/sglang/srt/dllm/mixin/req.py index 80b624f12..1c9cfc3b8 100644 --- a/python/sglang/srt/dllm/mixin/req.py +++ b/python/sglang/srt/dllm/mixin/req.py @@ -42,11 +42,11 @@ class ReqDllmMixin: prefix_length = len(self.prefix_indices) min_required_length = prefix_length + self.dllm_config.block_size - if len(self.fill_ids) < min_required_length: + if self.fill_len < min_required_length: # still incoming stage return - input_block = self.fill_ids[prefix_length:min_required_length] + input_block = self.get_fill_ids()[prefix_length:min_required_length] is_prefill_phase = self.dllm_config.mask_id not in input_block if is_prefill_phase: @@ -57,14 +57,15 @@ class ReqDllmMixin: def _init_fill_ids_for_dllm(self: Req): self.dllm_block_offset = ( 0 - if not self.fill_ids + if self.fill_len == 0 else self.dllm_block_offset + self.dllm_config.block_size ) - self.fill_ids = ( + self.full_untruncated_fill_ids = ( self.origin_input_ids + self.output_ids + array("q", [self.dllm_config.mask_id] * self.dllm_config.block_size) ) + self.fill_len = len(self.full_untruncated_fill_ids) def _update_block_offset_for_dllm(self): prefix_len = len(self.prefix_indices) diff --git a/python/sglang/srt/dllm/mixin/scheduler.py b/python/sglang/srt/dllm/mixin/scheduler.py index 3fbff7531..21321e723 100644 --- a/python/sglang/srt/dllm/mixin/scheduler.py +++ b/python/sglang/srt/dllm/mixin/scheduler.py @@ -80,7 +80,9 @@ class SchedulerDllmMixin: if new_tokens == 0: continue - req.fill_ids[-new_tokens:] = array("q", next_token_ids) + req.full_untruncated_fill_ids[ + req.fill_len - new_tokens : req.fill_len + ] = array("q", next_token_ids) self.metrics_reporter.num_generated_tokens += new_tokens req.output_ids.extend(next_token_ids) diff --git a/python/sglang/srt/hardware_backend/mlx/tp_worker.py b/python/sglang/srt/hardware_backend/mlx/tp_worker.py index 9488d1bb2..0f81881f0 100644 --- a/python/sglang/srt/hardware_backend/mlx/tp_worker.py +++ b/python/sglang/srt/hardware_backend/mlx/tp_worker.py @@ -185,7 +185,7 @@ class MlxTpModelWorker(TpModelWorker): else: # New prefill prefix_slot_ids = req.prefix_indices.tolist() - full_token_ids = list(req.fill_ids) + full_token_ids = list(req.get_fill_ids()) next_token = self._mlx_runner.prefill( req_id=req.rid, new_token_ids=req_token_ids, @@ -330,7 +330,7 @@ class MlxTpModelWorker(TpModelWorker): else: # New prefill prefix_slot_ids = req.prefix_indices.tolist() - full_token_ids = list(req.fill_ids) + full_token_ids = list(req.get_fill_ids()) pending_prefills.append( self._mlx_runner.prefill_start( req_id=req.rid, diff --git a/python/sglang/srt/managers/hisparse_coordinator.py b/python/sglang/srt/managers/hisparse_coordinator.py index 1720affb1..810b2f368 100644 --- a/python/sglang/srt/managers/hisparse_coordinator.py +++ b/python/sglang/srt/managers/hisparse_coordinator.py @@ -207,7 +207,7 @@ class HiSparseCoordinator: req.hisparse_staging = True full_kv_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, : len(req.fill_ids) + req.req_pool_idx, : req.fill_len ].to(dtype=torch.int64, copy=True) device_indices = ( self.mem_pool_device.translate_loc_from_full_to_hisparse_device( @@ -298,7 +298,7 @@ class HiSparseCoordinator: def alloc_device_buffer(self, req: Req) -> None: if self.is_dsv4_hisparse: - allocated_len = len(req.fill_ids) + allocated_len = req.fill_len alloc_size = self.padded_buffer_size else: allocated_len = req.kv_allocated_len @@ -707,7 +707,7 @@ class HiSparseCoordinator: # Wait for any in-flight staging DMA to complete before freeing self.write_staging_stream.synchronize() - prefill_len = len(req.fill_ids) + prefill_len = req.fill_len allocated_locs = self.req_to_token_pool.req_to_token[ req.req_pool_idx, :prefill_len ] diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 841427d80..3c6ad40aa 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -696,8 +696,11 @@ class Req(ReqDllmMixin): ) # Before image padding # Each decode stage's output ids self.output_ids = array("q") - # fill_ids = origin_input_ids + output_ids. Updated if chunked. - self.fill_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. + self.full_untruncated_fill_ids = array("q") + self.fill_len: int = 0 self.session = session self.input_embeds = input_embeds @@ -959,8 +962,8 @@ class Req(ReqDllmMixin): # the start index of the sent kv cache # We want to send it chunk by chunk for chunked prefill. # After every chunk forward, we do the following: - # kv_send(req.input_ids[req.start_send_idx:len(req.fill_ids)]) - # start_send_idx = len(req.fill_ids) + # kv_send(req.input_ids[req.start_send_idx:req.fill_len]) + # start_send_idx = req.fill_len self.start_send_idx: int = 0 # For overlap schedule, we delay the kv transfer until `process_batch_result_disagg_prefill` rather than `process_prefill_chunk` in non-overlap @@ -1062,6 +1065,9 @@ class Req(ReqDllmMixin): # Whether request reached finished condition return self.finished_reason is not None + def get_fill_ids(self) -> array: + return self.full_untruncated_fill_ids[: self.fill_len] + def init_next_round_input( self, tree_cache: Optional[BasePrefixCache] = None, @@ -1071,9 +1077,10 @@ class Req(ReqDllmMixin): self._init_fill_ids_for_dllm() self.determine_dllm_phase() else: - self.fill_ids = self.origin_input_ids + self.output_ids + self.full_untruncated_fill_ids = self.origin_input_ids + self.output_ids + self.fill_len = len(self.full_untruncated_fill_ids) - input_len = len(self.fill_ids) + input_len = self.fill_len # Streaming sessions reuse committed KV from the session slot, so # custom logprob_start_len is not supported — override to -1. @@ -1091,7 +1098,9 @@ class Req(ReqDllmMixin): ) self.logprob_start_len = -1 - token_ids_to_match = self.fill_ids[: self._compute_max_prefix_len(input_len)] + token_ids_to_match = self.get_fill_ids()[ + : 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. @@ -1154,7 +1163,7 @@ class Req(ReqDllmMixin): ) ) - self.set_extend_input_len(len(self.fill_ids) - len(self.prefix_indices)) + self.set_extend_input_len(self.fill_len - len(self.prefix_indices)) def _compute_max_prefix_len(self, input_len: int) -> int: # NOTE: the matched length is at most 1 less than the input length to enable logprob computation @@ -1420,7 +1429,7 @@ class Req(ReqDllmMixin): # - extend_input_len: Number of tokens that need to be processed in this extend batch self.extend_input_len = extend_input_len if self.logprob_start_len == -1: - logprob_start_len = len(self.fill_ids) + logprob_start_len = self.fill_len else: # logprob_start_len should be at least the length of the prefix indices logprob_start_len = max(self.logprob_start_len, len(self.prefix_indices)) @@ -1544,7 +1553,7 @@ def _compute_chunked_req_next_prompt_token( ) -> Optional[int]: if chunked_req is None: return None - fill_len = len(chunked_req.fill_ids) + fill_len = chunked_req.fill_len if fill_len >= len(chunked_req.origin_input_ids): return None return int(chunked_req.origin_input_ids[fill_len]) @@ -1892,10 +1901,10 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): # Init tensors reqs = self.reqs - input_ids = [r.fill_ids[len(r.prefix_indices) :] for r in reqs] + input_ids = [r.get_fill_ids()[len(r.prefix_indices) :] for r in reqs] extend_num_tokens = sum(len(ids) for ids in input_ids) - seq_lens = [len(r.fill_ids) for r in reqs] - orig_seq_lens = [max(len(r.fill_ids), len(r.origin_input_ids)) for r in reqs] + seq_lens = [r.fill_len for r in reqs] + orig_seq_lens = [max(r.fill_len, len(r.origin_input_ids)) for r in reqs] prefix_lens = [len(r.prefix_indices) for r in reqs] extend_lens = [r.extend_input_len for r in reqs] @@ -1948,7 +1957,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): # If input_embeds are available, store them if req.input_embeds is not None: # Slice to match extend_input_len — PrefillAdder truncates - # fill_ids/extend_input_len on chunk overflow but not input_embeds. + # fill_len/extend_input_len on chunk overflow but not input_embeds. input_embeds.extend( req.input_embeds[pre_len : pre_len + req.extend_input_len] ) @@ -2022,16 +2031,16 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): # to compute input logprobs. E.g., (chunk size 2) # # input_logprobs = [1, 2, 3, 4] - # fill_ids = [1, 2] + # get_fill_ids() = [1, 2] # extend_input_logprob_token_id = [2, 3] # # Note that it can also overflow. In this case, we pad it with 0. # input_logprobs = [1, 2, 3, 4] - # fill_ids = [3, 4] + # get_fill_ids() = [3, 4] # extend_input_logprob_token_id = [4, 0] global_start_idx, global_end_idx = ( len(req.prefix_indices), - len(req.fill_ids), + req.fill_len, ) if req.logprob_start_len == -1: logprob_start_len = len(req.origin_input_ids) @@ -2261,7 +2270,8 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): running_bs = running_batch.batch_size() for req in running_batch.reqs: - req.fill_ids = req.origin_input_ids + req.output_ids + req.full_untruncated_fill_ids = req.origin_input_ids + req.output_ids + req.fill_len = len(req.full_untruncated_fill_ids) req.set_extend_input_len(1) # Decode tokens of the running portion live in future_map.output_tokens_buf. diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 588261d7c..9fae0bab4 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -660,7 +660,7 @@ class PrefillAdder: ) req.extend_input_len = trunc_len - req.fill_ids = req.fill_ids[: prefix_len + trunc_len] + req.fill_len = prefix_len + trunc_len self.can_run_list.append(req) @@ -681,7 +681,7 @@ class PrefillAdder: # Truncate input length to available tokens and update request metadata truncated = req.extend_input_len > _rem_tokens req.extend_input_len = min(req.extend_input_len, _rem_tokens) - req.fill_ids = req.fill_ids[: len(req.prefix_indices) + req.extend_input_len] + req.fill_len = len(req.prefix_indices) + req.extend_input_len self.can_run_list.append(req) # Update budget: reserve max_new_tokens only if not truncated @@ -721,7 +721,7 @@ class PrefillAdder: truncated = req.extend_input_len > _rem_tokens 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] + req.fill_len = len(req.prefix_indices) + req.extend_input_len self.can_run_list.append(req) self._update_prefill_budget( 0, @@ -843,7 +843,7 @@ class PrefillAdder: trunc_len = self.rem_chunk_tokens req.set_extend_input_len(trunc_len) - req.fill_ids = req.fill_ids[:trunc_len] + req.fill_len = trunc_len self.can_run_list.append(req) self.new_chunked_req = req self._update_prefill_budget(0, trunc_len, 0, req.retracted_stain) @@ -932,7 +932,7 @@ class PrefillAdder: ) ) req.prefix_indices = torch.cat([req.prefix_indices, new_indices]) - req.set_extend_input_len(len(req.fill_ids) - len(req.prefix_indices)) + req.set_extend_input_len(req.fill_len - len(req.prefix_indices)) prefix_len = len(req.prefix_indices) req.cache_protected_len = prefix_len @@ -999,7 +999,7 @@ class PrefillAdder: # Chunked prefill req.set_extend_input_len(trunc_len) - req.fill_ids = req.fill_ids[: len(req.prefix_indices) + trunc_len] + req.fill_len = len(req.prefix_indices) + trunc_len self.can_run_list.append(req) self.new_chunked_req = req diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 96e21ec75..bbd32aaed 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -941,9 +941,9 @@ class Scheduler( # req_pool_idx from prepare_for_extend). Used to gate the # stash_chunked_request call at the top of get_next_batch_to_run: # if add_chunked_req early-returned under hybrid-SWA pressure, - # the req_pool_idx was already freed and fill_ids was reset by - # init_next_round_input, so running stash would double-free and - # corrupt prefix_indices. + # the req_pool_idx was already freed and the full_untruncated_fill_ids + # was rebuilt by init_next_round_input, so running stash would + # double-free and corrupt prefix_indices. self._chunked_req_scheduled_last_iter = False self.is_mixed_chunk = ( self.chunked_prefill_size is not None @@ -2159,7 +2159,7 @@ class Scheduler( if last_host_node.backuped or last_host_node is self.tree_cache.root_node: last_hash = last_host_node.get_last_hash_value() matched_len = len(req.prefix_indices) + req.host_hit_length - new_input_tokens = req.fill_ids[matched_len:] + new_input_tokens = req.get_fill_ids()[matched_len:] prefix_keys = ( last_host_node.get_prefix_hash_values(last_host_node.parent) diff --git a/python/sglang/srt/managers/scheduler_components/logprob_result_processor.py b/python/sglang/srt/managers/scheduler_components/logprob_result_processor.py index f58429f0a..d97d9ae80 100644 --- a/python/sglang/srt/managers/scheduler_components/logprob_result_processor.py +++ b/python/sglang/srt/managers/scheduler_components/logprob_result_processor.py @@ -184,7 +184,7 @@ class SchedulerLogprobResultProcessor: i: The request index in a batch. req: The request. Input logprobs inside req are modified as a consequence of the API - fill_ids: The prefill ids processed. + logprob_pt: Pointer into the prefill ids processed. output: Logit processor output that's used to compute input logprobs last_prefill_chunk: True if it is the last prefill (when chunked). Some of input logprob operation should only happen at the last diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py index c6aa22f16..8babb55cf 100644 --- a/python/sglang/srt/managers/scheduler_pp_mixin.py +++ b/python/sglang/srt/managers/scheduler_pp_mixin.py @@ -606,9 +606,10 @@ class SchedulerPPMixin: origin_input_ids=input_ids, sampling_params=sampling_params, ) - req.fill_ids = req.origin_input_ids + req.full_untruncated_fill_ids = req.origin_input_ids + req.fill_len = len(req.full_untruncated_fill_ids) req.logprob_start_len = -1 - req.set_extend_input_len(len(req.fill_ids) - len(req.prefix_indices)) + req.set_extend_input_len(req.fill_len - len(req.prefix_indices)) # Prepare batch batch = ScheduleBatch.init_new( @@ -621,7 +622,7 @@ class SchedulerPPMixin: self.spec_algorithm, ) - current_seq_len = len(req.fill_ids) + current_seq_len = req.fill_len if is_dp_attention_enabled(): # For profiling, we only have one request on PP0 @@ -685,7 +686,7 @@ class SchedulerPPMixin: # Release KV cache if req.req_pool_idx is not None: kv_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, : len(req.fill_ids) + req.req_pool_idx, : req.fill_len ] self.token_to_kv_pool_allocator.free(kv_indices) self.req_to_token_pool.free(req) diff --git a/python/sglang/srt/mem_cache/chunk_cache.py b/python/sglang/srt/mem_cache/chunk_cache.py index bd1351a38..cfccdc815 100644 --- a/python/sglang/srt/mem_cache/chunk_cache.py +++ b/python/sglang/srt/mem_cache/chunk_cache.py @@ -86,7 +86,7 @@ class ChunkCache(BasePrefixCache): def cache_unfinished_req(self, req: Req, chunked=False): kv_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, : len(req.fill_ids) + req.req_pool_idx, : req.fill_len ] # `req.prefix_indices` will be used in `PrefillAdder::add_chunked_req` later req.prefix_indices = kv_indices.to(dtype=torch.int64, copy=True) diff --git a/python/sglang/srt/mem_cache/mamba_radix_cache.py b/python/sglang/srt/mem_cache/mamba_radix_cache.py index 59903afe5..b0d80df19 100644 --- a/python/sglang/srt/mem_cache/mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/mamba_radix_cache.py @@ -611,14 +611,14 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache): def _skip_cache_unfinished_req(req: Req) -> None: kv_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, : len(req.fill_ids) + req.req_pool_idx, : req.fill_len ] # `req.prefix_indices` will be used in `PrefillAdder::add_chunked_req` later req.prefix_indices = kv_indices.to(dtype=torch.int64, copy=True) return - token_ids = req.fill_ids + token_ids = req.get_fill_ids() cache_len = ( req.mamba_last_track_seqlen if self.enable_mamba_extra_buffer diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index e7d5bade0..bd6adb6e3 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -466,7 +466,7 @@ class RadixCache(KVCacheEventMixin, BasePrefixCache): if self.disable: return - token_ids = req.fill_ids + token_ids = req.get_fill_ids() kv_indices = self.req_to_token_pool.req_to_token[ req.req_pool_idx, : len(token_ids) ] diff --git a/python/sglang/srt/mem_cache/radix_cache_cpp.py b/python/sglang/srt/mem_cache/radix_cache_cpp.py index 97bf02835..344b9a307 100644 --- a/python/sglang/srt/mem_cache/radix_cache_cpp.py +++ b/python/sglang/srt/mem_cache/radix_cache_cpp.py @@ -210,7 +210,7 @@ class RadixCacheCpp(BasePrefixCache): def cache_unfinished_req(self, req: Req, chunked=False): """Cache request when it is unfinished.""" assert req.req_pool_idx is not None - token_ids = req.fill_ids + token_ids = req.get_fill_ids() prefill_len = len(token_ids) # prefill only (maybe chunked) kv_indices = self.req_to_token_pool.req_to_token[ req.req_pool_idx, :prefill_len diff --git a/python/sglang/srt/mem_cache/swa_radix_cache.py b/python/sglang/srt/mem_cache/swa_radix_cache.py index cd428f86c..d4f6afa4c 100644 --- a/python/sglang/srt/mem_cache/swa_radix_cache.py +++ b/python/sglang/srt/mem_cache/swa_radix_cache.py @@ -488,14 +488,14 @@ class SWARadixCache(KVCacheEventMixin, BasePrefixCache): """Cache request when it is unfinished.""" if self.disable: kv_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, : len(req.fill_ids) + req.req_pool_idx, : req.fill_len ] # `req.prefix_indices` will be used in `PrefillAdder::add_chunked_req` later req.prefix_indices = kv_indices return - token_ids = req.fill_ids + token_ids = req.get_fill_ids() kv_indices = self.req_to_token_pool.req_to_token[ req.req_pool_idx, : len(token_ids) ] diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index f63254906..455a656fc 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -722,7 +722,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): if self.session.try_cache_unfinished_req(req, chunked=chunked, **kwargs): return - token_ids = req.fill_ids + token_ids = req.get_fill_ids() if self.disable: kv_indices = self.req_to_token_pool.req_to_token[ diff --git a/python/sglang/srt/session/streaming_session.py b/python/sglang/srt/session/streaming_session.py index e1b0b18ef..c3efbf0b0 100644 --- a/python/sglang/srt/session/streaming_session.py +++ b/python/sglang/srt/session/streaming_session.py @@ -244,9 +244,9 @@ class StreamingSession(BasePrefixCache): req = params.req slot.restore_to_req(req) - # token_ids = 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. + # 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)) # Streaming sessions are append-only (session_controller rollback @@ -353,7 +353,7 @@ class StreamingSession(BasePrefixCache): return False if chunked: kv_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, : len(req.fill_ids) + req.req_pool_idx, : req.fill_len ] req.prefix_indices = kv_indices.to(dtype=torch.int64, copy=True) return True diff --git a/test/manual/test_forward_split_prefill.py b/test/manual/test_forward_split_prefill.py index 80d6cc661..f48ef6b0a 100644 --- a/test/manual/test_forward_split_prefill.py +++ b/test/manual/test_forward_split_prefill.py @@ -95,9 +95,10 @@ class TestForwardSplitPrefill(CustomTestCase): origin_input_ids=array("q", input_ids[i]), sampling_params=sampling_params, ) - req.fill_ids = req.origin_input_ids + req.full_untruncated_fill_ids = req.origin_input_ids + req.fill_len = len(req.full_untruncated_fill_ids) req.logprob_start_len = -1 - req.set_extend_input_len(len(req.fill_ids) - len(req.prefix_indices)) + req.set_extend_input_len(req.fill_len - len(req.prefix_indices)) reqs.append(req) # Create dummy tree_cache for tests (no prefix caching, just allocation) diff --git a/test/registered/input_embedding/test_input_embeds_chunked.py b/test/registered/input_embedding/test_input_embeds_chunked.py index 878df0136..c26623068 100644 --- a/test/registered/input_embedding/test_input_embeds_chunked.py +++ b/test/registered/input_embedding/test_input_embeds_chunked.py @@ -3,14 +3,14 @@ Covers two bugs with the same crash signature (RuntimeError: shape mismatch in set_kv_buffer) but opposite polarity: -- Chunked prefill truncation (#20376): PrefillAdder truncates fill_ids and +- Chunked prefill truncation (#20376): PrefillAdder shrinks fill_len and extend_input_len on chunk overflow but not input_embeds, so the full array flows through while out_cache_loc is sized for the truncated length. Polarity: cache_k > loc. -- Retraction with output_ids (#14110): after retraction, fill_ids includes - accumulated output_ids but input_embeds only covers origin_input_ids. - Polarity: cache_k < loc. +- Retraction with output_ids (#14110): after retraction, get_fill_ids() + includes accumulated output_ids but input_embeds only covers + origin_input_ids. Polarity: cache_k < loc. """ import unittest @@ -162,8 +162,8 @@ class TestInputEmbedsChunkedAndRetract(CustomTestCase): SGLANG_TEST_RETRACT forces retraction every few scheduler iterations. Combined with ignore_eos and a reasonable max_new_tokens, at least one request is retracted mid-decode with non-empty output_ids, then - re-prefilled. Pre-#14110 this crashes (cache_k < loc) because fill_ids - includes output_ids but input_embeds does not. + re-prefilled. Pre-#14110 this crashes (cache_k < loc) because the + filled token sequence includes output_ids but input_embeds does not. """ text = "The quick brown fox jumps over the lazy dog. " * 4 embeds = _embeds_for(text) diff --git a/test/registered/unit/managers/test_hisparse_unit.py b/test/registered/unit/managers/test_hisparse_unit.py index 22996c54e..a31c39a72 100644 --- a/test/registered/unit/managers/test_hisparse_unit.py +++ b/test/registered/unit/managers/test_hisparse_unit.py @@ -9,6 +9,7 @@ Tests cover: import os import unittest +from array import array from types import SimpleNamespace import torch @@ -210,7 +211,8 @@ class TestHiSparseUnit(unittest.TestCase): self.req_to_token_pool.write((req.req_pool_idx, slice(0, len(kv_loc))), kv_loc) req.kv_allocated_len = fill_len req.kv_committed_len = fill_len - req.fill_ids = list(range(fill_len)) + req.full_untruncated_fill_ids = array("q", range(fill_len)) + req.fill_len = fill_len return kv_loc # ================================================================== diff --git a/test/registered/unit/managers/test_prefill_adder.py b/test/registered/unit/managers/test_prefill_adder.py index f9410f0c9..ded85b1fb 100644 --- a/test/registered/unit/managers/test_prefill_adder.py +++ b/test/registered/unit/managers/test_prefill_adder.py @@ -386,7 +386,8 @@ class TestPrefillAdder(CustomTestCase): req1.extend_input_len = 56 req1.host_hit_length = 0 req1.prefix_indices = [] - req1.fill_ids = list(range(56)) + req1.full_untruncated_fill_ids = list(range(56)) + req1.fill_len = 56 req1.last_node = MagicMock() req1.sampling_params.ignore_eos = False @@ -420,7 +421,8 @@ class TestPrefillAdder(CustomTestCase): req2.extend_input_len = 56 req2.host_hit_length = 0 req2.prefix_indices = [] - req2.fill_ids = list(range(56)) + req2.full_untruncated_fill_ids = list(range(56)) + req2.fill_len = 56 req2.last_node = MagicMock() req2.sampling_params.ignore_eos = False @@ -437,7 +439,8 @@ class TestPrefillAdder(CustomTestCase): req3.extend_input_len = 3 req3.host_hit_length = 0 req3.prefix_indices = [] - req3.fill_ids = list(range(3)) + req3.full_untruncated_fill_ids = list(range(3)) + req3.fill_len = 3 req3.last_node = MagicMock() req3.sampling_params.ignore_eos = False @@ -473,7 +476,8 @@ class TestPrefillAdder(CustomTestCase): req = self.create_mock_req("chunked", priority=0, max_new_tokens=128) req.extend_input_len = extend_input_len req.prefix_indices = [] - req.fill_ids = list(range(extend_input_len)) + req.full_untruncated_fill_ids = list(range(extend_input_len)) + req.fill_len = extend_input_len req.set_extend_input_len = MagicMock() return adder, req diff --git a/test/registered/unit/managers/test_scheduler_chunked_req_gate.py b/test/registered/unit/managers/test_scheduler_chunked_req_gate.py index 0263170bc..7e09b14e3 100644 --- a/test/registered/unit/managers/test_scheduler_chunked_req_gate.py +++ b/test/registered/unit/managers/test_scheduler_chunked_req_gate.py @@ -30,7 +30,8 @@ def _make_req( req.rid = "test-req" req.origin_input_ids = array("q", fill_ids) req.output_ids = array("q") - req.fill_ids = array("q", fill_ids) + req.full_untruncated_fill_ids = array("q", fill_ids) + req.fill_len = len(req.full_untruncated_fill_ids) req.prefix_indices = prefix_indices req.req_pool_idx = req_pool_idx req.extend_input_len = extend_input_len diff --git a/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py b/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py index 0d78edd0f..82943b8a1 100644 --- a/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py +++ b/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py @@ -67,7 +67,8 @@ class MockReq: """Minimal mock Req with fields needed by cache_unfinished/finished_req.""" def __init__(self, fill_ids, req_pool_idx=0, cache_protected_len=0, last_node=None): - self.fill_ids = array("q", fill_ids) + self.full_untruncated_fill_ids = array("q", fill_ids) + self.fill_len = len(self.full_untruncated_fill_ids) self.origin_input_ids = array( "q", fill_ids[:-1] if len(fill_ids) > 1 else fill_ids ) @@ -82,6 +83,9 @@ class MockReq: self.kv_allocated_len = len(fill_ids) self.kv_committed_freed = False + def get_fill_ids(self): + return self.full_untruncated_fill_ids[: self.fill_len] + def pop_committed_kv_cache(self): self.kv_committed_freed = True return self.kv_committed_len diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py b/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py index 04f6d28ec..01b1a7639 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py @@ -637,7 +637,8 @@ def bench_cache_finished( req = env.make_req() req.origin_input_ids = array("q", seq) req.output_ids = array("q") - req.fill_ids = array("q", seq) + req.full_untruncated_fill_ids = array("q", seq) + req.fill_len = len(req.full_untruncated_fill_ids) req.last_node = node req.cache_protected_len = matched_len req.kv_committed_len = len(seq) diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index f9847b873..2ecda3ed6 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -798,7 +798,8 @@ class UnifiedRadixCacheSuite: req.cache_protected_len = 0 req.swa_uuid_for_lock = None req.extra_key = None - req.fill_ids = array("q", input_ids + output_ids) + req.full_untruncated_fill_ids = array("q", input_ids + output_ids) + req.fill_len = len(req.full_untruncated_fill_ids) if self.cfg.has_mamba: req.mamba_last_track_seqlen = kv_len @@ -821,8 +822,9 @@ class UnifiedRadixCacheSuite: output_ids = self._make_seq(2000, 7) req.origin_input_ids = array("q", prompt_ids) req.output_ids = array("q", output_ids) - req.fill_ids = array("q", prompt_ids + output_ids) - kv_len = len(req.fill_ids) + req.full_untruncated_fill_ids = array("q", prompt_ids + output_ids) + req.fill_len = len(req.full_untruncated_fill_ids) + kv_len = req.fill_len kv_indices = self._alloc(allocator, kv_len) req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices) req.kv_committed_len = kv_len @@ -875,7 +877,8 @@ class UnifiedRadixCacheSuite: req.cache_protected_len = 0 req.swa_uuid_for_lock = None req.extra_key = None - req.fill_ids = array("q", tokens) + req.full_untruncated_fill_ids = array("q", tokens) + req.fill_len = len(req.full_untruncated_fill_ids) avail_before = allocator.available_size() tree.cache_finished_req(req, is_insert=False) @@ -892,7 +895,8 @@ class UnifiedRadixCacheSuite: tokens = self._make_seq(1, 3) req.origin_input_ids = array("q", tokens) req.output_ids = array("q") - req.fill_ids = array("q", tokens) + req.full_untruncated_fill_ids = array("q", tokens) + req.fill_len = len(req.full_untruncated_fill_ids) kv_len = len(tokens) kv_indices = self._alloc(allocator, kv_len) req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices) @@ -1026,7 +1030,8 @@ class UnifiedRadixCacheSuite: req.cache_protected_len = 0 req.swa_uuid_for_lock = None req.extra_key = None - req.fill_ids = array("q", input_ids) + req.full_untruncated_fill_ids = array("q", input_ids) + req.fill_len = len(req.full_untruncated_fill_ids) if self.cfg.has_mamba: req.mamba_last_track_seqlen = kv_len diff --git a/test/registered/unit/utils/test_common.py b/test/registered/unit/utils/test_common.py index a15f6b003..b38e1e7f7 100644 --- a/test/registered/unit/utils/test_common.py +++ b/test/registered/unit/utils/test_common.py @@ -14,8 +14,8 @@ register_cuda_ci(est_time=5, stage="base-b", runner_config="1-gpu-small") class TestFlattenArraysToInt64Tensor(CustomTestCase): """`flatten_arrays_to_int64_tensor` is invoked by `prepare_for_extend` to build the per-batch input_ids tensor (pinned, async H2D) from a - list of array.array('q') per-req fill_ids slices. Tests the full - matrix of (device, pin) the production code paths through. + list of array.array('q') per-req get_fill_ids() slices. Tests the + full matrix of (device, pin) the production code paths through. """ DEVICES = ("cpu", "cuda")