[Spec] Support ReplaySSM spec-verify for DFlash and DSPARK on hybrid GDN models

python/sglang/srt/mem_cache/kv_cache_configurator.py

@@ -1139,18 +1139,27 @@ def _build_hybrid_req_pool(
         max_num_reqs: int,
         extra_max_context_len: int,
     ) -> ReqToTokenPool:
-        # DSPARK/DFLASH commit routes through the backend fold (KDA-only); a
-        # non-KDA model there would scatter a None intermediate_ssm and crash.
+        # DFLASH and DSPARK route GDN ReplaySSM through the common
+        # fold-every-commit hook. DSPARK's compact/cap-accept layouts are ragged,
+        # while the GDN ring-write currently requires a dense fixed-width window.
         _algo = (get_spec().speculative_algorithm or "").upper()
         if (
             get_exec().mamba.enable_linear_replayssm_spec
-            and _algo in ("DSPARK", "DFLASH")
-            and self.hybrid_kda_config is None
+            and _algo == "DSPARK"
+            and self.hybrid_gdn_config is not None
         ):
-            raise ValueError(
-                "--enable-linear-replayssm-spec with DSPARK/DFLASH requires a KDA "
-                "model; got a non-KDA model."
+            from sglang.srt.speculative.ragged_verify import (
+                RaggedVerifyMode,
+                read_ragged_verify_mode,
             )
+
+            ragged_mode = read_ragged_verify_mode()
+            if ragged_mode is not RaggedVerifyMode.STATIC:
+                raise ValueError(
+                    "--enable-linear-replayssm-spec with DSPARK + GDN currently "
+                    "requires SGLANG_RAGGED_VERIFY_MODE=static; got "
+                    f"{ragged_mode.value!r}."
+                )
         req_to_token_pool = HybridReqToTokenPool(
             size=max_num_reqs,
             mamba_size=get_schedule().max_mamba_cache_size,

python/sglang/srt/speculative/dflash_worker_v2.py

@@ -95,6 +95,7 @@
     GrammarTree,
     assign_req_to_token_pool_func,
     build_grammar_vocab_mask,
+    commit_mamba_states_after_verify,
     draft_pp_context,
     draft_tp_context,
 )
@@ -2021,6 +2022,41 @@ def _update_target_mamba_state_after_verify(
         """
         if not self._need_mamba_verify_commit:
             return
+
+        bs = int(commit_lens.shape[0])
+        block_size = int(self.block_size)
+        req_pool = self.target_worker.model_runner.req_to_token_pool
+        mamba_pool = getattr(req_pool, "mamba_pool", None)
+        use_replayssm_commit = (
+            mamba_pool is not None
+            and (
+                getattr(mamba_pool, "replayssm_spec_fold", False)
+                or getattr(mamba_pool, "replayssm_cache_base", None) is not None
+            )
+            and not getattr(mamba_pool, "replayssm_is_kda", False)
+        )
+        if use_replayssm_commit:
+            offsets = torch.arange(
+                0,
+                bs * block_size,
+                step=block_size,
+                dtype=torch.int32,
+                device=commit_lens.device,
+            )
+            accept_index = offsets[:, None] + torch.arange(
+                block_size,
+                dtype=torch.int32,
+                device=commit_lens.device,
+            )
+            commit_mamba_states_after_verify(
+                target_worker=self.target_worker,
+                batch=batch,
+                accept_lens=commit_lens,
+                accept_index=accept_index,
+                draft_token_num=block_size,
+            )
+            return
+
         attn_backend = self.target_worker.model_runner.attn_backend
 
         last_correct_step_indices = commit_lens.to(torch.int64) - 1

python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py

@@ -79,6 +79,7 @@
 from sglang.srt.speculative.spec_utils import (
     GrammarTree,
     build_grammar_vocab_mask,
+    commit_mamba_states_after_verify,
     draft_pp_context,
     draft_tp_context,
     prepare_mamba_track_for_verify,
@@ -1098,6 +1099,31 @@ def _commit_target_mamba_states_after_verify(
         # layout would need the accept-index mapping the shared spec_utils
         # commit helper does.
         assert get_spec().speculative_eagle_topk in (None, 1)
+
+        req_pool = self.target_worker.model_runner.req_to_token_pool
+        mamba_pool = getattr(req_pool, "mamba_pool", None)
+        use_gdn_replayssm_fold = (
+            mamba_pool is not None
+            and getattr(mamba_pool, "replayssm_spec_fold", False)
+            and not getattr(mamba_pool, "replayssm_is_kda", False)
+        )
+        if use_gdn_replayssm_fold:
+            bs = int(commit_lens.shape[0])
+            draft_token_num = int(self.verify_num_draft_tokens)
+            accept_index = torch.arange(
+                bs * draft_token_num,
+                dtype=commit_lens.dtype,
+                device=commit_lens.device,
+            ).view(bs, draft_token_num)
+            commit_mamba_states_after_verify(
+                target_worker=self.target_worker,
+                batch=batch,
+                accept_lens=commit_lens,
+                accept_index=accept_index,
+                draft_token_num=draft_token_num,
+            )
+            return
+
         attn_backend = self.target_worker.model_runner.attn_backend
 
         last_correct_step_indices = commit_lens.to(torch.int64) - 1

test/registered/unit/spec/test_dflash_overlap_hostsync.py

@@ -5,6 +5,7 @@
 
 import unittest
 from types import SimpleNamespace
+from unittest.mock import patch
 
 import torch
 
@@ -74,6 +75,185 @@ def test_sawtooth_counterexample(self):
         self.assertGreaterEqual(int(bound), int(exact_true))
 
 
+class TestDFlashMambaCommit(CustomTestCase):
+    def test_fold_replayssm_uses_common_mamba_commit_hook(self):
+        from sglang.srt.speculative import dflash_worker_v2
+        from sglang.srt.speculative.dflash_worker_v2 import DFlashWorkerV2
+
+        calls = []
+
+        def fake_commit(**kwargs):
+            calls.append(kwargs)
+
+        old_commit = dflash_worker_v2.commit_mamba_states_after_verify
+        dflash_worker_v2.commit_mamba_states_after_verify = fake_commit
+        try:
+            worker = SimpleNamespace(
+                _need_mamba_verify_commit=True,
+                block_size=4,
+                target_worker=SimpleNamespace(
+                    model_runner=SimpleNamespace(
+                        req_to_token_pool=SimpleNamespace(
+                            mamba_pool=SimpleNamespace(
+                                replayssm_spec_fold=True,
+                                replayssm_is_kda=False,
+                            )
+                        )
+                    )
+                ),
+            )
+            batch = object()
+            commit_lens = torch.tensor([1, 3], dtype=torch.int32)
+            DFlashWorkerV2._update_target_mamba_state_after_verify(
+                worker,
+                batch=batch,
+                seq_lens_pre_verify=torch.tensor([10, 20], dtype=torch.int64),
+                seq_lens_post_verify=torch.tensor([11, 23], dtype=torch.int64),
+                commit_lens=commit_lens,
+            )
+        finally:
+            dflash_worker_v2.commit_mamba_states_after_verify = old_commit
+
+        self.assertEqual(len(calls), 1)
+        call = calls[0]
+        self.assertIs(call["target_worker"], worker.target_worker)
+        self.assertIs(call["batch"], batch)
+        self.assertIs(call["accept_lens"], commit_lens)
+        self.assertEqual(call["draft_token_num"], 4)
+        torch.testing.assert_close(
+            call["accept_index"],
+            torch.tensor([[0, 1, 2, 3], [4, 5, 6, 7]], dtype=torch.int32),
+        )
+
+    def test_non_replayssm_keeps_legacy_mamba_commit_path(self):
+        from sglang.srt.speculative.dflash_worker_v2 import DFlashWorkerV2
+
+        calls = []
+
+        class FakeAttnBackend:
+            def update_mamba_state_after_mtp_verify(self, **kwargs):
+                calls.append(kwargs)
+
+        worker = SimpleNamespace(
+            _need_mamba_verify_commit=True,
+            block_size=4,
+            target_worker=SimpleNamespace(
+                model_runner=SimpleNamespace(
+                    req_to_token_pool=SimpleNamespace(mamba_pool=SimpleNamespace()),
+                    attn_backend=FakeAttnBackend(),
+                    model=object(),
+                )
+            ),
+        )
+        batch = SimpleNamespace(
+            mamba_track_indices=None,
+            req_pool_indices=torch.tensor([0, 1], dtype=torch.int64),
+        )
+        commit_lens = torch.tensor([1, 3], dtype=torch.int32)
+
+        DFlashWorkerV2._update_target_mamba_state_after_verify(
+            worker,
+            batch=batch,
+            seq_lens_pre_verify=torch.tensor([10, 20], dtype=torch.int64),
+            seq_lens_post_verify=torch.tensor([11, 23], dtype=torch.int64),
+            commit_lens=commit_lens,
+        )
+
+        self.assertEqual(len(calls), 1)
+        torch.testing.assert_close(
+            calls[0]["last_correct_step_indices"],
+            torch.tensor([0, 2], dtype=torch.int64),
+        )
+        self.assertIsNone(calls[0]["mamba_track_indices"])
+        self.assertIsNone(calls[0]["mamba_steps_to_track"])
+
+
+class TestDSparkMambaCommit(CustomTestCase):
+    def _run_commit(self, *, is_kda):
+        from sglang.srt.speculative.dspark_components import dspark_worker_v2
+
+        fold_calls, backend_calls = [], []
+        backend = SimpleNamespace(
+            update_mamba_state_after_mtp_verify=lambda **kwargs: backend_calls.append(
+                kwargs
+            )
+        )
+        worker = SimpleNamespace(
+            _need_mamba_verify_commit=True,
+            verify_num_draft_tokens=4,
+            target_worker=SimpleNamespace(
+                model_runner=SimpleNamespace(
+                    req_to_token_pool=SimpleNamespace(
+                        mamba_pool=SimpleNamespace(
+                            replayssm_spec_fold=True,
+                            replayssm_is_kda=is_kda,
+                        )
+                    ),
+                    attn_backend=backend,
+                    model=object(),
+                )
+            ),
+        )
+        batch = SimpleNamespace(
+            mamba_track_indices=None, req_pool_indices=torch.tensor([2, 5])
+        )
+        lengths = torch.tensor([1, 3], dtype=torch.int32)
+        with (
+            patch.object(
+                dspark_worker_v2,
+                "get_spec",
+                return_value=SimpleNamespace(speculative_eagle_topk=1),
+            ),
+            patch.object(
+                dspark_worker_v2,
+                "commit_mamba_states_after_verify",
+                side_effect=lambda **kwargs: fold_calls.append(kwargs),
+            ),
+        ):
+            dspark_worker_v2.DSparkWorkerV2._commit_target_mamba_states_after_verify(
+                worker,
+                batch=batch,
+                seq_lens_pre_verify=torch.tensor([10, 20]),
+                seq_lens_post_verify=torch.tensor([11, 23]),
+                commit_lens=lengths,
+            )
+        return fold_calls, backend_calls, batch, lengths
+
+    def test_gdn_fold_uses_dense_request_indices(self):
+        fold_calls, backend_calls, batch, lengths = self._run_commit(is_kda=False)
+        self.assertEqual(len(fold_calls), 1)
+        self.assertEqual(backend_calls, [])
+        self.assertIs(fold_calls[0]["batch"], batch)
+        self.assertIs(fold_calls[0]["accept_lens"], lengths)
+        self.assertEqual(fold_calls[0]["draft_token_num"], 4)
+        torch.testing.assert_close(
+            fold_calls[0]["accept_index"],
+            torch.tensor([[0, 1, 2, 3], [4, 5, 6, 7]], dtype=torch.int32),
+        )
+
+    def test_kda_keeps_backend_commit(self):
+        fold_calls, backend_calls, batch, _ = self._run_commit(is_kda=True)
+        self.assertEqual(fold_calls, [])
+        self.assertEqual(len(backend_calls), 1)
+        torch.testing.assert_close(
+            backend_calls[0]["last_correct_step_indices"], torch.tensor([0, 2])
+        )
+        self.assertIs(backend_calls[0]["req_pool_indices"], batch.req_pool_indices)
+
+    def test_non_hybrid_commit_is_noop(self):
+        from sglang.srt.speculative.dspark_components.dspark_worker_v2 import (
+            DSparkWorkerV2,
+        )
+
+        DSparkWorkerV2._commit_target_mamba_states_after_verify(
+            SimpleNamespace(_need_mamba_verify_commit=False),
+            batch=object(),
+            seq_lens_pre_verify=torch.tensor([10]),
+            seq_lens_post_verify=torch.tensor([11]),
+            commit_lens=torch.tensor([1]),
+        )
+
+
 class _FakeTpGroup:
     """Single-process stand-in for the TP GroupCoordinator: replays the
     concatenation of all ranks' recorded all-gather inputs."""