diff --git a/python/sglang/srt/batch_overlap/two_batch_overlap.py b/python/sglang/srt/batch_overlap/two_batch_overlap.py index d6e2f0ef1..ca44b0021 100644 --- a/python/sglang/srt/batch_overlap/two_batch_overlap.py +++ b/python/sglang/srt/batch_overlap/two_batch_overlap.py @@ -723,7 +723,6 @@ class TboForwardBatchPreparer: for key in [ "forward_mode", "is_extend_in_batch", - "all_extend_in_batch", "return_logprob", "can_run_dp_cuda_graph", "can_run_dp_breakable_cuda_graph", diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index a92cfdc12..e808d3cf9 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -1892,7 +1892,6 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): # For DP attention is_extend_in_batch: bool = False - all_extend_in_batch: bool = False # plumbing for downstream forks (PR #19639) can_run_dp_cuda_graph: bool = False can_run_dp_breakable_cuda_graph: bool = False tbo_split_seq_index: Optional[int] = None @@ -1956,7 +1955,6 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): spec_info: Optional[SpecInput] = None # === One-shot per-forward overrides; init_new consumes and resets === - seq_lens_cpu_cache: torch.Tensor = None capture_hidden_mode: Optional[CaptureHiddenMode] = None return_hidden_states_before_norm: bool = False @@ -2811,16 +2809,11 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): req.kv_committed_len += 1 req.kv_allocated_len += 1 - if self.enable_overlap: - # New-tensor avoids racing model_worker_batch refs queued for - # overlap forward. - self.seq_lens = self.seq_lens + 1 - self.seq_lens_cpu = self.seq_lens_cpu + 1 - self.orig_seq_lens = self.orig_seq_lens + 1 - else: - self.seq_lens.add_(1) - self.seq_lens_cpu.add_(1) - self.orig_seq_lens.add_(1) + # New-tensor avoids racing model_worker_batch refs queued for + # overlap forward. + self.seq_lens = self.seq_lens + 1 + self.seq_lens_cpu = self.seq_lens_cpu + 1 + self.orig_seq_lens = self.orig_seq_lens + 1 # Sum is recomputed lazily by ForwardBatch.init_new. self.seq_lens_sum = None @@ -3013,7 +3006,6 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): can_run_dp_cuda_graph=self.can_run_dp_cuda_graph, can_run_dp_breakable_cuda_graph=self.can_run_dp_breakable_cuda_graph, is_extend_in_batch=self.is_extend_in_batch, - all_extend_in_batch=self.all_extend_in_batch, is_prefill_only=self.is_prefill_only, seq_lens_cpu=self.seq_lens_cpu, enable_overlap=self.enable_overlap, diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 988956e2f..357d80589 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -396,8 +396,6 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): # For DP attention is_extend_in_batch: bool = False - # Mirrors ScheduleBatch.all_extend_in_batch; kept for downstream forks. - all_extend_in_batch: bool = False can_run_dp_cuda_graph: bool = False can_run_dp_breakable_cuda_graph: bool = False global_forward_mode: Optional[ForwardMode] = None @@ -632,8 +630,6 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): # for the contract. capture_hidden_mode = batch.capture_hidden_mode batch.capture_hidden_mode = None - seq_lens_cpu_cache = batch.seq_lens_cpu_cache - batch.seq_lens_cpu_cache = None return_hidden_states_before_norm = batch.return_hidden_states_before_norm batch.return_hidden_states_before_norm = False @@ -668,18 +664,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): # ScheduleBatch.sampling_info is already swapped to the forward-only # copy by Scheduler.run_batch under overlap mode (see save/restore # block there). Use it directly. - if seq_lens_cpu_cache is not None: - # Stale-cache guard: shape must match current GPU seq_lens. Mismatch - # means caller forgot to refresh the override after batch size - # changed (e.g. filter/merge_batch); using a stale cache would - # propagate wrong CPU mirror to downstream DP / cudagraph logic. - assert seq_lens_cpu_cache.shape == batch.seq_lens.shape, ( - f"seq_lens_cpu_cache shape {seq_lens_cpu_cache.shape} != " - f"seq_lens {batch.seq_lens.shape}; stale override on batch?" - ) - seq_lens_cpu = seq_lens_cpu_cache - else: - seq_lens_cpu = batch.seq_lens_cpu + seq_lens_cpu = batch.seq_lens_cpu if batch.seq_lens_sum is None and seq_lens_cpu is not None: batch.seq_lens_sum = int(seq_lens_cpu.sum()) @@ -711,7 +696,6 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): # Scalar config / flags return_logprob=batch.return_logprob, is_extend_in_batch=batch.is_extend_in_batch, - all_extend_in_batch=batch.all_extend_in_batch, can_run_dp_cuda_graph=batch.can_run_dp_cuda_graph, can_run_dp_breakable_cuda_graph=batch.can_run_dp_breakable_cuda_graph, global_forward_mode=batch.global_forward_mode, diff --git a/test/registered/unit/managers/test_schedule_batch_prepare_for_decode.py b/test/registered/unit/managers/test_schedule_batch_prepare_for_decode.py new file mode 100644 index 000000000..a9b66b2bc --- /dev/null +++ b/test/registered/unit/managers/test_schedule_batch_prepare_for_decode.py @@ -0,0 +1,86 @@ +import types +import unittest +from unittest.mock import patch + +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 + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +def _make_req(): + return types.SimpleNamespace( + decode_batch_idx=0, + kv_committed_len=3, + kv_allocated_len=3, + ) + + +def _make_decode_batch(): + batch = ScheduleBatch(reqs=[_make_req(), _make_req()]) + batch.device = "cpu" + batch.model_config = types.SimpleNamespace(is_encoder_decoder=False) + batch.enable_overlap = False + batch.spec_algorithm = types.SimpleNamespace(is_none=lambda: True) + batch.sampling_info = types.SimpleNamespace( + penalizer_orchestrator=types.SimpleNamespace(is_required=False) + ) + batch.hisparse_coordinator = None + batch.seq_lens = torch.tensor([3, 5], dtype=torch.int64) + batch.seq_lens_cpu = torch.tensor([3, 5], dtype=torch.int64) + batch.orig_seq_lens = torch.tensor([3, 5], dtype=torch.int32) + return batch + + +class TestPrepareForDecodeSeqLensOwnership(unittest.TestCase): + def test_decode_seq_lens_bump_is_out_of_place(self): + """Each prepare_for_decode call rebinds seq-lens tensors to new +1 objects without mutating the old ones.""" + batch = _make_decode_batch() + + server_args = types.SimpleNamespace( + enable_mamba_extra_buffer=lambda: False, + ) + with ( + patch( + "sglang.srt.managers.schedule_batch.alloc_for_decode", + return_value=torch.tensor([6, 7], dtype=torch.int64), + ), + patch( + "sglang.srt.managers.schedule_batch.get_server_args", + return_value=server_args, + ), + ): + for step in range(1, 3): + prev_seq_lens = batch.seq_lens + prev_seq_lens_cpu = batch.seq_lens_cpu + prev_orig_seq_lens = batch.orig_seq_lens + prev_values = ( + prev_seq_lens.clone(), + prev_seq_lens_cpu.clone(), + prev_orig_seq_lens.clone(), + ) + + batch.prepare_for_decode() + + self.assertIsNot(batch.seq_lens, prev_seq_lens) + self.assertIsNot(batch.seq_lens_cpu, prev_seq_lens_cpu) + self.assertIsNot(batch.orig_seq_lens, prev_orig_seq_lens) + expected = torch.tensor([3 + step, 5 + step], dtype=torch.int64) + self.assertTrue(torch.equal(batch.seq_lens, expected)) + self.assertTrue(torch.equal(batch.seq_lens_cpu, expected)) + self.assertTrue( + torch.equal(batch.orig_seq_lens, expected.to(torch.int32)) + ) + self.assertTrue(torch.equal(prev_seq_lens, prev_values[0])) + self.assertTrue(torch.equal(prev_seq_lens_cpu, prev_values[1])) + self.assertTrue(torch.equal(prev_orig_seq_lens, prev_values[2])) + + +if __name__ == "__main__": + unittest.main()