diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 56c8ebd6d..22b8e1dde 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/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, diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 1ece6720f..8b92cfe48 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/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: diff --git a/python/sglang/srt/mem_cache/kv_cache_builder.py b/python/sglang/srt/mem_cache/kv_cache_builder.py index 7b6b0d9b6..7f36dab62 100644 --- a/python/sglang/srt/mem_cache/kv_cache_builder.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 9cb0fb363..20ec41b0b 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/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, diff --git a/test/registered/unit/mem_cache/test_retraction_mamba_backup.py b/test/registered/unit/mem_cache/test_retraction_mamba_backup.py new file mode 100644 index 000000000..7b678f35a --- /dev/null +++ b/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()