minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Inline extend_range accessors and remove the extend_input_len/fill_len properties (#27611)

Browse Source
This commit is contained in:
fzyzcjy
2026-06-25 08:56:44 +08:00
committed by GitHub
parent 9305d10099
commit 7002a37ea1
19 changed files with 87 additions and 97 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/benchmark/one_batch.py
+5 -6
View File
@@ -381,9 +381,8 @@ def prepare_inputs_for_correctness_test(bench_args, tokenizer, custom_prompts):
sampling_params=sampling_params, sampling_params=sampling_params,
) )
req.full_untruncated_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.logprob_start_len = -1
req.set_extend_input_len(req.fill_len - len(req.prefix_indices)) req.set_extend_range(len(req.prefix_indices), len(req.origin_input_ids))
reqs.append(req) reqs.append(req)
return input_ids, reqs return input_ids, reqs
@@ -395,14 +394,15 @@ def prepare_extend_inputs_for_correctness_test(
for i in range(len(reqs)): for i in range(len(reqs)):
req: Req = reqs[i] req: Req = reqs[i]
req.full_untruncated_fill_ids.extend(input_ids[i][bench_args.cut_len :]) req.full_untruncated_fill_ids.extend(input_ids[i][bench_args.cut_len :])
req.fill_len = len(req.full_untruncated_fill_ids)
if model_runner is not None: if model_runner is not None:
# Use req.req_pool_idx instead of i to handle slot 0 padding correctly # 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.prefix_indices = model_runner.req_to_token_pool.req_to_token[
req.req_pool_idx, : bench_args.cut_len req.req_pool_idx, : bench_args.cut_len
].to(req.prefix_indices.dtype) ].to(req.prefix_indices.dtype)
req.logprob_start_len = -1 req.logprob_start_len = -1
req.set_extend_input_len(req.fill_len - len(req.prefix_indices)) req.set_extend_range(
len(req.prefix_indices), len(req.full_untruncated_fill_ids)
)
return reqs return reqs
@@ -428,9 +428,8 @@ def prepare_synthetic_inputs_for_latency_test(
sampling_params=sampling_params, sampling_params=sampling_params,
) )
req.full_untruncated_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.logprob_start_len = -1
req.set_extend_input_len(req.fill_len - len(req.prefix_indices)) req.set_extend_range(len(req.prefix_indices), len(req.origin_input_ids))
reqs.append(req) reqs.append(req)
return reqs return reqs
python/sglang/srt/disaggregation/decode.py
+3 -1
View File
@@ -1085,7 +1085,9 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
@property @property
def num_tokens_pre_allocated(self): def num_tokens_pre_allocated(self):
return sum(decode_req.req.fill_len for decode_req in self.transfer_queue.queue) return sum(
decode_req.req.extend_range.end for decode_req in self.transfer_queue.queue
)
def _need_space_for_single_req( def _need_space_for_single_req(
self, retractable_tokens: Optional[int] = None self, retractable_tokens: Optional[int] = None
python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py
+9 -9
View File
@@ -37,7 +37,7 @@ class ScheduleBatchDisaggregationDecodeMixin:
req_pool_indices = [] req_pool_indices = []
# Pre-calculate total size # Pre-calculate total size
total_size = sum(req.extend_input_len for req in reqs) total_size = sum(req.extend_range.length for req in reqs)
out_cache_loc = torch.empty(total_size, dtype=torch.int64, device=self.device) out_cache_loc = torch.empty(total_size, dtype=torch.int64, device=self.device)
# Fill the tensor in one pass # Fill the tensor in one pass
@@ -47,20 +47,20 @@ class ScheduleBatchDisaggregationDecodeMixin:
pre_len = len(req.prefix_indices) pre_len = len(req.prefix_indices)
chunk = self.req_to_token_pool.req_to_token[req.req_pool_idx][ chunk = self.req_to_token_pool.req_to_token[req.req_pool_idx][
pre_len : pre_len + req.extend_input_len pre_len : pre_len + req.extend_range.length
] ]
assert ( assert (
offset + req.extend_input_len <= total_size offset + req.extend_range.length <= total_size
), f"Exceeds total size: offset={offset}, req.extend_input_len={req.extend_input_len}, total_size={total_size}" ), f"Exceeds total size: offset={offset}, req.extend_range.length={req.extend_range.length}, total_size={total_size}"
out_cache_loc[offset : offset + req.extend_input_len] = chunk out_cache_loc[offset : offset + req.extend_range.length] = chunk
offset += req.extend_input_len offset += req.extend_range.length
seq_len = len(req.origin_input_ids) + max(0, len(req.output_ids) - 1) seq_len = len(req.origin_input_ids) + max(0, len(req.output_ids) - 1)
seq_lens.append(seq_len) seq_lens.append(seq_len)
if len(req.output_ids) == 0: if len(req.output_ids) == 0:
assert ( assert (
seq_len - pre_len == req.extend_input_len seq_len - pre_len == req.extend_range.length
), f"seq_len={seq_len}, pre_len={pre_len}, req.extend_input_len={req.extend_input_len}" ), f"seq_len={seq_len}, pre_len={pre_len}, req.extend_range.length={req.extend_range.length}"
if not req.retracted_stain: if not req.retracted_stain:
# Clamp to avoid double-counting: already_computed is seeded from # Clamp to avoid double-counting: already_computed is seeded from
@@ -99,7 +99,7 @@ class ScheduleBatchDisaggregationDecodeMixin:
self.extend_num_tokens = extend_num_tokens self.extend_num_tokens = extend_num_tokens
self.prefix_lens = [len(r.prefix_indices) for r in reqs] self.prefix_lens = [len(r.prefix_indices) for r in reqs]
self.extend_lens = [r.extend_input_len for r in reqs] self.extend_lens = [r.extend_range.length for r in reqs]
self.extend_logprob_start_lens = [r.extend_logprob_start_len for r in reqs] self.extend_logprob_start_lens = [r.extend_logprob_start_len for r in reqs]
self.extend_input_logprob_token_ids = extend_input_logprob_token_ids self.extend_input_logprob_token_ids = extend_input_logprob_token_ids
self.multimodal_inputs = [r.multimodal_inputs for r in reqs] self.multimodal_inputs = [r.multimodal_inputs for r in reqs]
python/sglang/srt/disaggregation/prefill.py
+3 -3
View File
@@ -934,7 +934,7 @@ class SchedulerDisaggregationPrefillMixin:
elif self.enable_overlap: elif self.enable_overlap:
# Delay KV transfer to process_batch_result_disagg_prefill when overlap is enabled to ensure results are resolved # 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( self.chunked_req.tmp_end_idx = min(
self.chunked_req.fill_len, self.chunked_req.extend_range.end,
len(self.chunked_req.origin_input_ids), len(self.chunked_req.origin_input_ids),
) )
else: else:
@@ -970,7 +970,7 @@ class SchedulerDisaggregationPrefillMixin:
end_idx = ( end_idx = (
end_idx end_idx
if end_idx is not None if end_idx is not None
else min(req.fill_len, len(req.origin_input_ids)) else min(req.extend_range.end, len(req.origin_input_ids))
) )
if not last_chunk: if not last_chunk:
@@ -1001,7 +1001,7 @@ class SchedulerDisaggregationPrefillMixin:
# length here avoids emitting an extra state page when the sampled # length here avoids emitting an extra state page when the sampled
# token crosses a page boundary, which mismatched src/dst lengths in # token crosses a page boundary, which mismatched src/dst lengths in
# group_concurrent_contiguous. # group_concurrent_contiguous.
seq_len = min(req.fill_len, len(req.origin_input_ids)) seq_len = min(req.extend_range.end, len(req.origin_input_ids))
def _mamba_payload(): def _mamba_payload():
return [ return [
python/sglang/srt/dllm/mixin/scheduler.py
+1 -1
View File
@@ -81,7 +81,7 @@ class SchedulerDllmMixin:
continue continue
req.full_untruncated_fill_ids[ req.full_untruncated_fill_ids[
req.fill_len - new_tokens : req.fill_len req.extend_range.end - new_tokens : req.extend_range.end
] = array("q", next_token_ids) ] = array("q", next_token_ids)
self.metrics_reporter.num_generated_tokens += new_tokens self.metrics_reporter.num_generated_tokens += new_tokens
python/sglang/srt/managers/hisparse_coordinator.py
+3 -3
View File
@@ -217,7 +217,7 @@ class HiSparseCoordinator:
req.hisparse_staging = True req.hisparse_staging = True
full_kv_indices = self.req_to_token_pool.req_to_token[ full_kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, : req.fill_len req.req_pool_idx, : req.extend_range.end
].to(dtype=torch.int64, copy=True) ].to(dtype=torch.int64, copy=True)
device_indices = ( device_indices = (
self.mem_pool_device.translate_loc_from_full_to_hisparse_device( self.mem_pool_device.translate_loc_from_full_to_hisparse_device(
@@ -308,7 +308,7 @@ class HiSparseCoordinator:
def alloc_device_buffer(self, req: Req) -> None: def alloc_device_buffer(self, req: Req) -> None:
if self.is_dsv4_hisparse: if self.is_dsv4_hisparse:
allocated_len = req.fill_len allocated_len = req.extend_range.end
alloc_size = self.padded_buffer_size alloc_size = self.padded_buffer_size
else: else:
allocated_len = req.kv_allocated_len allocated_len = req.kv_allocated_len
@@ -729,7 +729,7 @@ class HiSparseCoordinator:
# Wait for any in-flight staging DMA to complete before freeing # Wait for any in-flight staging DMA to complete before freeing
self.write_staging_stream.synchronize() self.write_staging_stream.synchronize()
prefill_len = req.fill_len prefill_len = req.extend_range.end
allocated_locs = self.req_to_token_pool.req_to_token[ allocated_locs = self.req_to_token_pool.req_to_token[
req.req_pool_idx, :prefill_len req.req_pool_idx, :prefill_len
] ]
python/sglang/srt/managers/schedule_batch.py
+22 -30
View File
@@ -993,8 +993,8 @@ class Req(ReqDllmMixin):
# the start index of the sent kv cache # the start index of the sent kv cache
# We want to send it chunk by chunk for chunked prefill. # We want to send it chunk by chunk for chunked prefill.
# After every chunk forward, we do the following: # After every chunk forward, we do the following:
# kv_send(req.input_ids[req.start_send_idx:req.fill_len]) # kv_send(req.input_ids[req.start_send_idx:req.extend_range.end])
# start_send_idx = req.fill_len # start_send_idx = req.extend_range.end
self.start_send_idx: int = 0 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 # For overlap schedule, we delay the kv transfer until `process_batch_result_disagg_prefill` rather than `process_prefill_chunk` in non-overlap
@@ -1096,20 +1096,12 @@ class Req(ReqDllmMixin):
# Whether request reached finished condition # Whether request reached finished condition
return self.finished_reason is not None return self.finished_reason is not None
@property
def fill_len(self) -> int:
return self.extend_range.end
@property
def extend_input_len(self) -> int:
return self.extend_range.length
def set_extend_range(self, start: int, end: int) -> None: def set_extend_range(self, start: int, end: int) -> None:
self.extend_range = Range(start, end) self.extend_range = Range(start, end)
self._recompute_extend_logprob_start_len() self._recompute_extend_logprob_start_len()
def get_fill_ids(self) -> array: def get_fill_ids(self) -> array:
return self.full_untruncated_fill_ids[: self.fill_len] return self.full_untruncated_fill_ids[: self.extend_range.end]
def _refresh_fill_ids(self) -> None: def _refresh_fill_ids(self) -> None:
"""Keep full_untruncated_fill_ids == origin_input_ids + output_ids by """Keep full_untruncated_fill_ids == origin_input_ids + output_ids by
@@ -1546,7 +1538,7 @@ class Req(ReqDllmMixin):
logprob_start_len = max(self.logprob_start_len, len(self.prefix_indices)) logprob_start_len = max(self.logprob_start_len, len(self.prefix_indices))
self.extend_logprob_start_len = min( self.extend_logprob_start_len = min(
logprob_start_len - len(self.prefix_indices), logprob_start_len - len(self.prefix_indices),
self.extend_input_len, self.extend_range.length,
) )
def set_finish_with_abort(self, error_msg: str): def set_finish_with_abort(self, error_msg: str):
@@ -1667,7 +1659,7 @@ def _compute_chunked_req_next_prompt_token(
multimodal placeholder (hash) tokens that lie outside the model vocab.""" multimodal placeholder (hash) tokens that lie outside the model vocab."""
if chunked_req is None: if chunked_req is None:
return None return None
fill_len = chunked_req.fill_len fill_len = chunked_req.extend_range.end
origin_ids = chunked_req.origin_input_ids origin_ids = chunked_req.origin_input_ids
if fill_len >= len(origin_ids): if fill_len >= len(origin_ids):
return None return None
@@ -1932,17 +1924,17 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
input_ids[i] = input_ids[i][encoder_len:] input_ids[i] = input_ids[i][encoder_len:]
encoder_out_cache_loc.append(self.out_cache_loc[pt : pt + encoder_len]) encoder_out_cache_loc.append(self.out_cache_loc[pt : pt + encoder_len])
decoder_out_cache_loc.append( decoder_out_cache_loc.append(
self.out_cache_loc[pt + encoder_len : pt + req.extend_input_len] self.out_cache_loc[pt + encoder_len : pt + req.extend_range.length]
) )
self.extend_lens[i] -= encoder_len self.extend_lens[i] -= encoder_len
self.extend_num_tokens -= encoder_len self.extend_num_tokens -= encoder_len
else: else:
decoder_out_cache_loc.append( decoder_out_cache_loc.append(
self.out_cache_loc[pt : pt + req.extend_input_len] self.out_cache_loc[pt : pt + req.extend_range.length]
) )
self.prefix_lens[i] -= encoder_len self.prefix_lens[i] -= encoder_len
pt += req.extend_input_len pt += req.extend_range.length
# Reassign: ED stripping rebuilds prefill_input_ids_cpu (CPU pinned); # Reassign: ED stripping rebuilds prefill_input_ids_cpu (CPU pinned);
# resolve_forward_inputs will H2D this on forward stream. self.input_ids # resolve_forward_inputs will H2D this on forward stream. self.input_ids
@@ -1977,7 +1969,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
for i, req in enumerate(self.reqs): for i, req in enumerate(self.reqs):
encoder_len = self.encoder_lens_cpu[i] encoder_len = self.encoder_lens_cpu[i]
old_start_len = self.extend_logprob_start_lens[i] old_start_len = self.extend_logprob_start_lens[i]
old_contribution = req.extend_input_len - old_start_len old_contribution = req.extend_range.length - old_start_len
if len(req.prefix_indices) < encoder_len: if len(req.prefix_indices) < encoder_len:
tokens_to_strip = max(0, encoder_len - old_start_len) tokens_to_strip = max(0, encoder_len - old_start_len)
@@ -2028,10 +2020,10 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
reqs = self.reqs reqs = self.reqs
input_ids = [r.get_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) extend_num_tokens = sum(len(ids) for ids in input_ids)
seq_lens = [r.fill_len for r in reqs] seq_lens = [r.extend_range.end for r in reqs]
orig_seq_lens = [max(r.fill_len, len(r.origin_input_ids)) for r in reqs] orig_seq_lens = [max(r.extend_range.end, len(r.origin_input_ids)) for r in reqs]
prefix_lens = [len(r.prefix_indices) for r in reqs] prefix_lens = [len(r.prefix_indices) for r in reqs]
extend_lens = [r.extend_input_len for r in reqs] extend_lens = [r.extend_range.length for r in reqs]
_pin = is_pin_memory_available(self.device) _pin = is_pin_memory_available(self.device)
# Stay on pinned CPU; H2D is deferred to forward stream via # Stay on pinned CPU; H2D is deferred to forward stream via
@@ -2071,7 +2063,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
mamba_track_seqlens_cpu = [] mamba_track_seqlens_cpu = []
for i, (req, seq_len, pre_len) in enumerate(zip(reqs, seq_lens, prefix_lens)): for i, (req, seq_len, pre_len) in enumerate(zip(reqs, seq_lens, prefix_lens)):
assert seq_len - pre_len == req.extend_input_len assert seq_len - pre_len == req.extend_range.length
req.extend_batch_idx += 1 req.extend_batch_idx += 1
@@ -2084,7 +2076,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
# Slice to match extend_input_len — PrefillAdder truncates # Slice to match extend_input_len — PrefillAdder truncates
# fill_len/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( input_embeds.extend(
req.input_embeds[pre_len : pre_len + req.extend_input_len] req.input_embeds[pre_len : pre_len + req.extend_range.length]
) )
if req.positional_embed_overrides is not None: if req.positional_embed_overrides is not None:
@@ -2096,7 +2088,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
req.positional_embed_overrides.positions req.positional_embed_overrides.positions
): ):
extend_pos = pos - pre_len extend_pos = pos - pre_len
if extend_pos < 0 or extend_pos >= req.extend_input_len: if extend_pos < 0 or extend_pos >= req.extend_range.length:
continue # Outside current extend chunk, skip continue # Outside current extend chunk, skip
embeds_to_add.append((embed_idx, input_id_pointer + extend_pos)) embeds_to_add.append((embed_idx, input_id_pointer + extend_pos))
if embeds_to_add: if embeds_to_add:
@@ -2165,7 +2157,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
# extend_input_logprob_token_id = [4, 0] # extend_input_logprob_token_id = [4, 0]
global_start_idx, global_end_idx = ( global_start_idx, global_end_idx = (
len(req.prefix_indices), len(req.prefix_indices),
req.fill_len, req.extend_range.end,
) )
if req.logprob_start_len == -1: if req.logprob_start_len == -1:
logprob_start_len = len(req.origin_input_ids) logprob_start_len = len(req.origin_input_ids)
@@ -2180,12 +2172,12 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
] ]
extend_input_logprob_token_ids.extend(logprob_token_ids) extend_input_logprob_token_ids.extend(logprob_token_ids)
# We will need req.extend_input_len - req.extend_logprob_start_len number of # We will need req.extend_range.length - req.extend_logprob_start_len number of
# tokens, and logprob_token_ids is for input logprob, so pad the rest of them by 0. # tokens, and logprob_token_ids is for input logprob, so pad the rest of them by 0.
extend_input_logprob_token_ids.extend( extend_input_logprob_token_ids.extend(
[0] [0]
* ( * (
req.extend_input_len req.extend_range.length
- req.extend_logprob_start_len - req.extend_logprob_start_len
- len(logprob_token_ids) - len(logprob_token_ids)
) )
@@ -2295,7 +2287,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
# to force the math calculation to retrieve the correct mamba state from h. # to force the math calculation to retrieve the correct mamba state from h.
return i + 1 return i + 1
mask = req.extend_input_len >= mamba_cache_chunk_size mask = req.extend_range.length >= mamba_cache_chunk_size
track_index = req.mamba_ping_pong_track_buffer[req.mamba_next_track_idx].item() track_index = req.mamba_ping_pong_track_buffer[req.mamba_next_track_idx].item()
mamba_track_seqlen = -1 mamba_track_seqlen = -1
if mask: if mask:
@@ -2306,13 +2298,13 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
# otherwise retrieved from h (i.e. unaligned). # otherwise retrieved from h (i.e. unaligned).
# We need to pass the non-aligned seqlen to the calculation. Even though # We need to pass the non-aligned seqlen to the calculation. Even though
# we pass in mamba_track_seqlen, the actual tracked seqlen is mamba_last_track_seqlen. # we pass in mamba_track_seqlen, the actual tracked seqlen is mamba_last_track_seqlen.
mamba_track_seqlen = len(req.prefix_indices) + req.extend_input_len mamba_track_seqlen = len(req.prefix_indices) + req.extend_range.length
# mamba_track_seqlen_aligned/mamba_last_track_seqlen is actual tracked seqlen. Used to pass to # mamba_track_seqlen_aligned/mamba_last_track_seqlen is actual tracked seqlen. Used to pass to
# mamba radix cache to track which seqlen this mamba state should store at. # mamba radix cache to track which seqlen this mamba state should store at.
mamba_track_seqlen_aligned = ( mamba_track_seqlen_aligned = (
len(req.prefix_indices) len(req.prefix_indices)
+ (req.extend_input_len // mamba_cache_chunk_size) + (req.extend_range.length // mamba_cache_chunk_size)
* mamba_cache_chunk_size * mamba_cache_chunk_size
) )
@@ -2322,7 +2314,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
# by _force_track_h() # by _force_track_h()
mamba_track_fla_chunk_aligned = ( mamba_track_fla_chunk_aligned = (
len(req.prefix_indices) len(req.prefix_indices)
+ (req.extend_input_len // mamba_cache_chunk_size) + (req.extend_range.length // mamba_cache_chunk_size)
* mamba_cache_chunk_size * mamba_cache_chunk_size
) )
if mamba_track_fla_chunk_aligned != mamba_track_seqlen_aligned: if mamba_track_fla_chunk_aligned != mamba_track_seqlen_aligned:
python/sglang/srt/managers/schedule_policy.py
+3 -3
View File
@@ -693,7 +693,7 @@ class PrefillAdder:
else 0 else 0
) )
self._update_prefill_budget( self._update_prefill_budget(
0, req.extend_input_len, max_new_tokens, req.retracted_stain 0, req.extend_range.length, max_new_tokens, req.retracted_stain
) )
# Return based on remaining token availability # Return based on remaining token availability
@@ -730,7 +730,7 @@ class PrefillAdder:
self.can_run_list.append(req) self.can_run_list.append(req)
self._update_prefill_budget( self._update_prefill_budget(
0, 0,
req.extend_input_len, req.extend_range.length,
( (
min(req.sampling_params.max_new_tokens, CLIP_MAX_NEW_TOKENS) min(req.sampling_params.max_new_tokens, CLIP_MAX_NEW_TOKENS)
if not truncated if not truncated
@@ -842,7 +842,7 @@ class PrefillAdder:
self.can_run_list.append(req) self.can_run_list.append(req)
self._update_prefill_budget( self._update_prefill_budget(
0, 0,
req.extend_input_len, req.extend_range.length,
min(req.sampling_params.max_new_tokens, CLIP_MAX_NEW_TOKENS), min(req.sampling_params.max_new_tokens, CLIP_MAX_NEW_TOKENS),
req.retracted_stain, req.retracted_stain,
) )
python/sglang/srt/managers/scheduler.py
+4 -4
View File
@@ -1293,7 +1293,7 @@ class Scheduler(
request_lengths = [] request_lengths = []
for req in batch.reqs: for req in batch.reqs:
start = len(req.prefix_indices) start = len(req.prefix_indices)
end = start + req.extend_input_len end = start + req.extend_range.length
fill_ids = req.origin_input_ids + req.output_ids fill_ids = req.origin_input_ids + req.output_ids
if start == 0: if start == 0:
tokens = fill_ids[start:end] tokens = fill_ids[start:end]
@@ -2623,7 +2623,7 @@ class Scheduler(
# beyond what is already cached. A parked chunk (add_chunked_req # beyond what is already cached. A parked chunk (add_chunked_req
# hybrid-SWA early-return) leaves fill_len == len(prefix_indices), # hybrid-SWA early-return) leaves fill_len == len(prefix_indices),
# so there is nothing new to cache and stashing would be a no-op. # so there is nothing new to cache and stashing would be a no-op.
if self.chunked_req.fill_len > len(self.chunked_req.prefix_indices): if self.chunked_req.extend_range.end > len(self.chunked_req.prefix_indices):
self.stash_chunked_request(self.chunked_req) self.stash_chunked_request(self.chunked_req)
# HiSparse has its own prefill-to-decode transition; skip last_batch merge. # HiSparse has its own prefill-to-decode transition; skip last_batch merge.
@@ -2974,7 +2974,7 @@ class Scheduler(
self.enable_priority_scheduling, self.enable_priority_scheduling,
num_pending_tokens=self.load_inquirer._get_num_pending_tokens( num_pending_tokens=self.load_inquirer._get_num_pending_tokens(
chunk_deduct=( chunk_deduct=(
self.chunked_req.extend_input_len self.chunked_req.extend_range.length
if self.chunked_req is not None if self.chunked_req is not None
else 0 else 0
), ),
@@ -3337,7 +3337,7 @@ class Scheduler(
# we can use the correct values in output processing. # we can use the correct values in output processing.
if batch.return_logprob: if batch.return_logprob:
batch_result.extend_input_len_per_req = [ batch_result.extend_input_len_per_req = [
req.extend_input_len if req.extend_range is not None else 0 req.extend_range.length if req.extend_range is not None else 0
for req in batch.reqs for req in batch.reqs
] ]
batch_result.extend_logprob_start_len_per_req = [ batch_result.extend_logprob_start_len_per_req = [
python/sglang/srt/managers/scheduler_pp_mixin.py
+2 -2
View File
@@ -624,7 +624,7 @@ class SchedulerPPMixin:
self.spec_algorithm, self.spec_algorithm,
) )
current_seq_len = req.fill_len current_seq_len = req.extend_range.end
if is_dp_attention_enabled(): if is_dp_attention_enabled():
# For profiling, we only have one request on PP0 # For profiling, we only have one request on PP0
@@ -695,7 +695,7 @@ class SchedulerPPMixin:
# Release KV cache # Release KV cache
if req.req_pool_idx is not None: if req.req_pool_idx is not None:
kv_indices = self.req_to_token_pool.req_to_token[ kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, : req.fill_len req.req_pool_idx, : req.extend_range.end
] ]
self.token_to_kv_pool_allocator.free(kv_indices) self.token_to_kv_pool_allocator.free(kv_indices)
self.req_to_token_pool.free(req) self.req_to_token_pool.free(req)
python/sglang/srt/mem_cache/chunk_cache.py
+1 -1
View File
@@ -86,7 +86,7 @@ class ChunkCache(BasePrefixCache):
def cache_unfinished_req(self, req: Req, chunked=False): def cache_unfinished_req(self, req: Req, chunked=False):
kv_indices = self.req_to_token_pool.req_to_token[ kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, : req.fill_len req.req_pool_idx, : req.extend_range.end
] ]
# `req.prefix_indices` will be used in `PrefillAdder::add_chunked_req` later # `req.prefix_indices` will be used in `PrefillAdder::add_chunked_req` later
req.prefix_indices = kv_indices.to(dtype=torch.int64, copy=True) req.prefix_indices = kv_indices.to(dtype=torch.int64, copy=True)
python/sglang/srt/mem_cache/mamba_radix_cache.py
+1 -1
View File
@@ -627,7 +627,7 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache):
def _skip_cache_unfinished_req(req: Req) -> None: def _skip_cache_unfinished_req(req: Req) -> None:
kv_indices = self.req_to_token_pool.req_to_token[ kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, : req.fill_len req.req_pool_idx, : req.extend_range.end
] ]
# `req.prefix_indices` will be used in `PrefillAdder::add_chunked_req` later # `req.prefix_indices` will be used in `PrefillAdder::add_chunked_req` later
python/sglang/srt/mem_cache/swa_radix_cache.py
+1 -1
View File
@@ -488,7 +488,7 @@ class SWARadixCache(KVCacheEventMixin, BasePrefixCache):
"""Cache request when it is unfinished.""" """Cache request when it is unfinished."""
if self.disable: if self.disable:
kv_indices = self.req_to_token_pool.req_to_token[ kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, : req.fill_len req.req_pool_idx, : req.extend_range.end
] ]
# `req.prefix_indices` will be used in `PrefillAdder::add_chunked_req` later # `req.prefix_indices` will be used in `PrefillAdder::add_chunked_req` later
python/sglang/srt/session/streaming_session.py
+1 -1
View File
@@ -359,7 +359,7 @@ class StreamingSession(BasePrefixCache):
return False return False
if chunked: if chunked:
kv_indices = self.req_to_token_pool.req_to_token[ kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, : req.fill_len req.req_pool_idx, : req.extend_range.end
] ]
req.prefix_indices = kv_indices.to(dtype=torch.int64, copy=True) req.prefix_indices = kv_indices.to(dtype=torch.int64, copy=True)
return True return True
test/manual/chunked_prefill/test_scripted_special_case.py
+12 -9
View File
@@ -406,9 +406,9 @@ class TestSpecialCaseBasic(ScriptedTestCase):
r.is_chunking r.is_chunking
and chunked is not None and chunked is not None
and chunked.rid == r.rid and chunked.rid == r.rid
and chunked.extend_input_len > 0 and chunked.extend_range.length > 0
): ):
deduct = chunked.extend_input_len deduct = chunked.extend_range.length
base = s.load_inquirer._get_num_pending_tokens() base = s.load_inquirer._get_num_pending_tokens()
deducted = s.load_inquirer._get_num_pending_tokens(chunk_deduct=deduct) deducted = s.load_inquirer._get_num_pending_tokens(chunk_deduct=deduct)
assert deducted == base - deduct, ( assert deducted == base - deduct, (
@@ -473,15 +473,18 @@ class TestSpecialCaseBasic(ScriptedTestCase):
r.is_chunking r.is_chunking
and r.chunks_done >= 1 and r.chunks_done >= 1
and req is not None and req is not None
and req.extend_input_len is not None and req.extend_range is not None
): ):
saw_mid_chunk = True saw_mid_chunk = True
assert req.fill_len == len(req.prefix_indices) + req.extend_input_len, ( assert (
req.extend_range.end
== len(req.prefix_indices) + req.extend_range.length
), (
f"init_next_round_input must rebuild fill_ids to the committed " f"init_next_round_input must rebuild fill_ids to the committed "
f"prefix plus the in-flight chunk; " f"prefix plus the in-flight chunk; "
f"fill_ids_len={req.fill_len}, " f"fill_ids_len={req.extend_range.end}, "
f"prefix_indices_len={len(req.prefix_indices)}, " f"prefix_indices_len={len(req.prefix_indices)}, "
f"extend_input_len={req.extend_input_len}, " f"extend_input_len={req.extend_range.length}, "
f"chunks_done={r.chunks_done}" f"chunks_done={r.chunks_done}"
) )
if r.finished: if r.finished:
@@ -775,11 +778,11 @@ class TestSpecialCaseDeterministicFlashInfer(ScriptedTestCase):
page_size = 16 page_size = 16
saw_chunking = False saw_chunking = False
for _ in range(DEFAULT_MAX_STEPS): for _ in range(DEFAULT_MAX_STEPS):
if r.is_chunking and r.req.extend_input_len is not None: if r.is_chunking and r.req.extend_range is not None:
saw_chunking = True saw_chunking = True
assert r.req.extend_input_len % page_size == 0, ( assert r.req.extend_range.length % page_size == 0, (
f"deterministic chunk boundary must be page-aligned; " f"deterministic chunk boundary must be page-aligned; "
f"got extend_input_len={r.req.extend_input_len}, page_size={page_size}" f"got extend_input_len={r.req.extend_range.length}, page_size={page_size}"
) )
if r.finished: if r.finished:
break break
test/registered/unit/managers/test_hisparse_unit.py
+2 -13
View File
@@ -18,17 +18,6 @@ from sglang.srt.utils import is_cuda, is_hip, is_npu, is_xpu
from sglang.srt.utils.common import Range from sglang.srt.utils.common import Range
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
class _FakeReq(SimpleNamespace):
@property
def fill_len(self) -> int:
return self.extend_range.end
@property
def extend_input_len(self) -> int:
return self.extend_range.length
register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-small") register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-small")
register_amd_ci(est_time=10, suite="stage-b-test-1-gpu-small-amd") register_amd_ci(est_time=10, suite="stage-b-test-1-gpu-small-amd")
@@ -54,7 +43,7 @@ def _make_req(rid="test-req-0", origin_input_ids=None, output_ids=None):
origin_input_ids = list(range(64)) origin_input_ids = list(range(64))
if output_ids is None: if output_ids is None:
output_ids = [] output_ids = []
req = _FakeReq( req = SimpleNamespace(
rid=rid, rid=rid,
origin_input_ids=origin_input_ids, origin_input_ids=origin_input_ids,
output_ids=output_ids, output_ids=output_ids,
@@ -781,7 +770,7 @@ class TestHiSparseUnit(unittest.TestCase):
) )
self.assertEqual(req.kv_allocated_len, fill_len) self.assertEqual(req.kv_allocated_len, fill_len)
self.assertEqual(req.kv_committed_len, fill_len) self.assertEqual(req.kv_committed_len, fill_len)
self.assertEqual(req.extend_input_len, fill_len) self.assertEqual(req.extend_range.length, fill_len)
rounded_len = (fill_len + self.page_size - 1) // self.page_size * self.page_size rounded_len = (fill_len + self.page_size - 1) // self.page_size * self.page_size
self.assertEqual( self.assertEqual(
test/registered/unit/managers/test_prefill_adder.py
+10 -1
View File
@@ -9,6 +9,7 @@ from sglang.srt.mem_cache.base_prefix_cache import (
IncLockRefResult, IncLockRefResult,
) )
from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
from sglang.srt.utils.common import Range
from sglang.test.ci.ci_register import ( from sglang.test.ci.ci_register import (
register_amd_ci, register_amd_ci,
register_cpu_ci, register_cpu_ci,
@@ -471,7 +472,15 @@ class TestPrefillAdder(CustomTestCase):
req = self.create_mock_req("chunked", priority=0, max_new_tokens=128) req = self.create_mock_req("chunked", priority=0, max_new_tokens=128)
req.prefix_indices = [] req.prefix_indices = []
req.full_untruncated_fill_ids = list(range(extend_input_len)) req.full_untruncated_fill_ids = list(range(extend_input_len))
req.set_extend_range = MagicMock() # set_extend_range is the only writer of extend_range; the production
# path reads req.extend_range.length right after calling it, so the mock
# must actually set the attribute (a spec=Req mock has the method but
# not the instance attribute).
req.set_extend_range = MagicMock(
side_effect=lambda start, end: setattr(
req, "extend_range", Range(start, end)
)
)
return adder, req return adder, req
def test_add_chunked_req_hybrid_swa_reserves_page_for_alloc_extend(self): def test_add_chunked_req_hybrid_swa_reserves_page_for_alloc_extend(self):
test/registered/unit/mem_cache/test_decode_radix_lock_ref.py
+1 -5
View File
@@ -84,12 +84,8 @@ class MockReq:
self.kv_allocated_len = len(fill_ids) self.kv_allocated_len = len(fill_ids)
self.kv_committed_freed = False self.kv_committed_freed = False
@property
def fill_len(self):
return self.extend_range.end
def get_fill_ids(self): def get_fill_ids(self):
return self.full_untruncated_fill_ids[: self.fill_len] return self.full_untruncated_fill_ids[: self.extend_range.end]
def pop_committed_kv_cache(self): def pop_committed_kv_cache(self):
self.kv_committed_freed = True self.kv_committed_freed = True
test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py
+3 -3
View File
@@ -927,7 +927,7 @@ class UnifiedRadixCacheSuite:
req.set_extend_range( req.set_extend_range(
len(req.prefix_indices), len(req.full_untruncated_fill_ids) len(req.prefix_indices), len(req.full_untruncated_fill_ids)
) )
kv_len = req.fill_len kv_len = req.extend_range.end
kv_indices = self._alloc(allocator, kv_len) kv_indices = self._alloc(allocator, kv_len)
req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices) req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices)
req.kv_committed_len = kv_len req.kv_committed_len = kv_len
@@ -1815,7 +1815,7 @@ class UnifiedRadixCacheSuite:
req.origin_input_ids = tokens req.origin_input_ids = tokens
req.output_ids = [] req.output_ids = []
req.full_untruncated_fill_ids = array("q", tokens) req.full_untruncated_fill_ids = array("q", tokens)
req.fill_len = len(req.full_untruncated_fill_ids) req.set_extend_range(0, len(req.full_untruncated_fill_ids))
kv_indices = self._alloc(allocator, pre_len) kv_indices = self._alloc(allocator, pre_len)
req_to_token_pool.write((req.req_pool_idx, slice(0, pre_len)), kv_indices) req_to_token_pool.write((req.req_pool_idx, slice(0, pre_len)), kv_indices)
req.kv_committed_len = pre_len req.kv_committed_len = pre_len
@@ -1867,7 +1867,7 @@ class UnifiedRadixCacheSuite:
req.origin_input_ids = tokens req.origin_input_ids = tokens
req.output_ids = [] req.output_ids = []
req.full_untruncated_fill_ids = array("q", tokens) req.full_untruncated_fill_ids = array("q", tokens)
req.fill_len = len(req.full_untruncated_fill_ids) req.set_extend_range(0, len(req.full_untruncated_fill_ids))
kv_indices = self._alloc(allocator, pre_len) kv_indices = self._alloc(allocator, pre_len)
req_to_token_pool.write((req.req_pool_idx, slice(0, pre_len)), kv_indices) req_to_token_pool.write((req.req_pool_idx, slice(0, pre_len)), kv_indices)
req.kv_committed_len = pre_len req.kv_committed_len = pre_len