diff --git a/.claude/rules/schedule-batch-out-of-place-mutation.md b/.claude/rules/schedule-batch-out-of-place-mutation.md new file mode 100644 index 000000000..f7ddb02b6 --- /dev/null +++ b/.claude/rules/schedule-batch-out-of-place-mutation.md @@ -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. diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 72f90779a..8336955d1 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -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, diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 4c937c26f..5934cccfd 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -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. diff --git a/test/registered/unit/managers/test_schedule_batch_out_of_place.py b/test/registered/unit/managers/test_schedule_batch_out_of_place.py new file mode 100644 index 000000000..572648e79 --- /dev/null +++ b/test/registered/unit/managers/test_schedule_batch_out_of_place.py @@ -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()