diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index d5e85b986..05b6ba059 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -3345,6 +3345,10 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): # prefill-time tensor so it doesn't leak into ForwardBatch. self.input_embeds = None + self.mamba_cow_src_indices = None + self.mamba_cow_dst_indices = None + self.mamba_clear_indices = None + # Clear context parallel metadata - CP is only for prefill, not decode if hasattr(self, "attn_cp_metadata") and self.attn_cp_metadata is not None: self.attn_cp_metadata = None 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 index 4c3df0a42..ca47ad700 100644 --- a/test/registered/unit/managers/test_schedule_batch_prepare_for_decode.py +++ b/test/registered/unit/managers/test_schedule_batch_prepare_for_decode.py @@ -6,7 +6,7 @@ import torch from sglang.srt.runtime_context import get_context from sglang.test.ci.ci_register import register_cpu_ci -from sglang.test.test_utils import maybe_stub_sgl_kernel +from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel maybe_stub_sgl_kernel() @@ -39,6 +39,25 @@ def _make_decode_batch(): return batch +class TestPrepareForDecodeMambaInit(CustomTestCase): + def test_spec_decode_drops_extend_mamba_init_metadata(self): + batch = ScheduleBatch(reqs=[]) + batch.spec_algorithm = types.SimpleNamespace(is_none=lambda: False) + batch.mamba_cow_src_indices = torch.tensor([1]) + batch.mamba_cow_dst_indices = torch.tensor([2]) + batch.mamba_clear_indices = torch.tensor([3]) + + with patch( + "sglang.srt.speculative.spec_utils.spec_prepare_for_decode" + ) as prepare: + batch.prepare_for_decode() + + prepare.assert_called_once_with(batch) + self.assertIsNone(batch.mamba_cow_src_indices) + self.assertIsNone(batch.mamba_cow_dst_indices) + self.assertIsNone(batch.mamba_clear_indices) + + 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."""