[UnifiedTree]: Sync Replay SSM (#30636)

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