Remove dead ScheduleBatch fields and avoid inplace seq_lens bump (#30669)

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",

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
-            # New-tensor avoids racing model_worker_batch refs queued for
+        # overlap forward.
-            # overlap forward.
+        self.seq_lens = self.seq_lens + 1
-            self.seq_lens = self.seq_lens + 1
+        self.seq_lens_cpu = self.seq_lens_cpu + 1
-            self.seq_lens_cpu = self.seq_lens_cpu + 1
+        self.orig_seq_lens = self.orig_seq_lens + 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)
         # 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,

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:
+        seq_lens_cpu = batch.seq_lens_cpu
-            # 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
 
         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,

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()