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

Avoid mutating ScheduleBatch fields in place (#30672)

Browse Source
This commit is contained in:
fzyzcjy
2026-07-15 14:27:08 +08:00
committed by GitHub
parent e77d95c3d5
commit 52a88fb212
4 changed files with 317 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/schedule_batch.py
+46 -40
View File
@@ -2000,21 +2000,23 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
self, input_ids: List[array[int]], seq_lens: List[int]
):
_pin = is_pin_memory_available(self.device)
self.encoder_lens_cpu = []
self.encoder_cached = []
encoder_lens_cpu = []
encoder_cached = []
for req in self.reqs:
im = req.multimodal_inputs
if im is None or im.num_image_tokens is None:
# No image input
self.encoder_lens_cpu.append(0)
self.encoder_cached.append(True)
encoder_lens_cpu.append(0)
encoder_cached.append(True)
else:
self.encoder_lens_cpu.append(im.num_image_tokens)
self.encoder_cached.append(
encoder_lens_cpu.append(im.num_image_tokens)
encoder_cached.append(
self.forward_mode.is_decode()
or len(req.prefix_indices) >= im.num_image_tokens
)
self.encoder_lens_cpu = encoder_lens_cpu
self.encoder_cached = encoder_cached
self.encoder_lens = torch.tensor(
self.encoder_lens_cpu, dtype=torch.int64, pin_memory=_pin
@@ -2024,6 +2026,8 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
pt = 0
decoder_out_cache_loc = []
encoder_out_cache_loc = []
extend_lens = self.extend_lens[:]
prefix_lens = self.prefix_lens[:]
for i, req in enumerate(self.reqs):
encoder_len = self.encoder_lens_cpu[i]
seq_lens[i] -= encoder_len
@@ -2036,15 +2040,17 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
decoder_out_cache_loc.append(
self.out_cache_loc[pt + encoder_len : pt + req.extend_range.length]
)
self.extend_lens[i] -= encoder_len
self.extend_num_tokens -= encoder_len
extend_lens[i] -= encoder_len
self.extend_num_tokens = self.extend_num_tokens - encoder_len
else:
decoder_out_cache_loc.append(
self.out_cache_loc[pt : pt + req.extend_range.length]
)
self.prefix_lens[i] -= encoder_len
prefix_lens[i] -= encoder_len
pt += req.extend_range.length
self.extend_lens = extend_lens
self.prefix_lens = prefix_lens
# Reassign: ED stripping rebuilds prefill_input_ids_cpu (CPU pinned);
# resolve_forward_inputs will H2D this on forward stream. self.input_ids
@@ -2076,9 +2082,10 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
if self.extend_input_logprob_token_ids is not None:
new_token_ids_parts = []
offset = 0
extend_logprob_start_lens = self.extend_logprob_start_lens[:]
for i, req in enumerate(self.reqs):
encoder_len = self.encoder_lens_cpu[i]
old_start_len = self.extend_logprob_start_lens[i]
old_start_len = extend_logprob_start_lens[i]
old_contribution = req.extend_range.length - old_start_len
if len(req.prefix_indices) < encoder_len:
@@ -2088,9 +2095,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
offset + tokens_to_strip : offset + old_contribution
]
)
self.extend_logprob_start_lens[i] = max(
0, old_start_len - encoder_len
)
extend_logprob_start_lens[i] = max(0, old_start_len - encoder_len)
else:
new_token_ids_parts.append(
self.extend_input_logprob_token_ids[
@@ -2099,6 +2104,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
)
offset += old_contribution
self.extend_logprob_start_lens = extend_logprob_start_lens
if new_token_ids_parts:
self.extend_input_logprob_token_ids = torch.cat(new_token_ids_parts)
@@ -2521,16 +2527,16 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
delta = 0 if self.enable_overlap else -1
# NOTE: prefix_indices is what has been cached, but we don't cache each decode step
self.prefix_lens.extend(
[
len(r.origin_input_ids) + len(r.output_ids) + delta
for r in running_batch.reqs
]
)
self.extend_lens.extend([1] * running_bs)
self.extend_num_tokens += running_bs
self.prefix_lens = self.prefix_lens + [
len(r.origin_input_ids) + len(r.output_ids) + delta
for r in running_batch.reqs
]
self.extend_lens = self.extend_lens + [1] * running_bs
self.extend_num_tokens = self.extend_num_tokens + running_bs
# TODO (lianmin): Revisit this. It should be seq_len - 1
self.extend_logprob_start_lens.extend([0] * running_bs)
self.extend_logprob_start_lens = (
self.extend_logprob_start_lens + [0] * running_bs
)
self.is_prefill_only = False
def new_tokens_required_next_decode(
@@ -2923,7 +2929,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
# Encoder-decoder infos
if self.model_config.is_encoder_decoder:
self.encoder_lens = torch.cat([self.encoder_lens, other.encoder_lens])
self.encoder_lens_cpu.extend(other.encoder_lens_cpu)
self.encoder_lens_cpu = self.encoder_lens_cpu + other.encoder_lens_cpu
self.req_pool_indices = torch.cat(
[self.req_pool_indices, other.req_pool_indices]
)
@@ -2952,21 +2958,23 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
self.mamba_track_mask = None
self.mamba_track_seqlens = None
if self.return_logprob and other.return_logprob:
self.top_logprobs_nums.extend(other.top_logprobs_nums)
self.token_ids_logprobs.extend(other.token_ids_logprobs)
self.top_logprobs_nums = self.top_logprobs_nums + other.top_logprobs_nums
self.token_ids_logprobs = self.token_ids_logprobs + other.token_ids_logprobs
elif self.return_logprob:
self.top_logprobs_nums.extend([0] * len(other.reqs))
self.token_ids_logprobs.extend([None] * len(other.reqs))
self.top_logprobs_nums = self.top_logprobs_nums + [0] * len(other.reqs)
self.token_ids_logprobs = self.token_ids_logprobs + [None] * len(other.reqs)
elif other.return_logprob:
self.top_logprobs_nums = [0] * len(self.reqs) + other.top_logprobs_nums
self.token_ids_logprobs = [None] * len(self.reqs) + other.token_ids_logprobs
self.reqs.extend(other.reqs)
self.reqs = self.reqs + other.reqs
if self.multimodal_inputs is not None:
self.multimodal_inputs.extend(other.multimodal_inputs)
self.multimodal_inputs = self.multimodal_inputs + other.multimodal_inputs
self.return_logprob |= other.return_logprob
self.has_grammar |= other.has_grammar
self.return_hidden_states |= other.return_hidden_states
self.return_logprob = self.return_logprob or other.return_logprob
self.has_grammar = self.has_grammar or other.has_grammar
self.return_hidden_states = (
self.return_hidden_states or other.return_hidden_states
)
self.is_prefill_only = self.is_prefill_only and other.is_prefill_only
if self.spec_info:
@@ -2974,16 +2982,14 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
def copy(self):
# Only contain fields that will be used by process_batch_result.
# Shallow-copy the reqs list so that in-place mutations (filter_batch,
# merge_batch) on the original don't corrupt this snapshot.
# Shallow-copy the reqs list as a defensive snapshot. filter_batch and
# merge_batch historically mutated the list in place; they now rebind
# new lists, but the slice stays so this snapshot never aliases the
# original.
return ScheduleBatch(
reqs=self.reqs[:],
# Per-request extend/prefix lens, snapshotted (sliced like reqs) so the
# deferred prefill-stats report reads them after the original batch has
# moved on. prepare_for_extend sets these; mix_with_running mutates them
# in place. None for decode batches (no extend), which the reader skips.
extend_lens=self.extend_lens[:] if self.extend_lens is not None else None,
prefix_lens=self.prefix_lens[:] if self.prefix_lens is not None else None,
extend_lens=self.extend_lens,
prefix_lens=self.prefix_lens,
req_to_token_pool=self.req_to_token_pool,
req_pool_indices=self.req_pool_indices,
model_config=self.model_config,
python/sglang/srt/speculative/eagle_worker_v2.py
+5 -2
View File
@@ -777,13 +777,16 @@ class EagleDraftWorker(EagleDraftWorkerBase):
if not batch.forward_mode.is_idle():
# Chunked-prefill-aware tail tokens (see PR #26329).
tail_tokens = _eagle_prefill_tail_tokens(batch, next_token_ids)
new_input_ids = torch.empty_like(batch.input_ids)
pt = 0
for i, extend_len in enumerate(batch.extend_lens):
input_ids = batch.input_ids[pt : pt + extend_len]
batch.input_ids[pt : pt + extend_len] = torch.cat(
(input_ids[1:], tail_tokens[i].reshape(1))
new_input_ids[pt : pt + extend_len].copy_(
torch.cat((input_ids[1:], tail_tokens[i].reshape(1)))
)
pt += extend_len
assert pt == batch.input_ids.numel()
batch.input_ids = new_input_ids
# Draft-extend spec_info for the extend forward; carries only
# hidden_states + shape info.