From cc323848c9a022120c6fe0d45b8ca9f96f5aedf6 Mon Sep 17 00:00:00 2001 From: "Engram Session A (H100)" Date: Mon, 25 May 2026 04:06:59 +0000 Subject: [PATCH 1/2] Fix snapshot bootstrap for short generations --- python/sglang/srt/managers/scheduler.py | 2 + .../batch_result_processor.py | 56 ++++++++++++------- 2 files changed, 38 insertions(+), 20 deletions(-) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 5e8091326247..46c039b7b74c 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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 diff --git a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py index 35dd5370b664..981995df46c6 100644 --- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +++ b/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: From b15bba1406c9e8248f523521a9706eb091eb7717 Mon Sep 17 00:00:00 2001 From: "Engram Session A (H100)" Date: Mon, 25 May 2026 06:04:48 +0000 Subject: [PATCH 2/2] Fix WARM snapshot lookup by returned rid --- .../managers/scheduler_snapshot_handlers.py | 44 ++++++++++----- .../srt/managers/scheduler_snapshot_system.py | 5 ++ python/sglang/srt/snapshot/mamba_host_pool.py | 53 +++++++++++++++++-- 3 files changed, 84 insertions(+), 18 deletions(-) diff --git a/python/sglang/srt/managers/scheduler_snapshot_handlers.py b/python/sglang/srt/managers/scheduler_snapshot_handlers.py index bc92d60d6ff2..380a74f164c2 100644 --- a/python/sglang/srt/managers/scheduler_snapshot_handlers.py +++ b/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", ) diff --git a/python/sglang/srt/managers/scheduler_snapshot_system.py b/python/sglang/srt/managers/scheduler_snapshot_system.py index 854cd57efaa2..bae40a4373bb 100644 --- a/python/sglang/srt/managers/scheduler_snapshot_system.py +++ b/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}, " diff --git a/python/sglang/srt/snapshot/mamba_host_pool.py b/python/sglang/srt/snapshot/mamba_host_pool.py index b82ca6dc270e..b2c81bee984a 100644 --- a/python/sglang/srt/snapshot/mamba_host_pool.py +++ b/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")