From 8392c36bce22d3c54a6381aea6e003d6bf57719b Mon Sep 17 00:00:00 2001 From: Aurick Qiao Date: Mon, 7 Sep 2026 12:41:03 -0700 Subject: [PATCH] [Bugfix][Mamba] Clear deferred init metadata before speculative decode (#37165) --- python/sglang/srt/managers/schedule_batch.py | 4 ++++ .../test_schedule_batch_prepare_for_decode.py | 21 ++++++++++++++++++- 2 files changed, 24 insertions(+), 1 deletion(-) 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."""