From 212b2a39df80c5074abac87c84833da0dac7202c Mon Sep 17 00:00:00 2001 From: Mike Date: Fri, 14 Aug 2026 15:02:13 -0700 Subject: [PATCH 01/12] fix(gemma4): load vision tower unquantized under compressed-tensors MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Port of upstream PR sgl-project/sglang#33988 (gemma4_mm.py hunk only; the marlin shape-check renames in that PR do not apply to v0.5.17). compressed-tensors Gemma-4 checkpoints (e.g. RedHatAI/gemma-4-26B-A4B-it-FP8-dynamic) exclude the whole vision tower via the quantization ignore list, but the entries carry the checkpoint wrapper's '.linear' suffix and unfused q/k/v names, which never match SGLang's fused modules — so the bf16 vision tower was loaded through fp8 schemes with garbage scales and produced NaN image features. Observed blast radius on v0.5.17: the default multimodal warmup request runs the broken vision path, its NaN KV is cached at the shared radix prefix, and every subsequent request greedy-decodes to . This also makes --skip-server-warmup insufficient protection for a prod deployment: any real image request re-poisons the cache. Upstream issue: sgl-project/sglang#24927 (closed, but the fix PR is still unmerged as of 2026-08-14). Co-Authored-By: Claude Fable 5 (cherry picked from commit 29a730878ae07cd81ae5504c18a5bc2bb1efa6d0) --- python/sglang/srt/models/gemma4_mm.py | 35 ++++++++++++++++++++++++++- 1 file changed, 34 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/models/gemma4_mm.py b/python/sglang/srt/models/gemma4_mm.py index 56d033ff4b63..9c092df12ea2 100644 --- a/python/sglang/srt/models/gemma4_mm.py +++ b/python/sglang/srt/models/gemma4_mm.py @@ -173,6 +173,38 @@ class Gemma4ForConditionalGeneration(PreTrainedModel): embedding_padding_modules = [] supports_lora = True + @staticmethod + def _vision_tower_quant_config( + quant_config: Optional[QuantizationConfig], + ) -> Optional[QuantizationConfig]: + """Quantization config for the Gemma-4 vision tower. + + compressed-tensors Gemma-4 checkpoints exclude every vision-tower + linear through `ignore`, but the entries carry the `.linear` suffix of + the checkpoint's clip wrapper, which SGLang's fused `qkv_proj` / + `gate_up_proj` do not have, so the match fails and the vision tower + gets quantized anyway. Drop the config for the vision tower to honour + what the checkpoint asked for. + + CAI port of upstream PR sgl-project/sglang#33988. Without this, the + (bf16-on-disk) vision tower runs through fp8 kernels with garbage + scales and produces NaN image features; the multimodal warmup request + then poisons the radix cache at the shared prefix and every + subsequent request returns tokens. + """ + from sglang.srt.layers.quantization.compressed_tensors.compressed_tensors import ( + CompressedTensorsConfig, + ) + + if not isinstance(quant_config, CompressedTensorsConfig): + return quant_config + + logger.warning_once( + "Gemma-4 vision tower is loaded unquantized; the compressed-tensors " + "config does not apply to it." + ) + return None + def __init__( self, config: Gemma4Config, @@ -191,9 +223,10 @@ def __init__( # Vision/audio encoders + their projection embedders are only consumed # at the input-embedding stage, so they live on the first PP rank only. if self.pp_group.is_first_rank: + vision_tower_quant_config = self._vision_tower_quant_config(quant_config) self.vision_tower = Gemma4VisionEncoder( config=config.vision_config, - quant_config=quant_config, + quant_config=vision_tower_quant_config, prefix=add_prefix("vision_tower", prefix), ) self.embed_vision = Gemma4MultimodalEmbedder( From e036402b90d1eb5af031973d90c6719bae8cf6de Mon Sep 17 00:00:00 2001 From: Mike Date: Fri, 14 Aug 2026 20:03:20 -0700 Subject: [PATCH 02/12] fix(triton): dequantize shared-layer KV pulled from a quantized cache Port of unmerged upstream PR sgl-project/sglang#22615 (issue #22277). Gemma-4 KV-sharing layers call attention with k=None/v=None and the triton backend reads K/V straight out of the KV cache. With --kv-cache-dtype fp8_e4m3 those tensors arrive as fp8 and the extend kernel dies compiling tl.dot(bf16_q, fp8_k) ('Unsupported rhs dtype'). Convert to the query dtype and apply k/v scales, mirroring what the non-shared path does at store time. Validated on MI325x with RedHatAI/gemma-4-26B-A4B-it-FP8-dynamic: fp8 KV doubles the pool (2.20M full + 1.76M SWA tokens vs 1.11M/891k bf16) and improves the shared-prefix bench to 31.4/33.6 req/s (NP=512/2048) vs 28.6/29.9 with bf16 KV; greedy quality unchanged. Co-Authored-By: Claude Fable 5 (cherry picked from commit 9294fcc510d41c4e484c8535fd47d8c0b9e230d4) --- python/sglang/srt/layers/attention/triton_backend.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/python/sglang/srt/layers/attention/triton_backend.py b/python/sglang/srt/layers/attention/triton_backend.py index c378c80ab0f8..1fb91f398bfe 100644 --- a/python/sglang/srt/layers/attention/triton_backend.py +++ b/python/sglang/srt/layers/attention/triton_backend.py @@ -1279,6 +1279,17 @@ def forward_extend( k_buffer, v_buffer = pool.get_kv_buffer(layer.layer_id) k = k_buffer[cache_loc] v = v_buffer[cache_loc] + # CAI (upstream PR #22615): KV-sharing layers pull K/V straight from + # the cache; with a quantized KV cache these arrive as fp8 and the + # extend kernel's tl.dot(bf16_q, fp8_k) fails to compile. Dequantize + # to the compute dtype (issue #22277). + if k.dtype != q.dtype: + k = k.to(q.dtype) + v = v.to(q.dtype) + if layer.k_scale_float is not None: + k.mul_(layer.k_scale_float) + if layer.v_scale_float is not None: + v.mul_(layer.v_scale_float) elif k is None or v is None: raise ValueError("Both k and v should be None or not None") else: From 3b7d2b805f76214f9ae302a29ebc9cb022847728 Mon Sep 17 00:00:00 2001 From: Jialin Ouyang Date: Wed, 5 Aug 2026 11:39:28 -0700 Subject: [PATCH 03/12] [Unified Radix Cache] Complete the tree-core interface boundary (#33580) (cherry picked from commit 301a8a09c5e8ab6575a76efb0a55881888044118) --- .../unified_cache/components/README.md | 12 +++++------ .../components/mamba_component.py | 9 ++++---- .../unified_cache/unified_tree_core.py | 21 ++++++++++++++++--- .../unified_tree_core_interface.py | 21 +++++++++++++++++++ .../srt/mem_cache/unified_radix_cache.py | 7 +++++-- 5 files changed, 54 insertions(+), 16 deletions(-) diff --git a/python/sglang/srt/mem_cache/unified_cache/components/README.md b/python/sglang/srt/mem_cache/unified_cache/components/README.md index 336e4d1759ab..14f2f15c1d53 100644 --- a/python/sglang/srt/mem_cache/unified_cache/components/README.md +++ b/python/sglang/srt/mem_cache/unified_cache/components/README.md @@ -111,7 +111,7 @@ Find the longest cached prefix for a token sequence. - Promotes matched path to MRU in each component's LRU via `node_has_component_data()` as filter - Updates `last_access_time` with decreasing timestamps up the path (parent < child) - Concatenates matched device indices via `torch.cat` (concat length ≤ K, subsumed by O(K)) - - Calls `finalize_match_result_in_tree_core()` per component (tree-side: Full/SWA host-hit sums, Mamba `branching_seqlen`); the cache then routes the static `finalize_match_result_in_cache()` per component post-walk (Mamba performs copy-on-write: allocates new pool slot, copies SSM state) + - Calls `finalize_match_result_in_tree_core()` per component (tree-side: Full/SWA host-hit sums, Mamba `branching_seqlen`); the cache then routes `finalize_match_result_in_cache()` per component post-walk (Mamba performs copy-on-write: allocates new pool slot, copies SSM state) --- @@ -127,7 +127,7 @@ Insert a key-value pair into the tree. | **Mutation** | Creates new leaf nodes; updates component data on overlapping nodes; frees duplicate KV indices; may split nodes; updates LRU lists and evictable sizes | | **Complexity** | **O(K + D·C)** | -**Algorithm detail** (`_insert_helper`): +**Algorithm detail** (the resumable insert steps: `_insert_walk_step` / `_insert_commit_step` / `_insert_tail_step`): 1. At each existing node, calls `_touch_node` → promotes to MRU via `node_has_component_data()` 2. If key diverges mid-node, calls `_split_node` → `redistribute_on_node_split()` per component 3. For each overlapping node, calls `update_component_on_insert_overlap()` per component — returns `consumed_from` index; the tree frees `value[dup_start:consumed_from]` as duplicate pool indices @@ -267,15 +267,15 @@ Each component implements these hooks. See `tree_component.py` for the ABC and d |------|---------|-----------|----------| | `create_match_validator(match_device_only=False)` | Return a per-match stateful predicate that decides whether a node is a valid match boundary. Full: requires Full device data, or host backup when `match_device_only=False`. SWA: tracks accumulated window length across device/host data. Mamba: requires Mamba device data, or host backup when `match_device_only=False`. | `_match_prefix_helper` | *abstract* | | `finalize_match_result_in_tree_core()` | Tree-side post-processing inside the match walk. Full/SWA: host-hit sums. Mamba: records `branching_seqlen` + the host-hit bump. | `_match_post_processor` | pass-through | -| `finalize_match_result_in_cache()` | Static, cache-level finalize after the walk (receives the cache + NodeId-based result), dispatched class-level by `UnifiedRadixCache.match_prefix`. Mamba: copy-on-write — allocates a new mamba pool slot, copies SSM state into the request pool. | `UnifiedRadixCache.match_prefix` | pass-through | +| `finalize_match_result_in_cache()` | Cache-level finalize after the walk (receives the params + NodeId-based result), dispatched by `UnifiedRadixCache.match_prefix`. Mamba: copy-on-write — allocates a new mamba pool slot, copies SSM state into the request pool. | `UnifiedRadixCache.match_prefix` | pass-through | ### Insert Phase | Hook | Purpose | Called By | Default | |------|---------|-----------|----------| -| `update_component_on_insert_overlap()` | Handle key overlap with an existing node during insert. Returns the index within `value_slice` from which this component consumed (took ownership of) pool slots. Full/Mamba: no consumption (`prefix_len`). SWA: may recover tombstoned nodes within the sliding window boundary. | `_insert_helper` | returns `prefix_len` | -| `recover_after_unevict()` | Rebuild auxiliary component data after `_unevict_node_on_insert()` restores a Full device value from fresh KV indices. SWA uses this to rebuild in-window SWA data. | `_insert_helper` | no-op | -| `commit_insert_component_data()` | Finalize component data on the target node after the insert walk completes. Full: no-op (handled by `_add_new_node`). SWA: checks window boundary, may split node — parent becomes tombstone, child gets SWA data. Mamba: sets mamba pool indices and inserts into Mamba LRU. | `_insert_helper` | no-op | +| `update_component_on_insert_overlap()` | Handle key overlap with an existing node during insert. Returns the index within `value_slice` from which this component consumed (took ownership of) pool slots. Full/Mamba: no consumption (`prefix_len`). SWA: may recover tombstoned nodes within the sliding window boundary. | `_insert_walk_step` | returns `prefix_len` | +| `recover_after_unevict()` | Rebuild auxiliary component data after `_unevict_node_on_insert()` restores a Full device value from fresh KV indices. SWA uses this to rebuild in-window SWA data. | `_insert_walk_step` | no-op | +| `commit_insert_component_data()` | Finalize component data on the target node after the insert walk completes. Full: no-op (handled by `_add_new_node`). SWA: checks window boundary, may split node — parent becomes tombstone, child gets SWA data. Mamba: sets mamba pool indices and inserts into Mamba LRU. | `_insert_commit_step` | no-op | ### Node Split 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 a4efe3e4a78d..bf46238fcf67 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 @@ -890,12 +890,11 @@ def apply_component_action(self, action: ComponentAction) -> None: if isinstance(action, MambaEvictExcessPathStates): device_frees: dict[ComponentType, list[torch.Tensor]] = defaultdict(list) host_frees: dict[ComponentType, list[torch.Tensor]] = defaultdict(list) - # Drain even if the walk raises so tombstoned slots are not leaked. + # Drain even if the walk raises so tombstoned slots are not leaked; + # the walk runs behind the tree-core interface (Rust runs it natively). try: - self._evict_excess_path_states( - self.tree_core.node_by_id(action.tail_node_id), - device_frees, - host_frees, + self.tree_core.evict_excess_path_states( + action.tail_node_id, device_frees, host_frees ) finally: self.cache._free_values(device_frees, host_frees) diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py index e0110067cb0e..8627190a58fa 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py +++ b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py @@ -372,9 +372,6 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface): per-component LRUs, the size/leaf bookkeeping, and the component drivers, plus ``reset()``. - TODO(Jialin): the tree operations still live on ``UnifiedRadixCache`` and - reach this state through its proxy properties; they migrate onto this class - as the TreeCore split completes. """ def __init__( @@ -497,6 +494,14 @@ def get_prefix_hash_values(self, node_id: NodeId) -> list[str]: node = self.node_by_id(node_id) return node.get_prefix_hash_values(node.parent) + def get_hash_values(self, node_id: NodeId) -> list[str]: + """The hash values owned by this node, excluding its ancestors.""" + return self.node_by_id(node_id).hash_value or [] + + def root_node_handle(self, extra_key: Optional[str] = None) -> NodeId: + """The NodeId anchoring matches; the single root serves every namespace.""" + return self.root_node.id + def _new_node(self, priority: int = 0) -> UnifiedTreeNode: """Create and register a tree node in the arena.""" node = UnifiedTreeNode(self.component_types, priority=priority) @@ -1279,6 +1284,16 @@ def drive_host_eviction( ) return result + def evict_excess_path_states( + self, + tail_node_id: NodeId, + device_frees: dict[ComponentType, list[torch.Tensor]], + host_frees: dict[ComponentType, list[torch.Tensor]], + ) -> None: + self.components_by_type[ComponentType.MAMBA]._evict_excess_path_states( + self.node_by_id(tail_node_id), device_frees, host_frees + ) + def _evict_host_leaf( self, node: UnifiedTreeNode, diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py index cbd9910ffbed..2e16983880da 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py +++ b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py @@ -165,6 +165,16 @@ def get_prefix_hash_values(self, node_id: NodeId) -> list[str]: """The hash chain of the node's ancestors, in root-to-parent order.""" ... + @abstractmethod + def get_hash_values(self, node_id: NodeId) -> list[str]: + """The hash values owned by this node, excluding its ancestors.""" + ... + + @abstractmethod + def root_node_handle(self, extra_key: Optional[str] = None) -> NodeId: + """The NodeId anchoring matches for the namespace.""" + ... + @abstractmethod def inc_lock_ref( self, node_id: NodeId, skip_lock_components: Sequence[ComponentType] = () @@ -345,6 +355,17 @@ def drive_host_eviction( """Evict a component's host-side resources; no-op if the component is absent.""" ... + @abstractmethod + def evict_excess_path_states( + self, + tail_node_id: NodeId, + device_frees: dict[ComponentType, list[torch.Tensor]], + host_frees: dict[ComponentType, list[torch.Tensor]], + ) -> None: + """Evict shallow Mamba device checkpoints beyond the per-path cap on the + tail's root path, collecting freed values into the caller's dicts.""" + ... + # ==== HiCache ==== @abstractmethod diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index c11a02e30f68..d4362382658a 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -213,7 +213,10 @@ def __init__( self.hicache_storage_pass_prefix_keys = False self.reset() - logger.info(f"Init Unified RadixTree with components {self.tree_components}") + logger.info( + f"Init Unified Radix Cache. Components: {self.tree_components}. " + f"Tree Core: {type(self.tree_core).__name__}" + ) def _all_reduce_attn_groups(self, tensor: torch.Tensor, op): reduced = False @@ -2191,4 +2194,4 @@ def resolve_node_handle(self, node_handle): def root_node_handle(self, extra_key: Optional[str] = None) -> NodeId: """The root's NodeId -- URC match results carry NodeIds.""" - return self.tree_core.root_node.id + return self.tree_core.root_node_handle(extra_key) From 0789a8b210388bb28e717b8bb54b8ecfd44240aa Mon Sep 17 00:00:00 2001 From: Zhiqiang Xie Date: Fri, 7 Aug 2026 15:59:59 -0700 Subject: [PATCH 04/12] [HiCache] write_back: reclaim duplicated host copy first under host pressure (#33777) (cherry picked from commit f4f91ef254608f64102fb8cddfcc7a55bffc6c13) --- .../unified_cache/unified_tree_core.py | 181 +++++++++++++++++- .../unified_tree_core_interface.py | 9 + .../srt/mem_cache/unified_radix_cache.py | 28 ++- .../test_unified_radix_cache_unittest.py | 2 + 4 files changed, 209 insertions(+), 11 deletions(-) diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py index 8627190a58fa..d1acd88b1e0e 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py +++ b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py @@ -84,6 +84,11 @@ logger = logging.getLogger(__name__) +# 42 bits: digest * 1000003 (< 2^20) stays under 2^62, so the update never +# overflows int64 with plain (non-wrapping) arithmetic in the Rust port, and +# the TP consistency check can still all_reduce [digest, -digest] in int64. +_RECLAIM_DIGEST_MASK = (1 << 42) - 1 + class StorageBackupSpec(NamedTuple): """A node's device->storage backup spec, gathered tree-side.""" @@ -123,6 +128,9 @@ def __init__(self, tree_components: tuple[ComponentType, ...], priority: int = 0 self.id = UnifiedTreeNode.counter UnifiedTreeNode.counter += 1 self.write_through_pending_id: Optional[int] = None + # Anchor NodeId of an in-flight H->D load-back reading this node's + # host slots; such host copies must not be reclaimed until the ack. + self.load_back_pending_id: Optional[int] = None def component(self, component_type: ComponentType) -> ComponentData: return self.component_data[component_type] @@ -457,6 +465,11 @@ def reset(self) -> None: ) for ct in self.component_types } + # Full KV on both tiers -> redundant host copy, reclaimed first by + # write_back; insertion-ordered dict keeps victims TP-deterministic. + self.full_host_duplicates: dict[NodeId, UnifiedTreeNode] = {} + # Rolling digest of reclaim victim ids, cross-checked across TP ranks. + self.write_back_duplicate_reclaim_digest: int = 0 self._empty_match_result = MatchResult( device_indices=torch.empty( @@ -1021,6 +1034,8 @@ def _split_node( new_node.key = child.key[:split_len] new_node.hit_count = child.hit_count new_node.creation_time = child.creation_time + # Split fragments stay on the anchor's root path for the ack's walk. + new_node.load_back_pending_id = child.load_back_pending_id self._for_each_component_lru(child, UnifiedLRUList.remove_node) @@ -1056,6 +1071,8 @@ def _split_node( self._update_evictable_leaf_sets(new_node) self._update_evictable_leaf_sets(child) + # Only the new fragment needs qualifying; the child keeps its id. + self._update_duplicate_tracking(new_node) return new_node, action def _add_new_node( @@ -1091,6 +1108,8 @@ def _unevict_node_on_insert( cd.value = fresh_value.clone() self.component_evictable_size_[ct] += n self._update_evictable_leaf_sets(node) + # A backuped node restored from fresh KV is a duplicate right away. + self._update_duplicate_tracking(node) if node.parent is not None: self._update_evictable_leaf_sets(node.parent) self._record_store_event(node, medium=StorageMedium.GPU) @@ -1107,6 +1126,26 @@ def _update_evictable_leaf_sets(self, node: UnifiedTreeNode) -> None: else: self.evictable_host_leaves.discard(node) + def _update_duplicate_tracking(self, node: UnifiedTreeNode) -> None: + """Register where duplicates are born (acks, split, unevict); + deregistration is lazy, so entries may be stale and re-checked live.""" + if self._is_settled_full_host_duplicate(node): + self.full_host_duplicates.setdefault(node.id, node) + else: + self.full_host_duplicates.pop(node.id, None) + + def _is_settled_full_host_duplicate(self, node: UnifiedTreeNode) -> bool: + """Full KV present on both tiers with no in-flight DMA on the node's + host slots; mid-transfer nodes join the tracking at their ack.""" + cd = node.component_data[BASE_COMPONENT_TYPE] + return ( + node is not self.root_node + and cd.value is not None + and cd.host_value is not None + and node.write_through_pending_id is None + and node.load_back_pending_id is None + ) + def _for_each_component_lru( self, node: UnifiedTreeNode, @@ -1272,10 +1311,18 @@ def _delete_unbacked_device_leaf( def drive_host_eviction( self, component_type: ComponentType, num_tokens: int ) -> DriveHostEvictionResult: - """Evict a component's host-side resources; no-op if the component is absent.""" + """Evict a component's host-side resources; no-op if absent. Under + write_back, FULL pressure reclaims redundant Full host copies first.""" result = DriveHostEvictionResult() comp = self.components_by_type.get(component_type) if comp is not None: + if self.is_write_back and component_type == BASE_COMPONENT_TYPE: + self._reclaim_full_host_duplicates( + num_tokens, + result.tracker, + result.device_frees, + result.host_frees, + ) comp.drive_host_eviction( num_tokens, result.tracker, @@ -1294,6 +1341,74 @@ def evict_excess_path_states( self.node_by_id(tail_node_id), device_frees, host_frees ) + def _reclaim_full_host_duplicates( + self, + num_tokens: int, + tracker: dict[ComponentType, int], + device_frees: dict[ComponentType, list[torch.Tensor]], + host_frees: dict[ComponentType, list[torch.Tensor]], + ) -> None: + """Reclaim Full host duplicates until num_tokens are freed; pass 1 + spares evictable D-leaves (imminent free demotes), pass 2 takes them.""" + swept_ids: list[NodeId] = [] + for spare_imminent_demotes in (True, False): + if tracker[BASE_COMPONENT_TYPE] >= num_tokens: + break + for node in self.full_host_duplicates.values(): + if tracker[BASE_COMPONENT_TYPE] >= num_tokens: + break + cd = node.component_data[BASE_COMPONENT_TYPE] + if cd.value is None or cd.host_value is None: + swept_ids.append(node.id) # stale entry + continue + if spare_imminent_demotes and node in self.evictable_device_leaves: + continue + if not self._can_reclaim_full_host_duplicate(node): + continue + self._release_full_host_duplicate( + node, tracker, device_frees, host_frees + ) + swept_ids.append(node.id) # released -> no longer a duplicate + # Sweep after the walk: the dict must not be mutated mid-iteration. + for nid in swept_ids: + self.full_host_duplicates.pop(nid, None) + + def _can_reclaim_full_host_duplicate(self, node: UnifiedTreeNode) -> bool: + """Full on both tiers, no in-flight DMA, no Full host lock; checked + live because tracking may be stale.""" + cd = node.component_data[BASE_COMPONENT_TYPE] + if node is self.root_node or cd.value is None or cd.host_value is None: + return False + if ( + node.write_through_pending_id is not None + or node.load_back_pending_id is not None + ): + return False + return cd.host_lock_ref == 0 + + def _release_full_host_duplicate( + self, + node: UnifiedTreeNode, + tracker: dict[ComponentType, int], + device_frees: dict[ComponentType, list[torch.Tensor]], + host_frees: dict[ComponentType, list[torch.Tensor]], + ) -> None: + """Free only the Full host layer; aux host slices stay under their own + pools' LRU (a host-only aux slice may be a sole copy).""" + assert self._can_reclaim_full_host_duplicate(node) + self._record_remove_event(node, medium=StorageMedium.CPU) + self._evict_component_and_detach_lru( + node, + self.components_by_type[BASE_COMPONENT_TYPE], + target=EvictLayer.HOST, + tracker=tracker, + device_frees=device_frees, + host_frees=host_frees, + ) + self.write_back_duplicate_reclaim_digest = ( + self.write_back_duplicate_reclaim_digest * 1000003 + node.id + 1 + ) & _RECLAIM_DIGEST_MASK + def _evict_host_leaf( self, node: UnifiedTreeNode, @@ -1435,6 +1550,8 @@ def _remove_leaf_from_parent(self, node: UnifiedTreeNode): key = node.key.child_key(self.page_size) v = node.parent.children.pop(key, None) assert v == node + # Deleted nodes must not linger in duplicate tracking as ghosts. + self.full_host_duplicates.pop(node.id, None) self._unregister_node(node) def _evict_component_and_detach_lru( @@ -1552,7 +1669,8 @@ def _is_host_leaf(self, node: UnifiedTreeNode) -> bool: """H-leaf: evicted, Full host value present, no children, unlocked, not root. Only the Full (base) component host_value is required; auxiliary - components are not mandatory for H-leaf membership.""" + components are not mandatory for H-leaf membership. In-flight DMA + marks need no check: marked nodes are never ``evicted``.""" if node is self.root_node or not node.evicted: return False if not node.backuped: @@ -1803,6 +1921,19 @@ def commit_load_back( rebuild is deferred to the orchestration layer.""" node = self.node_by_id(node_id) cache_actions: list[CacheAction | ComponentAction] = [] + # Pin every node whose host slots the in-flight DMA reads (including + # aux-only nodes) against reclaim until the ack. + for xfers in ([kv_xfer], *comp_xfers.values()): + for xfer in xfers: + for nid in xfer.nodes_to_load or (): + pinned = self.node_by_id(nid) + # One live load-back per node; only the same anchor may + # re-pin (a node can sit in Full and aux transfer lists). + assert pinned.load_back_pending_id in (None, node_id), ( + f"node {nid} pinned by load-back " + f"{pinned.load_back_pending_id}, new anchor {node_id}" + ) + pinned.load_back_pending_id = node_id kv_xfer.device_indices = device_indices self.components_by_type[BASE_COMPONENT_TYPE].commit_hicache_transfer( node, @@ -1811,8 +1942,7 @@ def commit_load_back( cache_actions=cache_actions, ) for nid in kv_xfer.nodes_to_load or (): - loaded = self.node_by_id(nid) - self._record_store_event(loaded, medium=StorageMedium.GPU) + self._record_store_event(self.node_by_id(nid), medium=StorageMedium.GPU) for ct, xfers in comp_xfers.items(): self.components_by_type[ct].commit_hicache_transfer( node, @@ -1823,6 +1953,17 @@ def commit_load_back( self._update_evictable_leaf_sets(node) return cache_actions + def finish_load_back(self, anchor_node_id: NodeId) -> None: + """Clear the in-flight H->D marks along the anchor's root path at ack + time; split fragments stay on the path, so the walk covers them.""" + node = self.node_by_id(anchor_node_id) + while node is not None and node is not self.root_node: + if node.load_back_pending_id == anchor_node_id: + node.load_back_pending_id = None + # The loaded copies become tracked duplicates only now. + self._update_duplicate_tracking(node) + node = node.parent + def mark_write_through_pending(self, node_id: NodeId) -> None: """Mark a node as having an in-flight write-through backup.""" node = self.node_by_id(node_id) @@ -1835,6 +1976,8 @@ def finish_write_through(self, node_ids: list[NodeId], ack_id: int) -> None: node = self.node_by_id(node_id) if node.write_through_pending_id == ack_id: node.write_through_pending_id = None + # The backed-up copy becomes a tracked duplicate only now. + self._update_duplicate_tracking(node) self._record_store_event(node, medium=StorageMedium.CPU) def set_component_device_value( @@ -1902,6 +2045,7 @@ def sanity_check( # ── PART 2: Per-node state machine and leaf qualification ── expected_dev_leaves: set[UnifiedTreeNode] = set() expected_hst_leaves: set[UnifiedTreeNode] = set() + expected_duplicates: set[UnifiedTreeNode] = set() for node in all_nodes: if node is self.root_node: @@ -1918,7 +2062,10 @@ def sanity_check( if cd.value is not None and not full_dev: E(f"node {nid} {ct} device present but Full.value=None") if cd.host_value is not None and not full_hst: - E(f"node {nid} {ct} host present but Full.host_value=None") + # write_back reclaim takes only the Full host layer; an + # aux host slice may outlive it while Full device is live. + if not (self.is_write_back and full_dev): + E(f"node {nid} {ct} host present but Full.host_value=None") # Every node must keep Full data on at least one layer. if not full_dev and not full_hst: @@ -1951,6 +2098,8 @@ def sanity_check( expected_dev_leaves.add(node) if self._is_host_leaf(node): expected_hst_leaves.add(node) + if self._is_settled_full_host_duplicate(node): + expected_duplicates.add(node) # ── PART 3: Tracking structures ── @@ -1972,6 +2121,16 @@ def sanity_check( if missing: E(f"H-leaf missing: {[n.id for n in list(missing)[:5]]}") + # Lazy deregistration: stale extras are legal; settled duplicates must + # be tracked and entries must not outlive their node. + expected_ids = {n.id for n in expected_duplicates} + dup_ids = set(self.full_host_duplicates.keys()) + if expected_ids - dup_ids: + E(f"Duplicate missing: {list(expected_ids - dup_ids)[:5]}") + ghost_ids = dup_ids - {n.id for n in all_nodes} + if ghost_ids: + E(f"Duplicate ghosts: {list(ghost_ids)[:5]}") + # D-leaf ∩ H-leaf = ∅ overlap = self.evictable_device_leaves & self.evictable_host_leaves if overlap: @@ -2083,6 +2242,18 @@ def sanity_check( E( f"[Ongoing] load_back node {nid} lock_ref={n.component_data[FCT].lock_ref}" ) + # Every in-flight H->D mark must belong to a live load-back; a leaked + # mark would pin the node's host copy against reclaim forever. + ongoing_load_ids = {node_id for _, node_id in ongoing_load_back} + for node in all_nodes: + if ( + node.load_back_pending_id is not None + and node.load_back_pending_id not in ongoing_load_ids + ): + E( + f"[Ongoing] node {node.id} load_back_pending_id=" + f"{node.load_back_pending_id} has no live load-back" + ) if errors: msg = ( diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py index 2e16983880da..8e7e20e20e7d 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py +++ b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py @@ -460,6 +460,15 @@ def commit_load_back( """Commit a successful H->D load-back onto the node; returns any cache actions.""" ... + @abstractmethod + def finish_load_back(self, anchor_node_id: NodeId) -> None: + """Clear the in-flight H->D marks on the anchor's root path at ack time.""" + ... + + # Order-sensitive digest of write_back duplicate-reclaim victim ids, + # cross-checked across TP ranks; cores that never reclaim keep 0. + write_back_duplicate_reclaim_digest: int = 0 + @abstractmethod def mark_write_through_pending(self, node_id: NodeId) -> None: """Mark a node as having an in-flight write-through backup.""" diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index d4362382658a..f01b3c4f37b8 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -1797,22 +1797,30 @@ def _sync_hicache_ready_counts( else () ) + # Piggybacked TP check: [digest, -digest] MIN-reduces to [min, -max], + # equal iff reclaim victim order matched on every rank. + digest = self.tree_core.write_back_duplicate_reclaim_digest ready_counts = torch.tensor( [ write_acks, load_acks, *storage_queue_sizes, + digest, + -digest, ], - dtype=torch.int, + dtype=torch.int64, device="cpu", ) self._all_reduce(ready_counts, torch.distributed.ReduceOp.MIN) count_values = list(map(int, ready_counts.tolist())) + assert ( + count_values[-2] == -count_values[-1] + ), "write_back duplicate-reclaim victims diverged across TP ranks" return ( count_values[0], count_values[1], - tuple(count_values[2:]), + tuple(count_values[2:-2]), extra_pool_names, ) @@ -1867,11 +1875,17 @@ def loading_check(self, finish_count: Optional[int] = None) -> None: finish_count = 0 if self.pp_rank == 0: finish_count = self._count_ready_acks(cc.ack_load_queue) - finish_count_tensor = torch.tensor( - finish_count, dtype=torch.int, device="cpu" + # Piggybacked TP check: [digest, -digest] MIN-reduces to [min, -max], + # equal iff reclaim victim order matched on every rank. + digest = self.tree_core.write_back_duplicate_reclaim_digest + sync_tensor = torch.tensor( + [finish_count, digest, -digest], dtype=torch.int64, device="cpu" ) - self._all_reduce(finish_count_tensor, torch.distributed.ReduceOp.MIN) - finish_count = finish_count_tensor.item() + self._all_reduce(sync_tensor, torch.distributed.ReduceOp.MIN) + finish_count = int(sync_tensor[0].item()) + assert ( + sync_tensor[1].item() == -sync_tensor[2].item() + ), "write_back duplicate-reclaim victims diverged across TP ranks" while finish_count > 0: ack = cc.ack_load_queue.pop(0) @@ -1880,6 +1894,8 @@ def loading_check(self, finish_count: Optional[int] = None) -> None: node, lock_params, host_lock_params = self.ongoing_load_back.pop(ack_id) self.dec_lock_ref(node, lock_params) self.dec_host_lock_ref(node, host_lock_params) + # Unpin the loaded nodes; host copies stay as reclaimable duplicates. + self.tree_core.finish_load_back(node) if self.metrics_collector is not None: self.metrics_collector.increment_load_back_num_tokens(ack.num_tokens) diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index b875e4b9b74d..a8887b7153d7 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -2786,6 +2786,8 @@ def _simulate_backup(self, cache, node): cd = ancestor.component_data[ct] if cd.value is not None and cd.host_value is None: cd.host_value = cd.value.clone() + # A real backup registers duplicate tracking at its ack. + cache.tree_core._update_duplicate_tracking(ancestor) def _simulate_backup_tree(self, cache): """Backup all non-root nodes (simulates write-through).""" From 55eaf83d6f1ba13c50b443aa340a54559906ccd1 Mon Sep 17 00:00:00 2001 From: Ke Bao Date: Sat, 15 Aug 2026 01:12:38 +0800 Subject: [PATCH 05/12] Retain SWA down to the last state checkpoint (#34729) (cherry picked from commit bacab43e9a914513a778dc058dd702443d11abe1) --- python/sglang/srt/managers/schedule_batch.py | 1 + .../sglang/srt/mem_cache/base_prefix_cache.py | 7 + python/sglang/srt/mem_cache/common.py | 7 + .../unified_cache/components/swa_component.py | 1 + .../srt/mem_cache/unified_radix_cache.py | 8 + .../mem_cache/test_swa_eviction_boundary.py | 141 ++++++++++++++++++ 6 files changed, 165 insertions(+) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index a45d122f0ff5..3653ddd9e377 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -3271,6 +3271,7 @@ def _evict_swa(self, req: Req, pre_len: int): req_to_token_pool=self.req_to_token_pool, token_to_kv_pool_allocator=self.token_to_kv_pool_allocator, is_chunk_cache=self.tree_cache.is_chunk_cache(), + retain_floor=self.tree_cache.swa_retain_floor(req), ) def __str__(self): diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index 1f40409a416b..0c24cb2f18ff 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -377,6 +377,13 @@ def take_events(self): def supports_swa(self) -> bool: return False + def swa_retain_floor(self, req) -> int | None: + # A match lands on a state checkpoint rather than on the tail, so a cache + # that pairs SWA with mamba/conv checkpoints has to keep the window behind + # the last checkpoint. Those caches override this. Everyone else has + # nothing deeper than the tail to protect. + return None + def swa_reprefill_tail_tokens(self) -> int: # Only the unified_kv compress-only HiCache layout needs to hold back a # trailing sliding window for re-prefill; every other cache keeps SWA diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 9d18b37e67a9..19e1c3e120ed 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -53,6 +53,7 @@ def free_swa_out_of_window_slots( req_to_token_pool: ReqToTokenPool, token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator, is_chunk_cache: bool = False, + retain_floor: int | None = None, ) -> None: if req.kv is None: return @@ -76,6 +77,12 @@ def free_swa_out_of_window_slots( # boundary (page_floor(seq_len)) so the last leaf is never all-tombstone. # No extra page margin is needed. evict_threshold = pre_len - max(sliding_window_size, page_size) + if retain_floor is not None and not is_chunk_cache: + # The caller owns where the floor is (see BasePrefixCache.swa_retain_floor); + # this only promises not to free past it. Chunk cache has no tree, so a + # retained checkpoint could never be matched and holding it is pure cost. + evict_threshold = min(evict_threshold, retain_floor) + new_swa_evicted_seqlen = max( req.kv.swa_evicted_seqlen, evict_threshold, diff --git a/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py b/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py index d35ef0731b53..36920e5263eb 100644 --- a/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py +++ b/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py @@ -752,6 +752,7 @@ def free_out_of_window_slots( page_size=self.cache.page_size, req_to_token_pool=self.cache.req_to_token_pool, token_to_kv_pool_allocator=self.cache.token_to_kv_pool_allocator, + retain_floor=self.cache.swa_retain_floor(req), ) insert_params.swa_evicted_seqlen = req.kv.swa_evicted_seqlen diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index f01b3c4f37b8..6f52027438a9 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -2023,6 +2023,14 @@ def swa_reprefill_tail_tokens(self) -> int: ) return swa.sliding_window_size if unified_compress_only_hicache else 0 + def swa_retain_floor(self, req) -> int | None: + if not self.is_mamba_enabled or self._sliding_window_size is None: + return None + checkpoint = req.mamba_last_track_seqlen + if checkpoint is None: + return None + return checkpoint - self._sliding_window_size + def supports_swa(self) -> bool: return self.is_swa_enabled diff --git a/test/registered/unit/mem_cache/test_swa_eviction_boundary.py b/test/registered/unit/mem_cache/test_swa_eviction_boundary.py index cd81142378e1..934a800b0ecc 100644 --- a/test/registered/unit/mem_cache/test_swa_eviction_boundary.py +++ b/test/registered/unit/mem_cache/test_swa_eviction_boundary.py @@ -21,6 +21,7 @@ from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator from sglang.srt.mem_cache.cache_init_params import CacheInitParams +from sglang.srt.mem_cache.common import free_swa_out_of_window_slots from sglang.srt.mem_cache.memory_pool import ReqToTokenPool from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache @@ -195,6 +196,146 @@ def test_formula_page_leq_window(self): ) tree.sanity_check() + # -- Retention floor: never free past the last state checkpoint -- + + def test_retain_floor_clamps_eviction(self): + """A hybrid cache keeps SWA down to the last state checkpoint, not to the + window behind the tail, because that is where a prefix match lands. The + floor must clamp the frontier even though the tail has moved far past it.""" + page_size, window = 8, 16 + tree, allocator, pool = _build_swa_tree( + page_size=page_size, sliding_window_size=window + ) + seq_len = 200 + checkpoint = 96 + kv = _swa_alloc(allocator, seq_len) + pool.write((0, slice(0, seq_len)), kv) + req = _make_req(0, list(range(seq_len)), 0, tree) + batch = _make_batch(tree, allocator, pool) + + free_swa_out_of_window_slots( + req, + seq_len - 1, + sliding_window_size=window, + page_size=page_size, + req_to_token_pool=batch.req_to_token_pool, + token_to_kv_pool_allocator=batch.token_to_kv_pool_allocator, + retain_floor=checkpoint - window, + ) + + # Without the floor this would reach page_floor(199 - 16) = 176. + self.assertLessEqual(req.kv.swa_evicted_seqlen, checkpoint - window) + self.assertEqual(req.kv.swa_evicted_seqlen % page_size, 0) + + def test_retain_floor_ignored_for_chunk_cache(self): + """Chunk cache builds no tree, so a retained checkpoint could never be + matched. Holding it would cost SWA slots for nothing.""" + page_size, window = 8, 16 + seq_len = 200 + tree, allocator, pool = _build_swa_tree( + page_size=page_size, sliding_window_size=window + ) + kv = _swa_alloc(allocator, seq_len) + pool.write((0, slice(0, seq_len)), kv) + req = _make_req(0, list(range(seq_len)), 0, tree) + batch = _make_batch(tree, allocator, pool) + + free_swa_out_of_window_slots( + req, + seq_len - 1, + sliding_window_size=window, + page_size=page_size, + req_to_token_pool=batch.req_to_token_pool, + token_to_kv_pool_allocator=batch.token_to_kv_pool_allocator, + is_chunk_cache=True, + retain_floor=16, + ) + + expected = (seq_len - 1 - window) // page_size * page_size + self.assertEqual(req.kv.swa_evicted_seqlen, expected) + + def test_retain_floor_none_matches_old_behaviour(self): + """retain_floor=None must reproduce the pre-change frontier exactly, so a + cache without a second state stream is unaffected.""" + page_size, window = 8, 16 + seq_len = 200 + frontiers = [] + for floor in (None, "absent"): + tree, allocator, pool = _build_swa_tree( + page_size=page_size, sliding_window_size=window + ) + kv = _swa_alloc(allocator, seq_len) + pool.write((0, slice(0, seq_len)), kv) + req = _make_req(0, list(range(seq_len)), 0, tree) + batch = _make_batch(tree, allocator, pool) + kwargs = {} if floor == "absent" else {"retain_floor": None} + free_swa_out_of_window_slots( + req, + seq_len - 1, + sliding_window_size=window, + page_size=page_size, + req_to_token_pool=batch.req_to_token_pool, + token_to_kv_pool_allocator=batch.token_to_kv_pool_allocator, + **kwargs, + ) + frontiers.append(req.kv.swa_evicted_seqlen) + + expected = (seq_len - 1 - max(window, page_size)) // page_size * page_size + self.assertEqual(frontiers[0], expected) + self.assertEqual(frontiers[1], expected) + + def test_retain_floor_above_threshold_is_inert(self): + """The floor is a min(), so a checkpoint that is already inside the window + must not hold anything extra.""" + page_size, window = 8, 16 + seq_len = 200 + tree, allocator, pool = _build_swa_tree( + page_size=page_size, sliding_window_size=window + ) + kv = _swa_alloc(allocator, seq_len) + pool.write((0, slice(0, seq_len)), kv) + req = _make_req(0, list(range(seq_len)), 0, tree) + batch = _make_batch(tree, allocator, pool) + + free_swa_out_of_window_slots( + req, + seq_len - 1, + sliding_window_size=window, + page_size=page_size, + req_to_token_pool=batch.req_to_token_pool, + token_to_kv_pool_allocator=batch.token_to_kv_pool_allocator, + retain_floor=seq_len, + ) + + expected = (seq_len - 1 - max(window, page_size)) // page_size * page_size + self.assertEqual(req.kv.swa_evicted_seqlen, expected) + + def test_retain_floor_does_not_unfree(self): + """The frontier only advances. A floor arriving after slots were already + freed must not claim them back, which would double-free on the next pass.""" + page_size, window = 8, 16 + tree, allocator, pool = _build_swa_tree( + page_size=page_size, sliding_window_size=window + ) + seq_len = 200 + kv = _swa_alloc(allocator, seq_len) + pool.write((0, slice(0, seq_len)), kv) + req = _make_req(0, list(range(seq_len)), 0, tree) + batch = _make_batch(tree, allocator, pool) + common_kwargs = dict( + sliding_window_size=window, + page_size=page_size, + req_to_token_pool=batch.req_to_token_pool, + token_to_kv_pool_allocator=batch.token_to_kv_pool_allocator, + ) + + free_swa_out_of_window_slots(req, seq_len - 1, **common_kwargs) + advanced = req.kv.swa_evicted_seqlen + self.assertGreater(advanced, 0) + + free_swa_out_of_window_slots(req, seq_len - 1, retain_floor=0, **common_kwargs) + self.assertEqual(req.kv.swa_evicted_seqlen, advanced) + # -- Eviction formula: page_size == 1 -- def test_formula_page_size_1(self): From 81d902015abff76d0dcb64ad4d55f2878aefa6dd Mon Sep 17 00:00:00 2001 From: Ke Bao Date: Fri, 14 Aug 2026 17:13:21 +0800 Subject: [PATCH 06/12] Skip oow slot freeing under eagle (#34823) (cherry picked from commit fede84057fcf579620ef6921e01ed4d8cb1eb09d) --- python/sglang/srt/mem_cache/unified_radix_cache.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 6f52027438a9..3c6fe224f215 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -723,7 +723,12 @@ def cache_unfinished_req(self, req: Req, chunked: bool = False, **kwargs) -> Non if cl is not None: effective_cache_len = min(effective_cache_len, cl) - if envs.SGLANG_OPT_UNIFIED_CACHE_FREE_OUT_OF_WINDOW_SLOTS.get(): + # swa_evicted_seqlen is a raw-token length, but under EAGLE the insert key is + # bigram-indexed, so SWA would carve tombstones at the wrong offset (#34653). + if ( + envs.SGLANG_OPT_UNIFIED_CACHE_FREE_OUT_OF_WINDOW_SLOTS.get() + and not self.tree_core.is_eagle + ): for comp in self._components_tuple: comp.free_out_of_window_slots( req, effective_cache_len - 1, insert_params From 2cde519e589e84c2c989fedb935556c93d62c752 Mon Sep 17 00:00:00 2001 From: Ke Bao Date: Sun, 16 Aug 2026 11:52:14 +0800 Subject: [PATCH 07/12] Fix swa eviction frontier for bigram keys (#34870) (cherry picked from commit 7b0c65f7e5497167855824f639e20d3aea7517e3) --- .../srt/mem_cache/unified_radix_cache.py | 28 +++---- .../test_unified_radix_cache_unittest.py | 81 +++++++++++++++++++ 2 files changed, 95 insertions(+), 14 deletions(-) diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 3c6fe224f215..ddce09704aa9 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -723,16 +723,20 @@ def cache_unfinished_req(self, req: Req, chunked: bool = False, **kwargs) -> Non if cl is not None: effective_cache_len = min(effective_cache_len, cl) - # swa_evicted_seqlen is a raw-token length, but under EAGLE the insert key is - # bigram-indexed, so SWA would carve tombstones at the wrong offset (#34653). - if ( - envs.SGLANG_OPT_UNIFIED_CACHE_FREE_OUT_OF_WINDOW_SLOTS.get() - and not self.tree_core.is_eagle - ): + radix_key = RadixKey( + token_ids[:effective_cache_len], + req.extra_key, + is_bigram=self.tree_core.is_eagle, + ) + + if envs.SGLANG_OPT_UNIFIED_CACHE_FREE_OUT_OF_WINDOW_SLOTS.get(): + # The frontier lands a page below page_floor(pre_len + 1), which has to + # be where the insert stops, or the leaf it creates keeps less than a + # sliding window of live SWA and the match after the insert rejects it. + # The insert stops at page_floor(len(radix_key)), and a bigram key is + # one shorter than the tokens it spans, so measure the key. for comp in self._components_tuple: - comp.free_out_of_window_slots( - req, effective_cache_len - 1, insert_params - ) + comp.free_out_of_window_slots(req, len(radix_key) - 1, insert_params) if effective_cache_len <= 0: req.prefix_indices = kv_indices_orig.to(dtype=torch.int64, copy=True) @@ -744,11 +748,7 @@ def cache_unfinished_req(self, req: Req, chunked: bool = False, **kwargs) -> Non kv_indices = kv_indices_orig[:effective_cache_len] - radix_key = RadixKey( - token_ids[:effective_cache_len], - req.extra_key, - is_bigram=self.tree_core.is_eagle, - ).page_aligned(self.page_size) + radix_key = radix_key.page_aligned(self.page_size) page_aligned_len = len(radix_key) values = kv_indices[:page_aligned_len].to(dtype=torch.int64, copy=True) diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index a8887b7153d7..c72f92a76f2c 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -6408,5 +6408,86 @@ def test_prefetch_refill_kept_under_unbacked_parent_in_write_back(self): cache.sanity_check() +class TestSWAWindowUnderBigramKey(CustomTestCase): + """`cache_unfinished_req` has to leave the leaf it inserts holding a full + sliding window of live SWA. Otherwise the match that follows the insert + refuses that leaf, `cache_protected_len` never advances, and the next insert + frees KV the tree already owns as if it were the request's duplicate. + + An EAGLE bigram key holds one entry less than the tokens it spans, so the + leaf stops at page_floor(len(key)) rather than page_floor(seq_len), a page + lower. The eviction frontier has to be measured against the key. + """ + + cfg = CacheConfig( + page_size=4, + components=(ComponentType.FULL, ComponentType.SWA), + sliding_window_size=7, + is_eagle=True, + kv_size=256, + max_context_len=64, + ) + + def _alloc_paged(self, allocator, need_size): + ps = self.cfg.page_size + aligned = ((need_size + ps - 1) // ps) * ps + full_indices = allocator.full_attn_allocator.alloc(aligned) + swa_indices = allocator.swa_attn_allocator.alloc(aligned) + self.assertIsNotNone(full_indices) + self.assertIsNotNone(swa_indices) + allocator.full_to_swa_index_mapping[full_indices] = swa_indices + return full_indices[:need_size] + + def test_match_after_insert_reaches_the_new_leaf(self): + cache, allocator, req_to_token_pool = build_fixture(self.cfg) + page_size = self.cfg.page_size + # Page-aligned length is the shape that costs the bigram key a page. + seq_len = 4 * page_size + + req = Req( + rid=0, + origin_input_text="", + origin_input_ids=array("q"), + sampling_params=SamplingParams(temperature=0, max_new_tokens=1), + ) + req_to_token_pool.alloc([req]) + tokens = list(range(1, seq_len + 1)) + req.origin_input_ids = tokens + req.output_ids = [] + req.full_untruncated_fill_ids = array("q", tokens) + req.set_extend_range(0, len(req.full_untruncated_fill_ids)) + kv_indices = self._alloc_paged(allocator, seq_len) + req_to_token_pool.write((req.req_pool_idx, slice(0, seq_len)), kv_indices) + req.kv_committed_len = seq_len + req.last_node = cache.root_node.id + req.cache_protected_len = 0 + req.swa_uuid_for_lock = None + req.extra_key = None + req.kv = ReqKvInfo(kv_allocated_len=0, swa_evicted_seqlen=0) + + with envs.SGLANG_OPT_UNIFIED_CACHE_FREE_OUT_OF_WINDOW_SLOTS.override(True): + cache.cache_unfinished_req(req) + + boundary = (seq_len - 1) // page_size * page_size + self.assertGreaterEqual( + boundary - req.kv.swa_evicted_seqlen, + self.cfg.sliding_window_size, + f"leaf ending at {boundary} keeps only " + f"{boundary - req.kv.swa_evicted_seqlen} live SWA tokens against a " + f"{self.cfg.sliding_window_size} window", + ) + self.assertEqual( + req.cache_protected_len, + boundary, + "the match after the insert must reach the leaf the insert created", + ) + + cache.dec_lock_ref( + req.last_node, + DecLockRefParams(swa_uuid_for_lock=getattr(req, "swa_uuid_for_lock", None)), + ) + cache.sanity_check() + + if __name__ == "__main__": unittest.main() From 0e51034d20cec6dcbf7c57844e5ec7ccf4e1caf6 Mon Sep 17 00:00:00 2001 From: wangwenmingaa <30922691+wangwenmingaa@users.noreply.github.com> Date: Sun, 16 Aug 2026 12:17:33 +0800 Subject: [PATCH 08/12] [HiCache] Optimize LogicalHostPool free-list release (#33998) (cherry picked from commit b9c4c16925ca63802a3debc56cb8bc8015c5ca86) --- .../sglang/srt/mem_cache/memory_pool_host.py | 31 ++++++++++++++++--- .../unit/mem_cache/test_mem_pool_host.py | 16 ++++++++++ 2 files changed, 43 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index 2670de52c526..e05ae35464b2 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -669,9 +669,25 @@ def __init__(self, size: int, page_size: int, layout: str = "layer_first"): @synchronized def clear(self): self.free_slots = torch.arange(self.size, dtype=torch.int64) + # Match HostKVCache's lazy release path: defer large free-list merges + # until an allocation needs the released slots. + self.release_slots = [] + self.num_release_slots = 0 def available_size(self): - return len(self.free_slots) + return len(self.free_slots) + self.num_release_slots + + def _merge_release_slots(self): + if self.num_release_slots == 0: + return + + if len(self.free_slots) == 0 and len(self.release_slots) == 1: + self.free_slots = self.release_slots[0] + else: + self.free_slots = torch.cat([self.free_slots, *self.release_slots]) + + self.release_slots = [] + self.num_release_slots = 0 @synchronized def alloc(self, need_size: int) -> Optional[torch.Tensor]: @@ -682,6 +698,10 @@ def alloc(self, need_size: int) -> Optional[torch.Tensor]: ) if need_size > self.available_size(): return None + + if need_size > len(self.free_slots): + self._merge_release_slots() + select_index = self.free_slots[:need_size] self.free_slots = self.free_slots[need_size:] return select_index @@ -693,9 +713,12 @@ def free(self, indices: torch.Tensor) -> int: "LogicalHostPool free must be page-aligned, " f"got len(indices)={len(indices)}, page_size={self.page_size}" ) - self.free_slots = torch.cat( - [self.free_slots, indices.to(dtype=torch.int64, device="cpu").flatten()] - ) + indices_cpu = indices.to(dtype=torch.int64, device="cpu").flatten() + if indices_cpu.numel() == 0: + return 0 + + self.release_slots.append(indices_cpu) + self.num_release_slots += len(indices_cpu) return len(indices) def backup_from_device_all_layer( diff --git a/test/registered/unit/mem_cache/test_mem_pool_host.py b/test/registered/unit/mem_cache/test_mem_pool_host.py index ce94365d41bd..dae2acb2fbaa 100644 --- a/test/registered/unit/mem_cache/test_mem_pool_host.py +++ b/test/registered/unit/mem_cache/test_mem_pool_host.py @@ -8,6 +8,7 @@ from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool from sglang.srt.mem_cache.memory_pool_host import ( DeepSeekV4PagedHostPool, + LogicalHostPool, MambaPoolHost, ) from sglang.srt.mem_cache.pool_host.mha import MHATokenToKVPoolHost @@ -127,6 +128,10 @@ def _make_deepseek_v4_pool(): pool.clear() return pool + @staticmethod + def _make_logical_pool(): + return LogicalHostPool(size=8, page_size=2) + def _assert_lazy_release(self, pool): self.assertEqual(pool.free(torch.empty(0, dtype=torch.int64)), 0) self.assertEqual(pool.num_release_slots, 0) @@ -181,6 +186,17 @@ def test_deepseek_v4_pool_lazy_release(self): pool.clear() self.assertEqual(len(pool.alloc(1)), 2) + def test_logical_pool_lazy_release(self): + pool = self._make_logical_pool() + self._assert_lazy_release(pool) + + # Preserve the logical pool's strict page-alignment checks. + pool.clear() + with self.assertRaises(ValueError): + pool.alloc(1) + with self.assertRaises(ValueError): + pool.free(torch.tensor([0])) + if __name__ == "__main__": unittest.main() From fabf7c24fe403fb3c21692d7ee9ccf5f8306c9a3 Mon Sep 17 00:00:00 2001 From: Zhiqiang Xie Date: Wed, 5 Aug 2026 14:13:06 -0700 Subject: [PATCH 09/12] Observability enhancement for HiCache (#32388) (cherry picked from commit 41d1b33f12cf3f6dd4057d2e19e50be3485e7e1a) --- .../sglang/srt/managers/cache_controller.py | 32 +++- python/sglang/srt/managers/schedule_batch.py | 42 +++--- python/sglang/srt/managers/schedule_policy.py | 26 +++- .../scheduler_components/metrics_reporter.py | 24 +-- .../srt/mem_cache/hi_mamba_radix_cache.py | 8 +- python/sglang/srt/mem_cache/hiradix_cache.py | 31 +++- .../hybrid_cache/hybrid_cache_controller.py | 55 ++++++- .../srt/mem_cache/unified_radix_cache.py | 64 +++++++- .../srt/observability/metrics_collector.py | 139 +++++++++++++++++- .../unit/managers/test_prefill_adder.py | 2 + .../test_hicache_load_back_timing.py | 6 +- ...test_hicache_staged_write_back_dispatch.py | 40 ++++- 12 files changed, 420 insertions(+), 49 deletions(-) diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index 0f3234b3b358..beb4b398088e 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -163,6 +163,11 @@ class HiCacheAck(NamedTuple): node_ids: List[int] num_tokens: int = 0 timing_enabled: bool = False + # Tokens transferred per host pool (PoolName value -> count). + num_tokens_by_pool: Optional[dict[str, int]] = None + # Total bytes moved by the op across all pools, including draft piggyback + # and sidecar transfers that the per-pool token counts exclude. + num_bytes: int = 0 class StorageOperation: @@ -714,11 +719,12 @@ def start_writing(self) -> None: self.write_queue.clear() start_event = device_module.Event() - finish_event = device_module.Event() + ack_start_event, ack_finish_event, timing_enabled = make_timing_event_pair() start_event.record() with device_module.stream(self.write_stream): start_event.wait(self.write_stream) + ack_start_event.record() self.mem_pool_host.backup_from_device_all_layer( self.mem_pool_device, host_indices, device_indices, self.io_backend ) @@ -729,7 +735,7 @@ def start_writing(self) -> None: device_indices, self.io_backend, ) - finish_event.record() + ack_finish_event.record() # NOTE: We must save the host indices and device indices here, # this is because we need to guarantee that these tensors are # still alive when the write stream is executing. @@ -738,7 +744,25 @@ def start_writing(self) -> None: if device_indices.is_cuda: device_indices.record_stream(self.write_stream) - self.ack_write_queue.append(HiCacheAck(start_event, finish_event, op.node_ids)) + self.ack_write_queue.append( + HiCacheAck( + start_event=ack_start_event, + finish_event=ack_finish_event, + node_ids=op.node_ids, + num_tokens=len(op.device_indices), + timing_enabled=timing_enabled, + num_tokens_by_pool={PoolName.KV.value: len(op.device_indices)}, + num_bytes=self._transfer_num_bytes(op), + ) + ) + + def _transfer_num_bytes(self, op: CacheOperation) -> int: + """Total bytes moved by a merged transfer op (draft piggyback included).""" + num_tokens = len(op.device_indices) + num_bytes = num_tokens * self.mem_pool_host.size_per_token + if self.has_draft: + num_bytes += num_tokens * self.mem_pool_host_draft.size_per_token + return num_bytes def load( self, @@ -830,6 +854,8 @@ def start_loading(self) -> int: node_ids=op.node_ids, num_tokens=len(op.device_indices), timing_enabled=timing_enabled, + num_tokens_by_pool={PoolName.KV.value: len(op.device_indices)}, + num_bytes=self._transfer_num_bytes(op), ) ) return producer_id diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 3653ddd9e377..fa3ca04a7fce 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -191,6 +191,21 @@ def sanity_check_mm_pad_shift_value(vocab_size: int) -> None: ) +def split_cached_prefix_by_tier( + prefix_len: int, host_hit_len: int, storage_hit_len: int +) -> tuple[int, int, int]: + """Split a request's cached prefix into (device, host, storage) tokens. + + prefix_len is len(prefix_indices) AFTER host load-back, so it contains the + host-loaded portion; host_hit_len in turn contains the storage-prefetched + portion (storage is clamped to it to handle edge cases). + """ + storage = min(host_hit_len, storage_hit_len) + host = host_hit_len - storage + device = max(0, prefix_len - host_hit_len) + return device, host, storage + + def _compute_pad_value(hash: int) -> int: """Compute pad value from hash.""" return MM_PAD_SHIFT_VALUE + (hash % (1 << 30)) @@ -2393,24 +2408,17 @@ def prepare_for_extend(self): # Only compute once on FIRST chunk - subsequent chunks in chunked prefill # would incorrectly count previously computed tokens as cache hits. if not req._cache_breakdown_computed: - # At this point, prefix_indices has been extended with host data - # via init_load_back in schedule_policy, so: - # - len(prefix_indices) = device_original + host_loaded - # - host_hit_length = total tokens from host cache (including storage-prefetched) - # - storage_hit_length = tokens loaded from storage backend (L3 hits) - # - device_portion = len(prefix_indices) - host_hit_length - # - # Storage hits are now tracked via scheduler after prefetch completes. # storage_hit_length is set by scheduler.pop_prefetch_loaded_tokens() - host_total = req.host_hit_length - # Clamp storage to host_total to handle edge cases - storage_portion = min(host_total, req.storage_hit_length) - host_portion = host_total - storage_portion - device_portion = max(0, len(req.prefix_indices) - host_total) - - req.cached_tokens_device = device_portion - req.cached_tokens_host = host_portion - req.cached_tokens_storage = storage_portion + # after prefetch completes. + ( + req.cached_tokens_device, + req.cached_tokens_host, + req.cached_tokens_storage, + ) = split_cached_prefix_by_tier( + prefix_len=len(req.prefix_indices), + host_hit_len=req.host_hit_length, + storage_hit_len=req.storage_hit_length, + ) req._cache_breakdown_computed = True req.already_computed = seq_len diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 39bc02b780f9..65526e595c0d 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -38,7 +38,11 @@ from sglang.srt.dllm.config import DllmConfig from sglang.srt.layers.attention.dsa.utils import is_dsa_prefill_cp_in_seq_split from sglang.srt.layers.utils.cp_utils import is_prefill_context_parallel_enabled -from sglang.srt.managers.schedule_batch import Req, ScheduleBatch +from sglang.srt.managers.schedule_batch import ( + Req, + ScheduleBatch, + split_cached_prefix_by_tier, +) from sglang.srt.mem_cache.allocator.hisparse import ( DeepSeekV4HiSparseTokenToKVPoolAllocator, ) @@ -483,6 +487,9 @@ def __init__( self.new_chunked_req = None self.log_hit_tokens = 0 self.reprocessed_log_hit_tokens = 0 + self.log_device_hit_tokens = 0 + self.log_host_hit_tokens = 0 + self.log_storage_hit_tokens = 0 # TODO(lsyin): report the real input tokens excluding page alignment self.log_input_tokens = 0 self.reprocessed_log_input_tokens = 0 @@ -745,6 +752,8 @@ def _update_prefill_budget( max_new_tokens: int, retracted_stain: bool, mamba_gap_reserve: int = 0, + host_hit_len: int = 0, + storage_hit_len: int = 0, ): # TODO(lsyin): check this workaround logic, which only ensures the prefill will not out of memory, and may be too conservative extend_input_len = self.ceil_paged_tokens(extend_input_len) @@ -784,6 +793,15 @@ def _update_prefill_budget( if retracted_stain: self.reprocessed_log_hit_tokens += prefix_len self.reprocessed_log_input_tokens += extend_input_len + elif prefix_len > 0: + device_hit, host_hit, storage_hit = split_cached_prefix_by_tier( + prefix_len=prefix_len, + host_hit_len=host_hit_len, + storage_hit_len=storage_hit_len, + ) + self.log_device_hit_tokens += device_hit + self.log_host_hit_tokens += host_hit + self.log_storage_hit_tokens += storage_hit def _get_dllm_remain_tokens(self) -> int: _rem_tokens = min( @@ -816,6 +834,8 @@ def _add_dllm_req(self, req: Req, prefix_len: int): 0, req.retracted_stain, mamba_gap_reserve=self._mamba_gap_budget_for_req(req), + host_hit_len=req.host_hit_length, + storage_hit_len=req.storage_hit_length, ) def _req_inc_lock_ref(self, req: Req): @@ -1209,6 +1229,8 @@ def add_one_req( ), req.retracted_stain, mamba_gap_reserve=self._mamba_gap_budget_for_req(req), + host_hit_len=req.host_hit_length, + storage_hit_len=req.storage_hit_length, ) else: # Make sure at least one page is available @@ -1250,6 +1272,8 @@ def add_one_req( 0, req.retracted_stain, mamba_gap_reserve=self._mamba_gap_budget_for_req(req), + host_hit_len=req.host_hit_length, + storage_hit_len=req.storage_hit_length, ) return self.budget_state() diff --git a/python/sglang/srt/managers/scheduler_components/metrics_reporter.py b/python/sglang/srt/managers/scheduler_components/metrics_reporter.py index 13ba75da5bf2..9e1f28cef077 100644 --- a/python/sglang/srt/managers/scheduler_components/metrics_reporter.py +++ b/python/sglang/srt/managers/scheduler_components/metrics_reporter.py @@ -7,13 +7,7 @@ import time from collections import defaultdict from dataclasses import dataclass -from typing import ( - TYPE_CHECKING, - List, - Optional, - Tuple, - Union, -) +from typing import TYPE_CHECKING, List, Optional, Tuple, Union from sglang.srt.disaggregation.utils import DisaggregationMode from sglang.srt.environ import envs @@ -64,6 +58,9 @@ class PrefillStats: num_new_seqs: int # len(can_run_list) reprocessed_log_input_tokens: int = 0 reprocessed_log_hit_tokens: int = 0 + log_device_hit_tokens: int = 0 + log_host_hit_tokens: int = 0 + log_storage_hit_tokens: int = 0 num_pending_tokens: int = 0 @classmethod @@ -79,6 +76,9 @@ def from_adder( log_hit_tokens=adder.log_hit_tokens, reprocessed_log_input_tokens=adder.reprocessed_log_input_tokens, reprocessed_log_hit_tokens=adder.reprocessed_log_hit_tokens, + log_device_hit_tokens=adder.log_device_hit_tokens, + log_host_hit_tokens=adder.log_host_hit_tokens, + log_storage_hit_tokens=adder.log_storage_hit_tokens, new_token_ratio=adder.new_token_ratio, num_running_reqs=QueueCount.from_reqs( running_reqs, enable_priority_scheduling @@ -637,6 +637,12 @@ def report_prefill_stats( cache_hit_rate = ( effective_hit_tokens / total_tokens if total_tokens > 0 else 0.0 ) + self.metrics_collector.increment_effective_prefill_tokens( + input_tokens=effective_input_tokens, + device_hit_tokens=prefill_stats.log_device_hit_tokens, + host_hit_tokens=prefill_stats.log_host_hit_tokens, + storage_hit_tokens=prefill_stats.log_storage_hit_tokens, + ) # Basics if ( @@ -970,9 +976,7 @@ def _emit_forward_pass_metrics( if not self.scheduler.enable_fpm: return - from sglang.srt.observability.forward_pass_metrics import ( - ForwardPassMetrics, - ) + from sglang.srt.observability.forward_pass_metrics import ForwardPassMetrics if self.scheduler._fpm_uses_device_timer: self.forward_pass_device_timer._report() diff --git a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py index dbcec3547f5c..8671a7a75370 100644 --- a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py @@ -455,7 +455,13 @@ def loading_check(self): self.dec_lock_ref(end_node) if self.metrics_collector is not None: - self.metrics_collector.increment_load_back_num_tokens(ack.num_tokens) + for pool, num_tokens in (ack.num_tokens_by_pool or {}).items(): + if num_tokens > 0: + self.metrics_collector.increment_load_back_num_tokens( + num_tokens=num_tokens, pool=pool + ) + if ack.num_bytes > 0: + self.metrics_collector.increment_load_back_num_bytes(ack.num_bytes) if ack.timing_enabled: duration_ms = ack.start_event.elapsed_time(ack.finish_event) self.metrics_collector.observe_load_back_duration( diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 5ebac56f30eb..74dec738fccb 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -1028,6 +1028,7 @@ def writing_check(self, write_back=False, finish_count: Optional[int] = None): ack.finish_event.synchronize() for ack_id in ack.node_ids: self._finish_write_through_ack(ack_id, release_lock=False) + self._log_write_ack_metrics(ack) self.cache_controller.ack_write_queue.clear() assert len(self.ongoing_write_through) == 0 return @@ -1055,8 +1056,24 @@ def writing_check(self, write_back=False, finish_count: Optional[int] = None): ack.finish_event.synchronize() for ack_id in ack.node_ids: self._finish_write_through_ack(ack_id, release_lock=True) + self._log_write_ack_metrics(ack) finish_count -= 1 + def _log_write_ack_metrics(self, ack) -> None: + """Record D->H backup volume and duration for a completed write ack.""" + if self.metrics_collector is None: + return + for pool, num_tokens in (ack.num_tokens_by_pool or {}).items(): + if num_tokens > 0: + self.metrics_collector.increment_backup_num_tokens( + num_tokens=num_tokens, pool=pool + ) + if ack.num_bytes > 0: + self.metrics_collector.increment_backup_num_bytes(ack.num_bytes) + if ack.timing_enabled: + duration_ms = ack.start_event.elapsed_time(ack.finish_event) + self.metrics_collector.observe_backup_duration(duration_ms / 1000.0) + def loading_check(self, finish_count: Optional[int] = None): if finish_count is None: finish_count = 0 @@ -1080,7 +1097,13 @@ def loading_check(self, finish_count: Optional[int] = None): self.dec_lock_ref(end_node) if self.metrics_collector is not None: - self.metrics_collector.increment_load_back_num_tokens(ack.num_tokens) + for pool, num_tokens in (ack.num_tokens_by_pool or {}).items(): + if num_tokens > 0: + self.metrics_collector.increment_load_back_num_tokens( + num_tokens=num_tokens, pool=pool + ) + if ack.num_bytes > 0: + self.metrics_collector.increment_load_back_num_bytes(ack.num_bytes) if ack.timing_enabled: duration_ms = ack.start_event.elapsed_time(ack.finish_event) self.metrics_collector.observe_load_back_duration( @@ -1300,6 +1323,12 @@ def _drop_subtree_no_host(self, root: TreeNode) -> int: root.parent.children.pop(key, None) self._update_leaf_status(root.parent) self._update_host_leaf_status(root.parent) + if freed_device > 0 and self.metrics_collector is not None: + self.metrics_collector.increment_dropped_tokens( + num_tokens=freed_device, + reason="host_pressure", + pool=PoolName.KV.value, + ) return freed_device def evict_host(self, num_tokens: int): diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py index 5764a9f600e0..191cd8a9038b 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py @@ -419,10 +419,11 @@ def start_writing(self) -> None: ) self.write_queue.clear() start_event = device_module.Event() - finish_event = device_module.Event() + ack_start_event, ack_finish_event, timing_enabled = make_timing_event_pair() start_event.record() with device_module.stream(self.write_stream): start_event.wait(self.write_stream) + ack_start_event.record() self.mem_pool_host.backup_from_device_all_layer( self.mem_pool_device, host_indices, @@ -437,14 +438,60 @@ def start_writing(self) -> None: device_indices, self.io_backend, ) - finish_event.record() + ack_finish_event.record() self._record_transfer_indices_on_stream( self.write_stream, host_indices, device_indices, resolved_pool_transfers, ) - self.ack_write_queue.append(HiCacheAck(start_event, finish_event, op.node_ids)) + self.ack_write_queue.append( + HiCacheAck( + start_event=ack_start_event, + finish_event=ack_finish_event, + node_ids=op.node_ids, + num_tokens=len(op.device_indices), + timing_enabled=timing_enabled, + num_tokens_by_pool=self._num_tokens_by_pool(op), + num_bytes=self._transfer_num_bytes(op), + ) + ) + + def _num_tokens_by_pool(self, op: CacheOperation) -> dict[str, int]: + """Per-pool token counts for a merged transfer op (anchor + extra + pools), shared by D->H write and H->D load acks; sidecar transfers + reusing another pool's indices are excluded.""" + counts = {self.mem_pool_host.anchor_entry.name.value: len(op.device_indices)} + for transfer in op.pool_transfers or []: + if transfer.indices_from_pool is not None or transfer.host_indices is None: + continue + name = transfer.name.value + counts[name] = counts.get(name, 0) + len(transfer.host_indices) + return counts + + def _transfer_num_bytes(self, op: CacheOperation) -> int: + """Total bytes moved by a merged transfer op across all pools, + including draft piggyback and sidecar transfers riding another + pool's indices (both excluded from the per-pool token counts).""" + kv_tokens = len(op.device_indices) + num_bytes = kv_tokens * self.mem_pool_host.anchor_entry.host_pool.size_per_token + if self.has_draft: + num_bytes += kv_tokens * self.mem_pool_host_draft.size_per_token + # Slot counts of the pools sidecars can ride on. + source_len = {self.mem_pool_host.anchor_entry.name: kv_tokens} + for t in op.pool_transfers or []: + if t.indices_from_pool is None and t.host_indices is not None: + source_len[t.name] = len(t.host_indices) + for t in op.pool_transfers or []: + entry = self.mem_pool_host.entry_map.get(t.name) + if entry is None: + continue + if t.indices_from_pool is not None: + num_slots = source_len.get(t.indices_from_pool, 0) + else: + num_slots = len(t.host_indices) if t.host_indices is not None else 0 + num_bytes += num_slots * entry.host_pool.size_per_token + return num_bytes def load( self, @@ -542,6 +589,8 @@ def start_loading(self) -> int: op.node_ids, num_tokens=len(op.device_indices), timing_enabled=timing_enabled, + num_tokens_by_pool=self._num_tokens_by_pool(op), + num_bytes=self._transfer_num_bytes(op), ) ) return producer_id diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index ddce09704aa9..2e023eda0c6f 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -72,6 +72,7 @@ from sglang.srt.session.streaming_session import StreamingSession if TYPE_CHECKING: + from sglang.srt.managers.cache_controller import HiCacheAck from sglang.srt.managers.schedule_batch import Req from sglang.srt.mem_cache.cache_init_params import CacheInitParams from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import ( @@ -82,6 +83,14 @@ T = TypeVar("T") +# Metric label per component, matching the host pool names used by +# hicache_backup_tokens_total and the host occupancy gauges. +_COMPONENT_POOL_LABEL = { + ComponentType.FULL: PoolName.KV.value, + ComponentType.SWA: PoolName.SWA.value, + ComponentType.MAMBA: PoolName.MAMBA.value, +} + COMPONENT_REGISTRY: dict[ComponentType, type[TreeComponent]] = { ComponentType.FULL: FullComponent, @@ -360,6 +369,14 @@ def init_hicache(self, server_args: ServerArgs, params: CacheInitParams) -> None self.cache_controller is not None and self.cache_controller.write_policy == "write_back" ) + # Pre-seed the dropped-tokens series at 0 per pool + if self.metrics_collector is not None and self.cache_controller is not None: + for ct in self.tree_components: + self.metrics_collector.increment_dropped_tokens( + num_tokens=0, + reason="host_pressure", + pool=_COMPONENT_POOL_LABEL[ct], + ) self.load_back_threshold = 10 self.prefetch_stop_policy = server_args.hicache_storage_prefetch_policy @@ -436,7 +453,8 @@ def evict(self, params: EvictParams) -> EvictResult: ): self.writing_check(write_back=True) - self.update_eviction_metrics(sum(tracker.values()), start_time) + # Report full-layer tokens only + self.update_eviction_metrics(tracker[BASE_COMPONENT_TYPE], start_time) return EvictResult( num_tokens_evicted=tracker[BASE_COMPONENT_TYPE], swa_num_tokens_evicted=tracker.get(ComponentType.SWA, 0), @@ -519,10 +537,12 @@ def _evict_components( written = self._execute_and_commit_kv_backup( backup_kv, write_back=True ) + freed_before_drop = dict(tracker) if written > 0: self.writing_check(write_back=True) self._demote(node_id, tracker) elif self._drop_subtree_no_host(node_id, tracker): + self._record_dropped_tokens(tracker, freed_before_drop) logger.warning( "write_back: KV subtree dropped without backup " "due to host memory pressure, root node %d", @@ -539,6 +559,23 @@ def _evict_components( finally: self.tree_core.evict_device_end(ct) + def _record_dropped_tokens( + self, + tracker: dict[ComponentType, int], + freed_before_drop: dict[ComponentType, int], + ) -> None: + """Record per-pool tokens dropped without backup under host pressure.""" + if self.metrics_collector is None: + return + for ct, freed in tracker.items(): + dropped = freed - freed_before_drop[ct] + if dropped > 0: + self.metrics_collector.increment_dropped_tokens( + num_tokens=dropped, + reason="host_pressure", + pool=_COMPONENT_POOL_LABEL[ct], + ) + def inc_lock_ref( self, node_id: NodeId, skip_lock_components: Sequence[ComponentType] = () ) -> IncLockRefResult: @@ -1845,6 +1882,7 @@ def writing_check( for ack_id in ack.node_ids: if ack_id in self.ongoing_write_through: self._finish_write_through_ack(ack_id) + self._log_write_ack_metrics(ack) cc.ack_write_queue.clear() assert len(self.ongoing_write_through) == 0 return @@ -1867,8 +1905,24 @@ def writing_check( ack.finish_event.synchronize() for ack_id in ack.node_ids: self._finish_write_through_ack(ack_id) + self._log_write_ack_metrics(ack) finish_count -= 1 + def _log_write_ack_metrics(self, ack: HiCacheAck) -> None: + """Record D->H backup volume and duration for a completed write ack.""" + if self.metrics_collector is None: + return + for pool, num_tokens in (ack.num_tokens_by_pool or {}).items(): + if num_tokens > 0: + self.metrics_collector.increment_backup_num_tokens( + num_tokens=num_tokens, pool=pool + ) + if ack.num_bytes > 0: + self.metrics_collector.increment_backup_num_bytes(ack.num_bytes) + if ack.timing_enabled: + duration_ms = ack.start_event.elapsed_time(ack.finish_event) + self.metrics_collector.observe_backup_duration(duration_ms / 1000.0) + def loading_check(self, finish_count: Optional[int] = None) -> None: """Poll load-back completions.""" cc = self.cache_controller @@ -1903,7 +1957,13 @@ def loading_check(self, finish_count: Optional[int] = None) -> None: self.tree_core.finish_load_back(node) if self.metrics_collector is not None: - self.metrics_collector.increment_load_back_num_tokens(ack.num_tokens) + for pool, num_tokens in (ack.num_tokens_by_pool or {}).items(): + if num_tokens > 0: + self.metrics_collector.increment_load_back_num_tokens( + num_tokens=num_tokens, pool=pool + ) + if ack.num_bytes > 0: + self.metrics_collector.increment_load_back_num_bytes(ack.num_bytes) if ack.timing_enabled: duration_ms = ack.start_event.elapsed_time(ack.finish_event) self.metrics_collector.observe_load_back_duration( diff --git a/python/sglang/srt/observability/metrics_collector.py b/python/sglang/srt/observability/metrics_collector.py index d08ce2affae3..a5976c120525 100644 --- a/python/sglang/srt/observability/metrics_collector.py +++ b/python/sglang/srt/observability/metrics_collector.py @@ -889,6 +889,20 @@ def __init__( ), labelnames=list(labels.keys()) + ["mode"], ) + self.prefill_effective_tokens_total = Counter( + name="sglang:prefill_effective_tokens_total", + documentation=( + "Effective prefill tokens with retracted-request re-counts " + "excluded, updated on each log interval. mode: device_hit, " + "host_hit, storage_hit, input. Windowed prefix cache hit " + "rate = rate(sum of *_hit) / rate(sum of all modes); " + "per-tier rate uses a single *_hit mode in the numerator." + ), + labelnames=list(labels.keys()) + ["mode"], + ) + # Pre-seed every mode at 0 so per-tier ratio charts get a complete operand set + for mode in ("input", "device_hit", "host_hit", "storage_hit"): + self.prefill_effective_tokens_total.labels(**labels, mode=mode) self.forward_execution_seconds_total = Counter( name="sglang:forward_execution_seconds_total", documentation=( @@ -1248,6 +1262,24 @@ def increment_realtime_tokens( **dp_cooperation_info.to_labels(), ).inc(delta) + def increment_effective_prefill_tokens( + self, + input_tokens: int, + device_hit_tokens: int, + host_hit_tokens: int, + storage_hit_tokens: int, + ) -> None: + for mode, delta in [ + ("input", input_tokens), + ("device_hit", device_hit_tokens), + ("host_hit", host_hit_tokens), + ("storage_hit", storage_hit_tokens), + ]: + if delta > 0: + self.prefill_effective_tokens_total.labels( + **self.labels, mode=mode + ).inc(delta) + def increment_forward_execution_seconds( self, category: str, @@ -1964,6 +1996,11 @@ def __init__( 0.2, 0.5, 1.0, + 2.0, + 5.0, + 10.0, + 30.0, + 60.0, ] bucket_load_back_duration = get_histogram_conf_from_env( "SGLANG_BUCKET_LOAD_BACK_DURATION" @@ -1989,16 +2026,43 @@ def __init__( 0.5, 1.0, ] + # D->H backups include blocking merged ops issued during eviction under + # --hicache-write-policy write_back, which can run for seconds -- hence + # the wider default range than load-back. + bucket_backup_duration = [ + 0.001, + 0.002, + 0.005, + 0.01, + 0.02, + 0.05, + 0.1, + 0.2, + 0.5, + 1.0, + 2.0, + 5.0, + 10.0, + 30.0, + 60.0, + ] + self.eviction_duration_seconds = Histogram( name="sglang:eviction_duration_seconds", - documentation="Time taken to evict memory from GPU to CPU in seconds.", + documentation="End-to-end time of a device eviction pass in " + "seconds; under --hicache-write-policy write_back this includes " + "the blocking D->H backup (see " + "sglang:hicache_backup_duration_seconds for the copy alone).", labelnames=labels.keys(), buckets=bucket_eviction_duration, ) self.eviction_num_tokens = Counter( name="sglang:evicted_tokens_total", - documentation="The number of tokens evicted from GPU to CPU.", + documentation="The number of device KV token slots freed by " + "eviction, regardless of whether the data was backed up to host " + "(see sglang:hicache_backup_tokens_total) or destroyed (see " + "sglang:hicache_dropped_tokens_total).", labelnames=labels.keys(), ) @@ -2011,15 +2075,63 @@ def __init__( self.load_back_num_tokens = Counter( name="sglang:load_back_tokens_total", - documentation="The number of tokens loaded from CPU to GPU.", + documentation="The number of tokens loaded back from local host " + "DRAM (L2) to GPU, by host pool (kv, swa, mamba, ...).", + labelnames=list(labels.keys()) + ["pool"], + ) + + self.backup_duration_seconds = Histogram( + name="sglang:hicache_backup_duration_seconds", + documentation="Time taken to back up KV cache from GPU to local " + "host DRAM (L2) in seconds, per merged write op. Covers all D->H " + "backups regardless of --hicache-write-policy. Distinct from the " + "host-to-storage (L3) sglang:backuped_tokens_total.", + labelnames=labels.keys(), + buckets=bucket_backup_duration, + ) + + self.backup_num_bytes = Counter( + name="sglang:hicache_backup_bytes_total", + documentation="Bytes backed up from GPU to local host DRAM (L2), " + "all pools combined, including draft/sidecar transfers that the " + "token counter excludes. Divided by the rate of " + "hicache_backup_duration_seconds_sum, gives the achieved D->H " + "bandwidth while transferring.", labelnames=labels.keys(), ) + self.load_back_num_bytes = Counter( + name="sglang:load_back_bytes_total", + documentation="Bytes loaded back from local host DRAM (L2) to " + "GPU, all pools combined, including draft/sidecar transfers that " + "the token counter excludes. Divided by the rate of " + "load_back_duration_seconds_sum, gives the achieved H->D " + "bandwidth while transferring.", + labelnames=labels.keys(), + ) + + self.backup_num_tokens = Counter( + name="sglang:hicache_backup_tokens_total", + documentation="The number of tokens backed up from GPU to local " + "host DRAM (L2), by host pool (kv, swa, mamba, ...). Covers all " + "D->H backups regardless of --hicache-write-policy. Distinct from " + "the host-to-storage (L3) sglang:backuped_tokens_total.", + labelnames=list(labels.keys()) + ["pool"], + ) + + self.hicache_dropped_tokens = Counter( + name="sglang:hicache_dropped_tokens_total", + documentation="The number of device KV tokens destroyed without a " + "host backup, by pool (kv, swa, ...) and reason (e.g. write-back " + "backup failure under host memory pressure).", + labelnames=list(labels.keys()) + ["reason", "pool"], + ) + def increment_eviction_num_tokens(self, num_tokens: int) -> None: self.eviction_num_tokens.labels(**self.labels).inc(num_tokens) - def increment_load_back_num_tokens(self, num_tokens: int) -> None: - self.load_back_num_tokens.labels(**self.labels).inc(num_tokens) + def increment_load_back_num_tokens(self, num_tokens: int, pool: str) -> None: + self.load_back_num_tokens.labels(**self.labels, pool=pool).inc(num_tokens) def observe_eviction_duration(self, duration_seconds: float) -> None: self.eviction_duration_seconds.labels(**self.labels).observe(duration_seconds) @@ -2027,6 +2139,23 @@ def observe_eviction_duration(self, duration_seconds: float) -> None: def observe_load_back_duration(self, duration_seconds: float) -> None: self.load_back_duration_seconds.labels(**self.labels).observe(duration_seconds) + def increment_backup_num_tokens(self, num_tokens: int, pool: str) -> None: + self.backup_num_tokens.labels(**self.labels, pool=pool).inc(num_tokens) + + def increment_backup_num_bytes(self, num_bytes: int) -> None: + self.backup_num_bytes.labels(**self.labels).inc(num_bytes) + + def increment_load_back_num_bytes(self, num_bytes: int) -> None: + self.load_back_num_bytes.labels(**self.labels).inc(num_bytes) + + def observe_backup_duration(self, duration_seconds: float) -> None: + self.backup_duration_seconds.labels(**self.labels).observe(duration_seconds) + + def increment_dropped_tokens(self, num_tokens: int, reason: str, pool: str) -> None: + self.hicache_dropped_tokens.labels(**self.labels, reason=reason, pool=pool).inc( + num_tokens + ) + class EncoderMetricsCollector(_StatLoggerDIMixin): """Metrics collector for the EPD encoder server (--encoder-only).""" diff --git a/test/registered/unit/managers/test_prefill_adder.py b/test/registered/unit/managers/test_prefill_adder.py index c88075581e0a..0e240ed0f764 100644 --- a/test/registered/unit/managers/test_prefill_adder.py +++ b/test/registered/unit/managers/test_prefill_adder.py @@ -97,6 +97,8 @@ def create_mock_req(self, rid, priority, max_new_tokens, output_len=0, wait_time req.sampling_params = SimpleNamespace(max_new_tokens=max_new_tokens) req.time_stats = SimpleNamespace(wait_queue_entry_time=wait_time) req.retracted_stain = False + req.host_hit_length = 0 + req.storage_hit_length = 0 req.finished.return_value = False req.needs_host_load_back.return_value = False return req diff --git a/test/registered/unit/mem_cache/test_hicache_load_back_timing.py b/test/registered/unit/mem_cache/test_hicache_load_back_timing.py index e056a07a56ee..b92e63475b32 100644 --- a/test/registered/unit/mem_cache/test_hicache_load_back_timing.py +++ b/test/registered/unit/mem_cache/test_hicache_load_back_timing.py @@ -65,6 +65,7 @@ def test_loading_check_observes_duration_and_tokens(self): node_ids=[1, 2], num_tokens=1024, timing_enabled=True, + num_tokens_by_pool={"kv": 1024}, ) stub = object.__new__(HiRadixCache) stub.cache_controller = SimpleNamespace(ack_load_queue=[ack]) @@ -77,7 +78,7 @@ def test_loading_check_observes_duration_and_tokens(self): stub.loading_check() stub.metrics_collector.increment_load_back_num_tokens.assert_called_once_with( - 1024 + num_tokens=1024, pool="kv" ) stub.metrics_collector.observe_load_back_duration.assert_called_once() (observed,), _ = stub.metrics_collector.observe_load_back_duration.call_args @@ -100,6 +101,7 @@ def test_loading_check_fallback_when_timing_unsupported(self): node_ids=[7], num_tokens=512, timing_enabled=False, + num_tokens_by_pool={"kv": 512}, ) stub = object.__new__(HiRadixCache) stub.cache_controller = SimpleNamespace(ack_load_queue=[ack]) @@ -112,7 +114,7 @@ def test_loading_check_fallback_when_timing_unsupported(self): stub.loading_check() stub.metrics_collector.increment_load_back_num_tokens.assert_called_once_with( - 512 + num_tokens=512, pool="kv" ) stub.metrics_collector.observe_load_back_duration.assert_not_called() self.assertEqual(stub.cache_controller.ack_load_queue, []) diff --git a/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py b/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py index 8c79c3565752..57c0353bd6ba 100644 --- a/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py +++ b/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py @@ -139,6 +139,9 @@ def _cpu_per_layer_pf_lf_copy( class _FakeEvent: + def __init__(self, enable_timing=False): + self.enable_timing = enable_timing + def record(self): pass @@ -156,6 +159,15 @@ def stream(stream): class TestHiCacheStagedWriteBackDispatch(unittest.TestCase): + def setUp(self): + # start_writing probes timing support via a module-cached check; + # clear it on both sides so results from (or against) the fake + # device module never leak across tests. + manager_cache_controller._timing_events_supported.cache_clear() + + def tearDown(self): + manager_cache_controller._timing_events_supported.cache_clear() + def _patched_transfers(self, src_registry=None, module=MEMORY_POOL_HOST_MODULE): staged_side_effect = None if src_registry is not None: @@ -678,6 +690,10 @@ def test_write_back_jit_hybrid_write_keeps_extra_host_indices_on_cpu(self): class FakeHostGroup: layout = "page_first" can_use_write_back_jit = True + anchor_entry = SimpleNamespace( + name=PoolName.KV, host_pool=SimpleNamespace(size_per_token=2) + ) + entry_map = {} def backup_from_device_all_layer( self, @@ -718,8 +734,13 @@ def backup_from_device_all_layer( ) ) - with mock.patch.object( - hybrid_cache_controller, "device_module", _FakeDeviceModule + with ( + mock.patch.object( + hybrid_cache_controller, "device_module", _FakeDeviceModule + ), + mock.patch.object( + manager_cache_controller, "device_module", _FakeDeviceModule + ), ): controller.start_writing() @@ -733,6 +754,10 @@ def test_hybrid_write_moves_indices_without_write_back_jit(self): class FakeHostGroup: layout = "page_first" can_use_write_back_jit = False + anchor_entry = SimpleNamespace( + name=PoolName.KV, host_pool=SimpleNamespace(size_per_token=2) + ) + entry_map = {} def backup_from_device_all_layer( self, @@ -770,8 +795,13 @@ def backup_from_device_all_layer( return_value=(op.host_indices, op.device_indices, op.pool_transfers) ) - with mock.patch.object( - hybrid_cache_controller, "device_module", _FakeDeviceModule + with ( + mock.patch.object( + hybrid_cache_controller, "device_module", _FakeDeviceModule + ), + mock.patch.object( + manager_cache_controller, "device_module", _FakeDeviceModule + ), ): controller.start_writing() @@ -785,6 +815,7 @@ def test_write_back_jit_cache_controller_keeps_host_indices_on_cpu(self): class FakeHostPool: layout = "page_first" can_use_write_back_jit = True + size_per_token = 2 def backup_from_device_all_layer( self, device_pool, host_indices, device_indices, io_backend @@ -825,6 +856,7 @@ def test_cache_controller_moves_indices_without_write_back_jit(self): class FakeHostPool: layout = "page_first" can_use_write_back_jit = False + size_per_token = 2 def backup_from_device_all_layer( self, device_pool, host_indices, device_indices, io_backend From 70f3127b8c3d72892dfbdeda38a936c5aa7a2e66 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Wed, 12 Aug 2026 20:50:46 -0700 Subject: [PATCH 10/12] [Fix] Snapshot `req.prefix_indices` when the prefix cache is disabled (#34644) (cherry picked from commit 26627e999ded02398a7337697e7281c9b8c567fa) (cherry picked from commit 53668167ff1334028affe0ffe47423fccbed7b59) --- python/sglang/srt/mem_cache/swa_radix_cache.py | 2 +- python/sglang/srt/mem_cache/unified_radix_cache.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/mem_cache/swa_radix_cache.py b/python/sglang/srt/mem_cache/swa_radix_cache.py index 31cedca763b9..98d945c3034b 100644 --- a/python/sglang/srt/mem_cache/swa_radix_cache.py +++ b/python/sglang/srt/mem_cache/swa_radix_cache.py @@ -514,7 +514,7 @@ def cache_unfinished_req(self, req: Req, chunked=False) -> None: ] # `req.prefix_indices` will be used in `PrefillAdder::add_chunked_req` later - req.prefix_indices = kv_indices + req.prefix_indices = kv_indices.to(dtype=torch.int64, copy=True) return token_ids = req.get_fill_ids() diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 2e023eda0c6f..2d5cc13a0126 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -736,7 +736,7 @@ def cache_unfinished_req(self, req: Req, chunked: bool = False, **kwargs) -> Non kv_indices = self.req_to_token_pool.req_to_token[ req.req_pool_idx, : len(token_ids) ] - req.prefix_indices = kv_indices + req.prefix_indices = kv_indices.to(dtype=torch.int64, copy=True) return kv_indices_orig = self.req_to_token_pool.req_to_token[ From 7d42fa13bb2967cdce23d1a633dba44ed8fbf62a Mon Sep 17 00:00:00 2001 From: Ke Bao Date: Thu, 13 Aug 2026 02:15:36 +0800 Subject: [PATCH 11/12] Add bit-exact unified radix cache KL test for hybrid SWA + mamba (#34607) (cherry picked from commit 3974b00359776e24aa257031b951ce8841c9b64a) (cherry picked from commit 2a7a56d3df71becd7a031a390a6c4502ba34342f) --- ..._unified_radix_cache_kl_hybrid_bitexact.py | 251 ++++++++++++++++++ 1 file changed, 251 insertions(+) create mode 100644 test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_hybrid_bitexact.py diff --git a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_hybrid_bitexact.py b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_hybrid_bitexact.py new file mode 100644 index 000000000000..e26b873a2e45 --- /dev/null +++ b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_hybrid_bitexact.py @@ -0,0 +1,251 @@ +"""Bit-exact KL guards for the unified radix cache on a hybrid SWA + mamba model. + +The other KL tests in this directory gate on a loose threshold because their +models cannot score a token identically twice: Qwen3-Next's chunkwise prefill +scan and its decode recurrence are different algorithms and land an ulp apart, so +a tight floor there would fail on float noise. The shrunken Inkling checkpoint +reproduces every logprob exactly under deterministic inference, which turns the +same comparison into an exact one -- any nonzero KL is a state-reuse bug. It also +fits on one GPU, so these run per-commit rather than on a 4-GPU stage. + +Each class below reproduces a specific merged regression when its fix is +reverted; the measured pre-fix divergence is recorded in the class docstring so a +later threshold change has to argue with a number. + +These classes do not use UnifiedRadixTreeTestMixin: it bundles gsm8k and mmlu, +which an undertrained checkpoint cannot gate on, and each class here runs the +harness its regression was actually reproduced with. + +The imported `test_`-prefixed helpers are aliased so pytest does not collect them +as tests. +""" + +import os +import random +import unittest + +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kl_multiturn_utils import ( + make_mamba_decode_assert, +) +from sglang.test.kl_multiturn_utils import ( + test_input_output_logprobs_match_decode_cache_hit_helper as assert_multiturn_decode_cache_hit, +) +from sglang.test.kl_test_utils import ( + get_input_ids, +) +from sglang.test.kl_test_utils import ( + test_input_output_logprobs_match_decode_cache_hit_helper as assert_decode_cache_hit, +) +from sglang.test.kl_test_utils import ( + test_input_output_logprobs_match_helper as assert_logprobs_match, +) +from sglang.test.kl_test_utils import ( + test_input_output_logprobs_match_prefill_cache_hit_helper as assert_prefill_cache_hit, +) +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + +register_cuda_ci(est_time=600, stage="base-b", runner_config="1-gpu-large") + +_MODEL_PATH = os.environ.get("INKLING_TEST_MODEL_PATH", "thinkingmachines/Inkling") +_MODEL_REVISION = os.environ.get("INKLING_TEST_MODEL_REVISION", "test") + +# Both classes measure exactly 0 in their fixed state -- every logprob matches bit +# for bit. The floor only keeps a stray ulp from failing the run; a state-reuse +# bug lands orders of magnitude above it. It cannot be 0.0: the comparison is a +# strict `<`, so an exact 0 would fail its own threshold. +KL_DIV_THRESHOLD = 1e-9 + +# Equal to the page size below. Out-of-window SWA slots are freed a page at a +# time, so only a checkpoint sitting on a page boundary still has a full window of +# SWA data below it -- at the default 256 half the sequence lengths land off that +# boundary and lose their decode prefix entirely. +TRACK_INTERVAL = 128 +PAGE_SIZE = 128 + +# Past the 512-token sliding window, so decode carries the window through the +# handover from prompt tokens to generated ones. +MAX_NEW_TOKENS = 1024 + + +def _random_suffixes(n: int, length: int, seed: int) -> list[list[int]]: + rng = random.Random(seed) + return [[rng.randint(1, 30000) for _ in range(length)] for _ in range(n)] + + +def _base_args() -> list[str]: + return [ + "--trust-remote-code", + "--attention-backend", + "fa4", + "--page-size", + str(PAGE_SIZE), + "--mamba-radix-cache-strategy", + "extra_buffer", + "--swa-full-tokens-ratio", + "0.1", + "--mamba-full-memory-ratio", + "0.1", + # 0.85 was carried over from the 4-GPU B200 test and OOMs an 80 GB card: + # the static pool leaves ~19 GB for the prefill graphs, the fa4 workspace + # and the chunked-prefill activations, which is what this config needs. + "--mem-fraction-static", + "0.6", + "--mamba-track-interval", + str(TRACK_INTERVAL), + "--enable-deterministic-inference", + ] + + +class TestUnifiedHybridBitExact(CustomTestCase): + """Prefill and decode must score a token identically once every kernel on the + path is batch-invariant, so any drift is a stale conv/mamba checkpoint or a + prefix the cache restored wrong. + + Guards #34184 (stale track rows corrupting conv checkpoints under the prefill + graph). Reverting that fix here measures avg_kl_div 5.58e-07 on + test_logprobs_match and 6.22e-06 on test_prefill_cache_hit, against 0.0 with + it in place. test_decode_cache_hit is 0.0 either way -- it guards decode-region + state reuse in general, not that regression. + """ + + @classmethod + def setUpClass(cls): + cls.model = _MODEL_PATH + cls.base_url = DEFAULT_URL_FOR_TEST + other_args = _base_args() + [ + # Pinned, not incidental: the prefill graph derives its fixed + # request-slot count from this (chunked_prefill_size // 512), and those + # slots are exactly what #34184 left stale. Lowering it shrinks the + # sentinel tail and the guard stops firing while still passing. + "--chunked-prefill-size", + "16384", + ] + if _MODEL_REVISION: + other_args += ["--revision", _MODEL_REVISION] + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=other_args, + env={**os.environ, "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"}, + ) + + @classmethod + def tearDownClass(cls): + if getattr(cls, "process", None) is not None: + kill_process_tree(cls.process.pid) + + def _run(self, helper): + helper( + self.base_url, + {self.model: {"kl_div": KL_DIV_THRESHOLD}}, + self.model, + max_samples=32, + max_new_tokens=MAX_NEW_TOKENS, + trust_remote_code=True, + ) + + def test_logprobs_match(self): + self._run(assert_logprobs_match) + + def test_prefill_cache_hit(self): + self._run(assert_prefill_cache_hit) + + def test_decode_cache_hit(self): + self._run(assert_decode_cache_hit) + + +class TestUnifiedHybridHiCacheBitExact(CustomTestCase): + """Same exactness bar with the host tier in the loop, over interleaved + branches so hits land at many prefix lengths rather than one aligned one. + + Guards #29792 (decode track save picking its slot from the producer-side + pointer). Without that fix this measures avg_kl_div 9.43e-06 and 1.16e-05 over + two rounds, with 3 of 9 samples dirty and the rest exactly 0; with it in place + both rounds are 0.0. + + Runs the multi-turn branching harness because the single-turn helpers above + cannot produce a non-aligned hit length, which this regression needs. + """ + + @classmethod + def setUpClass(cls): + cls.model = _MODEL_PATH + cls.base_url = DEFAULT_URL_FOR_TEST + other_args = _base_args() + [ + "--enable-hierarchical-cache", + "--hicache-ratio", + "4", + "--hicache-write-policy", + "write_through", + "--hicache-io-backend", + "direct", + # The mamba host pool only supports page_first and page_first_direct. + "--hicache-mem-layout", + "page_first_direct", + # Tight pools and a small budget so decode crosses a track boundary and + # the host tier is actually exercised instead of everything staying + # resident on device. + "--chunked-prefill-size", + "2048", + "--max-total-tokens", + "65536", + "--max-mamba-cache-size", + "500", + "--max-running-requests", + "4", + ] + if _MODEL_REVISION: + other_args += ["--revision", _MODEL_REVISION] + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=other_args, + env={**os.environ, "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"}, + ) + cls.input_ids = get_input_ids( + tokenizer_path=cls.model, num_samples=9, trust_remote_code=True + ) + + @classmethod + def tearDownClass(cls): + if getattr(cls, "process", None) is not None: + kill_process_tree(cls.process.pid) + + def test_multiturn_decode_cache_hit_branching(self): + groups, branches = 3, 3 + n = groups * branches + first_turn = [] + for g in range(groups): + base = self.input_ids[g][:512] + for _ in range(branches): + first_turn.append(list(base)) + + assert_multiturn_decode_cache_hit( + self.base_url, + self.model, + KL_DIV_THRESHOLD, + first_turn, + turn_suffixes=[ + _random_suffixes(n, 512, seed=300), + _random_suffixes(n, 256, seed=400), + ], + # Not the default exact equality: a mamba checkpoint lands on a track + # boundary, so the reusable prefix is floor-aligned to the interval. + assert_decode_cached_tokens=make_mamba_decode_assert(TRACK_INTERVAL), + branches_per_group=branches, + max_new_tokens=512, + sampling_temperature=0, + ) + + +if __name__ == "__main__": + unittest.main() From 1b131d1ea8d8d929a26cf211cde40f2709da9234 Mon Sep 17 00:00:00 2001 From: "Po-Han Huang (NVIDIA)" <53919306+nvpohanh@users.noreply.github.com> Date: Thu, 20 Aug 2026 03:01:36 +0800 Subject: [PATCH 12/12] [Scheduler] Add configurable decode interval after prefill (#35017) (cherry picked from commit c34e0793a66eda02e87c36a247389c9457ef369f) --- python/sglang/srt/managers/scheduler.py | 27 +++++++ python/sglang/srt/server_args.py | 10 +++ .../test_scheduler_chunked_req_gate.py | 2 + .../test_scheduler_prefill_decode_interval.py | 75 +++++++++++++++++++ .../unit/server_args/test_server_args.py | 9 +++ 5 files changed, 123 insertions(+) create mode 100644 test/registered/unit/managers/test_scheduler_prefill_decode_interval.py diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index a9615a14c512..31ab440ddb78 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -1135,6 +1135,8 @@ def init_running_status(self): def init_chunked_prefill(self): self.chunked_prefill_size = get_schedule().chunked_prefill_size + self.prefill_decode_interval = get_schedule().prefill_decode_interval + self._prefill_decode_interval_remaining = 0 uses_transformers_backend = ( get_resolved_model_impl(self.model_config) == ModelImpl.TRANSFORMERS ) @@ -1171,6 +1173,28 @@ def init_chunked_prefill(self): ) self.enable_dynamic_chunking = False + def _should_defer_prefill(self) -> bool: + if self._prefill_decode_interval_remaining == 0: + return False + + self._prefill_decode_interval_remaining -= 1 + return True + + def _arm_prefill_decode_interval(self, batch: Optional[ScheduleBatch]) -> None: + if self.prefill_decode_interval == 0 or batch is None: + return + + # DP attention synchronizes this flag across ranks. This keeps every + # rank on the same prefill/decode cadence even when only one rank has + # local prefill work. Non-DP scheduling can use the local mode directly. + is_extend = ( + batch.is_extend_in_batch + if self.require_mlp_sync + else batch.forward_mode.is_extend() + ) + if is_extend: + self._prefill_decode_interval_remaining = self.prefill_decode_interval + def init_metrics_reporter( self, tp_rank: int, pp_rank: int, dp_rank: Optional[int] ) -> None: @@ -3013,6 +3037,8 @@ def get_next_batch_to_run( if self.dllm_config is not None: new_batch = self.get_new_batch_dllm(running_batch) + elif self._should_defer_prefill(): + new_batch = None else: prefill_plan = self.get_new_batch_prefill(running_batch) new_batch = prefill_plan.batch_to_run @@ -3046,6 +3072,7 @@ def get_next_batch_to_run( ret = self.dp_attn_adapter.maybe_prepare_mlp_sync_batch( ret, need_sync=need_mlp_sync ) + self._arm_prefill_decode_interval(ret) # Handle ngram embedding ret = self.ngram_embedding_manager.prepare_for_forward( diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 18664959e959..feb3f4dff045 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -794,6 +794,11 @@ class ServerArgs: "The maximum number of tokens in a chunk for the chunked prefill. Setting this to -1 means disabling chunked prefill.", NS("schedule"), ] = None + prefill_decode_interval: A[ + int, + "The number of decode rounds to run after a prefill batch before scheduling the next prefill. In data-parallel attention mode, the interval is synchronized across all DP ranks. Set to 0 to disable.", + NS("schedule"), + ] = 0 enable_dynamic_chunking: A[ bool, "Enable dynamic chunk size adjustment for pipeline parallelism. When enabled, chunk sizes are dynamically calculated based on fitted function to maintain consistent execution time across chunks.", @@ -3463,6 +3468,7 @@ def __post_init__(self): self._resolved_overrides = [] self._handle_return_hidden_states_mode() + self._validate_prefill_decode_interval() if self.model_path.lower() in ["none", "dummy"]: return @@ -8089,6 +8095,10 @@ def _handle_asr_validation(self): f"(got {self.asr_max_concurrent_sessions})." ) + def _validate_prefill_decode_interval(self): + if self.prefill_decode_interval < 0: + raise ValueError("--prefill-decode-interval must be non-negative.") + def _handle_other_validations(self): # Handle optimistic prefill validation diff --git a/test/registered/unit/managers/test_scheduler_chunked_req_gate.py b/test/registered/unit/managers/test_scheduler_chunked_req_gate.py index 14ed036ec46b..4911f207ce79 100644 --- a/test/registered/unit/managers/test_scheduler_chunked_req_gate.py +++ b/test/registered/unit/managers/test_scheduler_chunked_req_gate.py @@ -90,6 +90,8 @@ def _scheduler_for_get_next_batch(*, tree_cache, chunked_req) -> Scheduler: s.running_batch.is_prefill_only = False s.running_batch.batch_is_full = False s.running_batch.reqs = [] + s.prefill_decode_interval = 0 + s._prefill_decode_interval_remaining = 0 s.get_new_batch_prefill = MagicMock( return_value=NextBatchPlan(batch_to_run=None, running_batch=s.running_batch) ) diff --git a/test/registered/unit/managers/test_scheduler_prefill_decode_interval.py b/test/registered/unit/managers/test_scheduler_prefill_decode_interval.py new file mode 100644 index 000000000000..8b0d683c43b4 --- /dev/null +++ b/test/registered/unit/managers/test_scheduler_prefill_decode_interval.py @@ -0,0 +1,75 @@ +"""Tests for scheduler prefill/decode interleaving.""" + +import unittest +from types import SimpleNamespace + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import maybe_stub_sgl_kernel + +maybe_stub_sgl_kernel() + +from sglang.srt.managers.scheduler import Scheduler + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + + +def _make_scheduler(*, interval: int, require_mlp_sync: bool) -> Scheduler: + scheduler = Scheduler.__new__(Scheduler) + scheduler.prefill_decode_interval = interval + scheduler._prefill_decode_interval_remaining = 0 + scheduler.require_mlp_sync = require_mlp_sync + return scheduler + + +def _make_batch(*, local_extend: bool, global_extend: bool): + return SimpleNamespace( + forward_mode=SimpleNamespace(is_extend=lambda: local_extend), + is_extend_in_batch=global_extend, + ) + + +class TestPrefillDecodeInterval(unittest.TestCase): + def test_disabled_interval_does_not_arm(self): + scheduler = _make_scheduler(interval=0, require_mlp_sync=False) + + scheduler._arm_prefill_decode_interval( + _make_batch(local_extend=True, global_extend=False) + ) + + self.assertFalse(scheduler._should_defer_prefill()) + + def test_non_dp_interval_uses_local_forward_mode(self): + scheduler = _make_scheduler(interval=2, require_mlp_sync=False) + + scheduler._arm_prefill_decode_interval( + _make_batch(local_extend=True, global_extend=False) + ) + + self.assertTrue(scheduler._should_defer_prefill()) + self.assertTrue(scheduler._should_defer_prefill()) + self.assertFalse(scheduler._should_defer_prefill()) + + def test_dp_interval_uses_globally_synchronized_extend_flag(self): + scheduler = _make_scheduler(interval=2, require_mlp_sync=True) + + # This rank is locally decoding, but another DP rank is prefilling. + scheduler._arm_prefill_decode_interval( + _make_batch(local_extend=False, global_extend=True) + ) + + self.assertEqual(scheduler._prefill_decode_interval_remaining, 2) + self.assertTrue(scheduler._should_defer_prefill()) + + def test_decode_batch_does_not_rearm_interval(self): + scheduler = _make_scheduler(interval=2, require_mlp_sync=True) + scheduler._prefill_decode_interval_remaining = 1 + + scheduler._arm_prefill_decode_interval( + _make_batch(local_extend=False, global_extend=False) + ) + + self.assertEqual(scheduler._prefill_decode_interval_remaining, 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index a7102bc25aac..35f6ba463b55 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -42,6 +42,15 @@ class TestPrepareServerArgs(CustomTestCase): + def test_prefill_decode_interval(self): + args = ServerArgs(model_path="dummy", prefill_decode_interval=16) + self.assertEqual(args.prefill_decode_interval, 16) + + with self.assertRaisesRegex( + ValueError, "--prefill-decode-interval must be non-negative" + ): + ServerArgs(model_path="dummy", prefill_decode_interval=-1) + def test_return_hidden_states_mode_configuration(self): disabled = ServerArgs(model_path="dummy") self.assertFalse(disabled.enable_return_hidden_states)