[Bugfix][Mamba] Clear deferred init metadata before speculative decode (#37165)
This commit is contained in:
@@ -3345,6 +3345,10 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
|
|||||||
# prefill-time tensor so it doesn't leak into ForwardBatch.
|
# prefill-time tensor so it doesn't leak into ForwardBatch.
|
||||||
self.input_embeds = None
|
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
|
# 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:
|
if hasattr(self, "attn_cp_metadata") and self.attn_cp_metadata is not None:
|
||||||
self.attn_cp_metadata = None
|
self.attn_cp_metadata = None
|
||||||
|
|||||||
@@ -6,7 +6,7 @@ import torch
|
|||||||
|
|
||||||
from sglang.srt.runtime_context import get_context
|
from sglang.srt.runtime_context import get_context
|
||||||
from sglang.test.ci.ci_register import register_cpu_ci
|
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()
|
maybe_stub_sgl_kernel()
|
||||||
|
|
||||||
@@ -39,6 +39,25 @@ def _make_decode_batch():
|
|||||||
return 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):
|
class TestPrepareForDecodeSeqLensOwnership(unittest.TestCase):
|
||||||
def test_decode_seq_lens_bump_is_out_of_place(self):
|
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."""
|
"""Each prepare_for_decode call rebinds seq-lens tensors to new +1 objects without mutating the old ones."""
|
||||||
|
|||||||
Reference in New Issue
Block a user