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
.claude/rules/schedule-batch-out-of-place-mutation.md
+30
View File
@@ -0,0 +1,30 @@
---
paths:
- "python/sglang/srt/**/*.py"
---
# Mutate ScheduleBatch fields out of place
Never mutate a `ScheduleBatch` direct field in place — no `.extend()`/`.append()`,
`+=`/`-=`/`|=`, tensor `.add_()`/`.fill_()`, or slice/index assignment. Build the
new value and rebind the field:
```python
# Bad
self.reqs.extend(other.reqs)
self.seq_lens.add_(1)
self.extend_lens[i] -= encoder_len
# Good
self.reqs = self.reqs + other.reqs
self.seq_lens = self.seq_lens + 1
lens = self.extend_lens[:]; lens[i] -= encoder_len; self.extend_lens = lens # loop a copy, rebind once
```
Why: `copy()` snapshots and the overlap scheduler's queued references rely on old
objects staying frozen (the prerequisite for the per-step immutable ScheduleBatch
refactor); derived values like `seq_lens_sum` only recompute safely when fields
are rebound.
Out of scope (mutation intended): penalizer buffers, CUDA graph static buffers,
`ForwardBatch` fields, attention-backend metadata.
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.
test/registered/unit/managers/test_schedule_batch_out_of_place.py
+236
View File
@@ -0,0 +1,236 @@
import dataclasses
import types
import unittest
from array import array
from unittest.mock import MagicMock
import torch
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import maybe_stub_sgl_kernel
maybe_stub_sgl_kernel()
from sglang.srt.managers.schedule_batch import ScheduleBatch # noqa: E402
from sglang.srt.model_executor.forward_batch_info import ForwardMode # noqa: E402
from sglang.srt.utils.common import Range # noqa: E402
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
AUTO_FILL_EXCLUDED_FIELDS = ["reqs"]
def make_schedule_batch(bs: int, **overrides) -> ScheduleBatch:
batch = ScheduleBatch(reqs=overrides.pop("reqs"))
for field in dataclasses.fields(ScheduleBatch):
name = field.name
if name in overrides or name in AUTO_FILL_EXCLUDED_FIELDS:
continue
annotation = str(field.type)
if "List" in annotation or "list[" in annotation:
setattr(batch, name, [f"{name}-{i}" for i in range(bs)])
elif "Tensor" in annotation:
setattr(batch, name, torch.arange(bs, dtype=torch.int64))
for name, value in overrides.items():
setattr(batch, name, value)
return batch
def _snapshot_mutable_fields(batch):
snapshot = []
for field in dataclasses.fields(batch):
value = getattr(batch, field.name, None)
if isinstance(value, MagicMock):
continue
if isinstance(value, list):
snapshot.append((field.name, value, list(value)))
elif isinstance(value, torch.Tensor):
snapshot.append((field.name, value, value.clone()))
return snapshot
def _assert_snapshot_not_mutated(test_case, snapshot):
for name, obj, value_copy in snapshot:
if isinstance(obj, list):
test_case.assertEqual(obj, value_copy, f"field {name} was mutated in place")
else:
test_case.assertTrue(
torch.equal(obj, value_copy), f"field {name} was mutated in place"
)
class _FakeReq:
def __init__(self, rid, origin_len, output_len):
self.rid = rid
self.origin_input_ids = list(range(origin_len))
self.output_ids = list(range(output_len))
self.full_untruncated_fill_ids = list(range(origin_len + output_len))
self.extend_range = None
def _refresh_fill_ids(self):
self.full_untruncated_fill_ids = self.origin_input_ids + self.output_ids
def set_extend_range(self, start, end):
self.extend_range = Range(start, end)
class TestMergeBatchOutOfPlace(unittest.TestCase):
def test_merge_batch_rebinds_lists_without_mutating_either_side(self):
"""merge_batch must build new list objects; no field of either side may be mutated in place."""
self_batch = make_schedule_batch(
2,
reqs=[types.SimpleNamespace(rid="a"), types.SimpleNamespace(rid="b")],
model_config=types.SimpleNamespace(is_encoder_decoder=False),
sampling_info=MagicMock(),
return_logprob=True,
top_logprobs_nums=[1, 2],
token_ids_logprobs=[[10], [20]],
)
other_batch = make_schedule_batch(
1,
reqs=[types.SimpleNamespace(rid="c")],
model_config=types.SimpleNamespace(is_encoder_decoder=False),
sampling_info=MagicMock(),
return_logprob=True,
top_logprobs_nums=[3],
token_ids_logprobs=[[30]],
)
self_reqs_before = self_batch.reqs
self_top_before = self_batch.top_logprobs_nums
self_snapshot = _snapshot_mutable_fields(self_batch)
other_snapshot = _snapshot_mutable_fields(other_batch)
self_batch.merge_batch(other_batch)
self.assertEqual([r.rid for r in self_batch.reqs], ["a", "b", "c"])
self.assertEqual(self_batch.top_logprobs_nums, [1, 2, 3])
self.assertEqual(self_batch.token_ids_logprobs, [[10], [20], [30]])
self.assertIsNot(self_batch.reqs, self_reqs_before)
self.assertIsNot(self_batch.top_logprobs_nums, self_top_before)
_assert_snapshot_not_mutated(self, self_snapshot)
_assert_snapshot_not_mutated(self, other_snapshot)
class TestMixWithRunningOutOfPlace(unittest.TestCase):
def test_mix_with_running_rebinds_extend_fields_without_mutating_either_side(self):
"""mix_with_running must append via rebound lists; no field of either side may be mutated in place."""
extend_batch = make_schedule_batch(
2,
reqs=[types.SimpleNamespace(rid="e1"), types.SimpleNamespace(rid="e2")],
model_config=types.SimpleNamespace(is_encoder_decoder=False),
sampling_info=MagicMock(),
return_logprob=False,
forward_mode=ForwardMode.EXTEND,
enable_overlap=False,
is_prefill_only=True,
out_cache_loc=torch.arange(6, dtype=torch.int64),
prefix_lens=[0, 0],
extend_lens=[3, 3],
extend_num_tokens=6,
extend_logprob_start_lens=[0, 0],
)
running_batch = make_schedule_batch(
1,
reqs=[_FakeReq("r1", origin_len=4, output_len=2)],
model_config=types.SimpleNamespace(is_encoder_decoder=False),
sampling_info=MagicMock(),
return_logprob=False,
forward_mode=ForwardMode.DECODE,
out_cache_loc=torch.arange(6, 7, dtype=torch.int64),
)
extend_prefix_before = extend_batch.prefix_lens
extend_lens_before = extend_batch.extend_lens
extend_snapshot = _snapshot_mutable_fields(extend_batch)
running_snapshot = _snapshot_mutable_fields(running_batch)
extend_batch.mix_with_running(running_batch)
self.assertEqual(extend_batch.forward_mode, ForwardMode.MIXED)
self.assertIs(extend_batch.mix_running_indices, running_batch.req_pool_indices)
self.assertEqual([r.rid for r in extend_batch.reqs], ["e1", "e2", "r1"])
self.assertTrue(
torch.equal(extend_batch.out_cache_loc, torch.arange(7, dtype=torch.int64))
)
# delta is -1 without overlap: 4 origin + 2 output - 1
self.assertEqual(extend_batch.prefix_lens, [0, 0, 5])
self.assertEqual(extend_batch.extend_lens, [3, 3, 1])
self.assertEqual(extend_batch.extend_num_tokens, 7)
self.assertEqual(extend_batch.extend_logprob_start_lens, [0, 0, 0])
self.assertFalse(extend_batch.is_prefill_only)
self.assertIsNot(extend_batch.prefix_lens, extend_prefix_before)
self.assertIsNot(extend_batch.extend_lens, extend_lens_before)
_assert_snapshot_not_mutated(self, extend_snapshot)
_assert_snapshot_not_mutated(self, running_snapshot)
class TestPrepareEncoderInfoExtendOutOfPlace(unittest.TestCase):
def test_prepare_encoder_info_extend_rebinds_lens_without_mutating_old_lists(self):
"""prepare_encoder_info_extend must strip encoder tokens via rebound lists; old list objects stay intact."""
req_with_image = types.SimpleNamespace(
rid="img",
multimodal_inputs=types.SimpleNamespace(num_image_tokens=2),
prefix_indices=[],
extend_range=Range(0, 5),
logprob_start_len=0,
)
req_text_only = types.SimpleNamespace(
rid="txt",
multimodal_inputs=None,
prefix_indices=[],
extend_range=Range(0, 4),
logprob_start_len=0,
)
batch = make_schedule_batch(
2,
reqs=[req_with_image, req_text_only],
device="cpu",
forward_mode=ForwardMode.EXTEND,
out_cache_loc=torch.arange(9, dtype=torch.int64),
prefix_lens=[0, 0],
extend_lens=[5, 4],
extend_num_tokens=9,
extend_logprob_start_lens=[0, 0],
extend_input_logprob_token_ids=torch.arange(9, dtype=torch.int64),
)
prefix_before = batch.prefix_lens
extend_before = batch.extend_lens
logprob_start_before = batch.extend_logprob_start_lens
snapshot = _snapshot_mutable_fields(batch)
batch.prepare_encoder_info_extend(
input_ids=[array("q", range(5)), array("q", range(4))],
seq_lens=[5, 4],
)
self.assertEqual(batch.encoder_lens_cpu, [2, 0])
self.assertEqual(batch.encoder_cached, [False, True])
self.assertEqual(batch.extend_lens, [3, 4])
self.assertEqual(batch.prefix_lens, [0, 0])
self.assertEqual(batch.extend_num_tokens, 7)
self.assertTrue(
torch.equal(batch.out_cache_loc, torch.arange(2, 9, dtype=torch.int64))
)
self.assertTrue(
torch.equal(batch.encoder_out_cache_loc, torch.arange(2, dtype=torch.int64))
)
self.assertTrue(
torch.equal(batch.seq_lens_cpu, torch.tensor([3, 4], dtype=torch.int64))
)
self.assertTrue(
torch.equal(
batch.extend_input_logprob_token_ids,
torch.arange(2, 9, dtype=torch.int64),
)
)
self.assertEqual(batch.extend_logprob_start_lens, [0, 0])
self.assertIsNot(batch.prefix_lens, prefix_before)
self.assertIsNot(batch.extend_lens, extend_before)
self.assertIsNot(batch.extend_logprob_start_lens, logprob_start_before)
_assert_snapshot_not_mutated(self, snapshot)
if __name__ == "__main__":
unittest.main()