Fix recurrent state loss on decode retraction (#35957)

python/sglang/srt/managers/schedule_batch.py

@@ -109,7 +109,7 @@ from sglang.srt.mem_cache.common import (
     release_kv_cache,
     retraction_backup,
 )
-from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
+from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool, ReqToTokenPool
 from sglang.srt.mem_cache.radix_cache import RadixKey
 from sglang.srt.model_executor.forward_batch_info import (
     CaptureHiddenMode,
@@ -1726,15 +1726,30 @@ class Req(ReqDllmMixin):
                 self.weight_version_events, num_kept_tokens=self.send_token_offset
             )
 
+    def _mamba_pool_needing_backup(self, req_to_token_pool, allocator):
+        if allocator.get_kvcache().cpu_copy_carries_mamba:
+            return None
+        if not isinstance(req_to_token_pool, HybridReqToTokenPool):
+            return None
+        return req_to_token_pool.mamba_pool
+
     def offload_kv_cache(self, req_to_token_pool, token_to_kv_pool_allocator):
         token_indices = req_to_token_pool.req_to_token[
             self.req_pool_idx, : self.seqlen - 1
         ]
         # Copies over both the kv cache and mamba state if available
+        mamba_pool = self._mamba_pool_needing_backup(
+            req_to_token_pool, token_to_kv_pool_allocator
+        )
         self.retraction_backup = RetractionBackup(
             cpu_tensors=token_to_kv_pool_allocator.get_cpu_copy(
                 token_indices, mamba_indices=self.mamba_pool_idx
-            )
+            ),
+            mamba_cpu=(
+                mamba_pool.get_cpu_copy(self.mamba_pool_idx.unsqueeze(0))
+                if mamba_pool is not None and self.mamba_pool_idx is not None
+                else None
+            ),
         )
 
     def load_kv_cache(self, req_to_token_pool, token_to_kv_pool_allocator):
@@ -1743,6 +1758,11 @@ class Req(ReqDllmMixin):
             self.req_pool_idx, : self.seqlen - 1
         ]
         # Loads both the kv cache and mamba state if exists
+        mamba_cpu = self.retraction_backup.mamba_cpu
+        if mamba_cpu is not None and self.mamba_pool_idx is not None:
+            req_to_token_pool.mamba_pool.load_cpu_copy(
+                mamba_cpu, self.mamba_pool_idx.unsqueeze(0)
+            )
         token_to_kv_pool_allocator.load_cpu_copy(
             self.retraction_backup.cpu_tensors,
             token_indices,

python/sglang/srt/mem_cache/common.py

@@ -35,6 +35,8 @@ class RetractionBackup(NamedTuple):
     cpu_tensors: Any = None
     host_indices: Optional[torch.Tensor] = None
     pool_transfers: Optional[list[PoolTransfer]] = None
+    # Set when the KV pool leaves the recurrent state to the caller.
+    mamba_cpu: Any = None
 
 
 def kv_to_page_indices(kv_indices: torch.Tensor, page_size: int) -> np.ndarray:

python/sglang/srt/mem_cache/kv_cache_builder.py

@@ -146,6 +146,18 @@ def _register_legacy_hicache_draft(
 BACKUP_ONLY_HICACHE_RATIO = 0.2
 
 
+def uses_ssm_state(model_config) -> bool:
+    """Whether the model keeps recurrent/conv state alongside its attention KV."""
+    spec = linear_attn_model_spec(model_config)
+    return (
+        hybrid_gdn_config(model_config) is not None
+        or mamba2_config(model_config) is not None
+        or (spec.uses_mamba_radix_cache if spec is not None else False)
+        or kimi_linear_config(model_config) is not None
+        or hybrid_lightning_config(model_config) is not None
+    )
+
+
 def resolve_decode_retraction_backup(*, tp_worker: BaseTpWorker) -> str:
     """Resolve the retraction backend onto the config bags and return it.
 
@@ -165,8 +177,13 @@ def resolve_decode_retraction_backup(*, tp_worker: BaseTpWorker) -> str:
             if tp_worker.is_hybrid_swa
             else None
         )
-        supports_host_pool = isinstance(kv_cache, MHATokenToKVPool) or (
-            isinstance(kv_cache, SWAKVPool) and full_tokens_per_layer > 0
+        # Host-pool retraction transfers full and sliding-window components
+        # only, so a model with recurrent state stays on cpu_tensor.
+        supports_host_pool = not uses_ssm_state(
+            tp_worker.model_runner.model_config
+        ) and (
+            isinstance(kv_cache, MHATokenToKVPool)
+            or (isinstance(kv_cache, SWAKVPool) and full_tokens_per_layer > 0)
         )
         schedule = get_schedule()
         priority_preemption = (
@@ -228,15 +245,7 @@ def build_kv_cache(
 
     # Hybrid memory pool
     is_hybrid_swa = tp_worker.is_hybrid_swa
-    _spec = linear_attn_model_spec(tp_worker.model_runner.model_config)
-    _registry_needs_mamba = _spec.uses_mamba_radix_cache if _spec is not None else False
-    is_hybrid_ssm = (
-        hybrid_gdn_config(tp_worker.model_runner.model_config) is not None
-        or mamba2_config(tp_worker.model_runner.model_config) is not None
-        or _registry_needs_mamba
-        or kimi_linear_config(tp_worker.model_runner.model_config) is not None
-        or hybrid_lightning_config(tp_worker.model_runner.model_config) is not None
-    )
+    is_hybrid_ssm = uses_ssm_state(tp_worker.model_runner.model_config)
     is_dsa = is_deepseek_dsa(model_config.hf_config)
 
     sliding_window_size = None

python/sglang/srt/mem_cache/memory_pool.py

@@ -1658,6 +1658,9 @@ class KvBufferDesc:
 class KVCache(abc.ABC):
     layer_shard_enabled: bool = False
     post_capture_active: bool = False
+    # Whether get_cpu_copy/load_cpu_copy carry the recurrent state. False when the
+    # state lives on the request pool instead, and the caller has to move it.
+    cpu_copy_carries_mamba: bool = False
 
     @abc.abstractmethod
     def __init__(
@@ -3684,6 +3687,8 @@ class MHATokenToKVPoolMXFP8(MHATokenToKVPool):
 class HybridLinearKVPool(KVCache):
     """KV cache with separate pools for full and linear attention layers."""
 
+    cpu_copy_carries_mamba = True
+
     def __init__(
         self,
         size: int,

test/registered/unit/mem_cache/test_retraction_mamba_backup.py

@@ -0,0 +1,79 @@
+import unittest
+
+import torch
+
+from sglang.srt.managers.schedule_batch import Req
+from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+MAMBA_STATE = object()
+
+
+class _MambaPool:
+    def __init__(self):
+        self.loaded = None
+
+    def get_cpu_copy(self, indices):
+        return MAMBA_STATE
+
+    def load_cpu_copy(self, state, indices):
+        self.loaded = state
+
+
+class _Allocator:
+    def __init__(self, carries_mamba: bool):
+        self._kv = type("_KV", (), {"cpu_copy_carries_mamba": carries_mamba})()
+        self.loaded_kv = None
+
+    def get_kvcache(self):
+        return self._kv
+
+    def get_cpu_copy(self, indices, mamba_indices=None):
+        return "kv"
+
+    def load_cpu_copy(self, cpu_tensors, indices, mamba_indices=None):
+        self.loaded_kv = cpu_tensors
+
+
+def _req_and_pool():
+    req = object.__new__(Req)
+    req.req_pool_idx = 0
+    req.origin_input_ids = [1, 2]
+    req.output_ids = [3]
+    req.mamba_pool_idx = torch.tensor(1)
+
+    pool = object.__new__(HybridReqToTokenPool)
+    pool.req_to_token = torch.zeros(1, 8, dtype=torch.int64)
+    pool.mamba_pool = _MambaPool()
+    return req, pool
+
+
+class TestRetractionMambaBackup(unittest.TestCase):
+    def test_state_travels_when_kv_pool_leaves_it_behind(self):
+        """A sliding-window KV pool accepts mamba_indices and ignores them, so a
+        retracted request whose recurrent state is not backed up separately
+        resumes on whatever state the reused slot happens to hold."""
+        req, pool = _req_and_pool()
+        allocator = _Allocator(carries_mamba=False)
+
+        req.offload_kv_cache(pool, allocator)
+        self.assertIs(req.retraction_backup.mamba_cpu, MAMBA_STATE)
+
+        req.load_kv_cache(pool, allocator)
+        self.assertIs(pool.mamba_pool.loaded, MAMBA_STATE)
+
+    def test_state_is_not_copied_twice_when_the_kv_pool_carries_it(self):
+        req, pool = _req_and_pool()
+        allocator = _Allocator(carries_mamba=True)
+
+        req.offload_kv_cache(pool, allocator)
+        self.assertIsNone(req.retraction_backup.mamba_cpu)
+
+        req.load_kv_cache(pool, allocator)
+        self.assertIsNone(pool.mamba_pool.loaded)
+
+
+if __name__ == "__main__":
+    unittest.main()