Fix snapshot bootstrap for short generations

python/sglang/srt/managers/scheduler.py

@@ -3187,6 +3187,8 @@ def _trigger_snapshot_hooks(self, batch: ScheduleBatch):
             # Skip if request doesn't have mamba_pool_idx
             if not hasattr(req, "mamba_pool_idx") or req.mamba_pool_idx is None:
                 continue
+            if getattr(req, "req_pool_idx", None) is None:
+                continue
 
             # Calculate turn number (approximate based on output length)
             turn_number = len(req.output_ids) if hasattr(req, "output_ids") else 0

python/sglang/srt/managers/scheduler_components/batch_result_processor.py

@@ -180,6 +180,39 @@ def _maybe_collect_customized_info(
                     elem = elem.copy()
                 req.customized_info[k].append(elem)
 
+    def _snapshot_finished_req_before_release(self, req: Req):
+        """Persist Mamba state while finished request pool slots are still live."""
+        if (
+            self.snapshot_hook_manager is None
+            or getattr(req, "mamba_pool_idx", None) is None
+        ):
+            return
+
+        if getattr(req, "req_pool_idx", None) is None:
+            logger.warning(
+                "Skipping finished-request snapshot for rid=%s: req_pool_idx is already released",
+                getattr(req, "rid", None),
+            )
+            return
+
+        mamba_pool = getattr(self.req_to_token_pool, "mamba_pool", None)
+        if mamba_pool is None:
+            return
+
+        fill_ids = getattr(req, "fill_ids", None)
+        turn_number = (
+            len(fill_ids)
+            if fill_ids is not None and hasattr(fill_ids, "__len__")
+            else len(req.output_ids)
+        )
+        self.snapshot_hook_manager.trigger_post_forward(
+            req=req,
+            mamba_pool=mamba_pool,
+            req_pool=self.req_to_token_pool,
+            turn_number=turn_number,
+            additional_context=None,
+        )
+
     def process_batch_result_prefill(
         self,
         batch: ScheduleBatch,
@@ -235,6 +268,7 @@ def process_batch_result_prefill(
                     if req.finished():
                         self._maybe_collect_routed_experts(req)
                         self._maybe_collect_indexer_topk(req)
+                        self._snapshot_finished_req_before_release(req)
                         release_kv_cache(req, self.tree_cache)
                         req.time_stats.set_completion_time()
                     elif not batch.decoding_reqs or req not in batch.decoding_reqs:
@@ -319,6 +353,7 @@ def process_batch_result_prefill(
                     req.update_finish_state()
 
                     if req.finished():
+                        self._snapshot_finished_req_before_release(req)
                         release_kv_cache(req, self.tree_cache)
                         req.time_stats.set_completion_time()
                     else:
@@ -788,26 +823,7 @@ def _handle_finished_req(
             # _trigger_snapshot_hooks (called later) fires after free_mamba_cache
             # has already set mamba_pool_idx = None, so it misses finished reqs.
             # We snapshot here while the state is still live in the pool.
-            if (
-                self.snapshot_hook_manager is not None
-                and hasattr(req, "mamba_pool_idx")
-                and req.mamba_pool_idx is not None
-            ):
-                _mamba_pool = getattr(self.req_to_token_pool, "mamba_pool", None)
-                if _mamba_pool is not None:
-                    _fill_ids = getattr(req, "fill_ids", None)
-                    _turn_number = (
-                        len(_fill_ids)
-                        if _fill_ids is not None and hasattr(_fill_ids, "__len__")
-                        else len(req.output_ids)
-                    )
-                    self.snapshot_hook_manager.trigger_post_forward(
-                        req=req,
-                        mamba_pool=_mamba_pool,
-                        req_pool=self.req_to_token_pool,
-                        turn_number=_turn_number,
-                        additional_context=None,
-                    )
+            self._snapshot_finished_req_before_release(req)
             # --- END ENGRAM ---
 
             if self.server_args.disaggregation_decode_enable_offload_kvcache:

python/sglang/srt/managers/scheduler_snapshot_handlers.py

@@ -395,9 +395,6 @@ def handle_save_snapshot(scheduler, recv_req):
 
     try:
         effective_conv_id = recv_req.conversation_id or recv_req.rid
-        effective_snapshot_id = (
-            recv_req.snapshot_id or f"{effective_conv_id}-t{recv_req.turn_number or 0}"
-        )
 
         # Find the request by rid
         req = scheduler._find_request_by_rid(recv_req.rid)
@@ -413,10 +410,21 @@ def handle_save_snapshot(scheduler, recv_req):
                 )
                 if warm_result is not None:
                     conv_states_w, temporal_states_w, meta_dict = warm_result
+                    warm_conversation_id = (
+                        recv_req.conversation_id
+                        or meta_dict.get("conversation_id")
+                        or effective_conv_id
+                    )
+                    warm_turn_number = recv_req.turn_number or meta_dict.get(
+                        "turn_number", 0
+                    )
+                    warm_snapshot_id = (
+                        recv_req.snapshot_id
+                        or f"{warm_conversation_id}-t{warm_turn_number or 0}"
+                    )
                     snap_meta = MambaSnapshotMetadata(
-                        conversation_id=effective_conv_id,
-                        turn_number=recv_req.turn_number
-                        or meta_dict.get("turn_number", 0),
+                        conversation_id=warm_conversation_id,
+                        turn_number=warm_turn_number,
                         branch_name=recv_req.branch_name,
                         timestamp=time.time(),
                         token_count=int(meta_dict.get("token_count", 0)),
@@ -434,12 +442,13 @@ def handle_save_snapshot(scheduler, recv_req):
                         conv_states_w, temporal_states_w, snap_meta
                     )
                     logger.info(
-                        f"Manual snapshot saved from WARM tier: conversation={effective_conv_id}, "
-                        f"snapshot_id={effective_snapshot_id}"
+                        "Manual snapshot saved from WARM tier: "
+                        f"conversation={warm_conversation_id}, "
+                        f"snapshot_id={warm_snapshot_id}",
                     )
                     return SaveSnapshotReqOutput(
                         success=True,
-                        snapshot_id=effective_snapshot_id,
+                        snapshot_id=warm_snapshot_id,
                         message="Snapshot saved successfully (from WARM tier)",
                     )
             return SaveSnapshotReqOutput(
@@ -471,13 +480,22 @@ def handle_save_snapshot(scheduler, recv_req):
         )
 
         # Build metadata
+        live_conversation_id = (
+            recv_req.conversation_id
+            or getattr(req, "conversation_id", None)
+            or recv_req.rid
+        )
+        live_snapshot_id = (
+            recv_req.snapshot_id
+            or f"{live_conversation_id}-t{recv_req.turn_number or 0}"
+        )
         layer_config = {
             "num_layers": mamba_pool.num_mamba_layers,
             "model_type": "hybrid" if hasattr(req, "mamba_pool_idx") else "mamba",
         }
 
         metadata = MambaSnapshotMetadata(
-            conversation_id=effective_conv_id,
+            conversation_id=live_conversation_id,
             turn_number=recv_req.turn_number,
             branch_name=recv_req.branch_name,
             timestamp=time.time(),
@@ -501,13 +519,13 @@ def handle_save_snapshot(scheduler, recv_req):
         scheduler.snapshot_manager.save_snapshot(conv_states, temporal_states, metadata)
 
         logger.info(
-            f"Manual snapshot saved: conversation={effective_conv_id}, "
-            f"snapshot_id={effective_snapshot_id}, turn={recv_req.turn_number}"
+            f"Manual snapshot saved: conversation={live_conversation_id}, "
+            f"snapshot_id={live_snapshot_id}, turn={recv_req.turn_number}"
         )
 
         return SaveSnapshotReqOutput(
             success=True,
-            snapshot_id=effective_snapshot_id,
+            snapshot_id=live_snapshot_id,
             message="Snapshot saved successfully",
         )
 

python/sglang/srt/managers/scheduler_snapshot_system.py

@@ -324,6 +324,11 @@ def post_forward_snapshot_callback(trigger):
                 )
 
                 if success:
+                    request_id = getattr(req, "rid", None)
+                    if request_id and request_id != conversation_id:
+                        scheduler.tier_manager.host_pool.alias_state(
+                            request_id, conversation_id
+                        )
                     scheduler.snapshot_policy.mark_snapshot_taken(conversation_id)
                     logger.debug(
                         f"Snapshot saved to WARM tier: conversation={conversation_id}, "

python/sglang/srt/snapshot/mamba_host_pool.py

@@ -105,6 +105,7 @@ def __init__(
 
         # Storage: OrderedDict maintains LRU order
         self._pool: OrderedDict[str, HostPoolEntry] = OrderedDict()
+        self._aliases: dict[str, str] = {}
 
         # Thread safety
         self._lock = threading.RLock()
@@ -182,6 +183,7 @@ def save_state(
             # Add entry (at end for LRU)
             self._pool[conversation_id] = entry
             self._current_memory_bytes += entry_size
+            self._aliases.pop(conversation_id, None)
 
             logger.debug(
                 f"Saved state to host pool: {conversation_id}, "
@@ -191,6 +193,30 @@ def save_state(
 
             return True
 
+    def resolve_key(self, conversation_id: str) -> str:
+        """Resolve a host-pool alias to its canonical conversation id."""
+        with self._lock:
+            return self._aliases.get(conversation_id, conversation_id)
+
+    def alias_state(self, alias_id: str, conversation_id: str) -> bool:
+        """Make an existing WARM state addressable by an additional id."""
+        with self._lock:
+            canonical_id = self._aliases.get(conversation_id, conversation_id)
+            if alias_id == canonical_id:
+                return canonical_id in self._pool
+            if canonical_id not in self._pool:
+                return False
+            if alias_id in self._pool:
+                logger.warning(
+                    "Cannot alias host-pool state: alias %s already has canonical state",
+                    alias_id,
+                )
+                return False
+
+            self._aliases[alias_id] = canonical_id
+            logger.debug("Aliased host-pool state: %s -> %s", alias_id, canonical_id)
+            return True
+
     def get_state(
         self, conversation_id: str
     ) -> Optional[Tuple[List[torch.Tensor], torch.Tensor, dict]]:
@@ -209,19 +235,20 @@ def get_state(
             Tuple of (conv_states, temporal_states, metadata) or None if not found
         """
         with self._lock:
-            if conversation_id not in self._pool:
+            canonical_id = self._aliases.get(conversation_id, conversation_id)
+            if canonical_id not in self._pool:
                 self._misses += 1
                 logger.debug(f"Host pool miss: {conversation_id}")
                 return None
 
-            entry = self._pool.pop(conversation_id)
+            entry = self._pool.pop(canonical_id)
 
             # Update access metadata
             entry.last_access_time = time.time()
             entry.access_count += 1
 
             # Move to end (most recently used)
-            self._pool[conversation_id] = entry
+            self._pool[canonical_id] = entry
 
             # Update metrics
             self._hits += 1
@@ -271,7 +298,7 @@ def get_state_for_reference(
     def has_state(self, conversation_id: str) -> bool:
         """Check if conversation state exists in host pool."""
         with self._lock:
-            return conversation_id in self._pool
+            return self._aliases.get(conversation_id, conversation_id) in self._pool
 
     def remove_state(self, conversation_id: str) -> bool:
         """
@@ -284,11 +311,20 @@ def remove_state(self, conversation_id: str) -> bool:
             True if removed, False if not found
         """
         with self._lock:
+            if conversation_id in self._aliases:
+                self._aliases.pop(conversation_id, None)
+                return True
+
             if conversation_id not in self._pool:
                 return False
 
             entry = self._pool.pop(conversation_id)
             self._current_memory_bytes -= entry.memory_bytes()
+            self._aliases = {
+                alias: canonical
+                for alias, canonical in self._aliases.items()
+                if canonical != conversation_id
+            }
 
             logger.debug(
                 f"Removed state from host pool: {conversation_id}, "
@@ -311,6 +347,11 @@ def _evict_lru(self) -> bool:
         conversation_id, entry = self._pool.popitem(last=False)
         self._current_memory_bytes -= entry.memory_bytes()
         self._evictions += 1
+        self._aliases = {
+            alias: canonical
+            for alias, canonical in self._aliases.items()
+            if canonical != conversation_id
+        }
 
         logger.info(
             f"Evicted from host pool (LRU): {conversation_id}, "
@@ -328,7 +369,8 @@ def list_conversations(self) -> List[str]:
     def get_conversation_metadata(self, conversation_id: str) -> Optional[dict]:
         """Get metadata for a conversation without accessing state."""
         with self._lock:
-            entry = self._pool.get(conversation_id)
+            canonical_id = self._aliases.get(conversation_id, conversation_id)
+            entry = self._pool.get(canonical_id)
             if entry:
                 return {
                     "conversation_id": entry.conversation_id,
@@ -369,6 +411,7 @@ def clear(self):
         """Clear all entries from host pool."""
         with self._lock:
             self._pool.clear()
+            self._aliases.clear()
             self._current_memory_bytes = 0
             logger.info("Host pool cleared")