From 2286e25a213670aceef50fb82393a887a8efcd28 Mon Sep 17 00:00:00 2001 From: Zhangheng Date: Fri, 10 Jul 2026 21:54:59 +0800 Subject: [PATCH] [UnifiedTree]: Sync Replay SSM (#30636) --- .../mamba_component.py | 23 +++++++++++++++---- 1 file changed, 18 insertions(+), 5 deletions(-) diff --git a/python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py b/python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py index dae605f35..1828f2974 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py +++ b/python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py @@ -301,11 +301,24 @@ class MambaComponent(TreeComponent): token_ids_len: int, is_finished: bool, ) -> Optional[int]: - cache_len = ( - req.mamba_last_track_seqlen - if self.enable_mamba_extra_buffer - else token_ids_len - ) + if self.enable_mamba_extra_buffer: + cache_len = req.mamba_last_track_seqlen + else: + cache_len = token_ids_len + # ReplaySSM (no_buffer): `temporal[slot]` lags the live state by the + # slot's unflushed ring depth (`write_pos`), so on request finish cap + # the donate to the last flush boundary (where temporal is current) + # and reset the cursor, keeping the donated checkpoint consistent with + # its key length. page_size is asserted == 1, so no realign. Mirrors + # MambaRadixCache.cache_finished_req. + if is_finished: + write_pos_buf = ( + self.cache.req_to_token_pool.mamba_pool.replayssm_write_pos + ) + if write_pos_buf is not None: + cache_len -= int(write_pos_buf[req.mamba_pool_idx].item()) + write_pos_buf[req.mamba_pool_idx] = 0 + if is_finished: if cache_len is None: cache_len = 0