diff --git a/python/sglang/srt/hardware_backend/mlx/kv_cache/auxiliary_state.py b/python/sglang/srt/hardware_backend/mlx/kv_cache/auxiliary_state.py index f75d8832c..9a45f49f2 100644 --- a/python/sglang/srt/hardware_backend/mlx/kv_cache/auxiliary_state.py +++ b/python/sglang/srt/hardware_backend/mlx/kv_cache/auxiliary_state.py @@ -17,10 +17,10 @@ import torch from sglang.srt.mem_cache.base_prefix_cache import EvictParams, InsertResult from sglang.srt.mem_cache.memory_pool import ReqToTokenPool -from sglang.srt.mem_cache.unified_cache.components.mamba_component import ( +from sglang.srt.mem_cache.unified_cache.components.base import TreeComponent +from sglang.srt.mem_cache.unified_cache.components.mamba import ( MambaComponent, ) -from sglang.srt.mem_cache.unified_cache.components.tree_component import TreeComponent _CACHE_ATTRS = ("offset", "lengths", "left_padding") _MISSING = object() diff --git a/python/sglang/srt/mem_cache/cache_init_params.py b/python/sglang/srt/mem_cache/cache_init_params.py index b3aa2a636..75c54a074 100644 --- a/python/sglang/srt/mem_cache/cache_init_params.py +++ b/python/sglang/srt/mem_cache/cache_init_params.py @@ -9,7 +9,7 @@ if TYPE_CHECKING: from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator from sglang.srt.mem_cache.memory_pool import ReqToTokenPool from sglang.srt.mem_cache.unified_cache.components import ComponentType - from sglang.srt.mem_cache.unified_cache.components.tree_component import ( + from sglang.srt.mem_cache.unified_cache.components.base import ( TreeComponent, ) 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 f7556ba44..f0b1c61b6 100644 --- a/python/sglang/srt/mem_cache/unified_cache/components/README.md +++ b/python/sglang/srt/mem_cache/unified_cache/components/README.md @@ -74,10 +74,10 @@ node.component_data[ComponentType.MAMBA] # MambaComponent data | `../unified_cache/unified_tree_core.py` | `UnifiedTreeCore` — the tree, LRUs, and size counters; `UnifiedTreeNode`, `UnifiedLRUList` | | `../unified_cache/unified_tree_core_interface.py` | `UnifiedTreeCoreInterface`, `NodeId` — the tree/cache boundary contract | | `../unified_cache/cache_action.py` | Deferred `CacheAction`/`ComponentAction` types emitted by the tree | -| `tree_component.py` | `TreeComponent` ABC, `ComponentType`, `ComponentData`, `get_and_increase_time_counter`, `next_component_uuid` | -| `full_component.py` | `FullComponent` — standard full-attention KV cache component | -| `swa_component.py` | `SWAComponent` — sliding-window attention component with tombstone/window tracking | -| `mamba_component.py` | `MambaComponent` — Mamba/SSM state component with copy-on-write | +| `base.py` | `TreeComponent` ABC, `ComponentType`, `ComponentData`, `get_and_increase_time_counter`, `next_component_uuid` | +| `full.py` | `FullComponent` — standard full-attention KV cache component | +| `swa.py` | `SWAComponent` — sliding-window attention component with tombstone/window tracking | +| `mamba.py` | `MambaComponent` — Mamba/SSM state component with copy-on-write | | `../hybrid_cache/hybrid_cache_controller.py` | `HybridCacheController` — HiCache multi-pool controller (L1 GPU → L2 CPU, optional L3 storage) | | `__init__.py` | Re-exports: `ComponentType`, `ComponentData`, `TreeComponent`, `FullComponent`, `SWAComponent`, `MambaComponent` | @@ -276,7 +276,7 @@ Cache an in-progress request's partial KV data (chunked prefill). ## TreeComponent Hook Reference -Each component implements these hooks. See `tree_component.py` for the ABC and docstrings. +Each component implements these hooks. See `base.py` for the ABC and docstrings. ### Match Phase diff --git a/python/sglang/srt/mem_cache/unified_cache/components/__init__.py b/python/sglang/srt/mem_cache/unified_cache/components/__init__.py index 763f515ec..147f57205 100644 --- a/python/sglang/srt/mem_cache/unified_cache/components/__init__.py +++ b/python/sglang/srt/mem_cache/unified_cache/components/__init__.py @@ -1,7 +1,4 @@ -from sglang.srt.mem_cache.unified_cache.components.full_component import FullComponent -from sglang.srt.mem_cache.unified_cache.components.mamba_component import MambaComponent -from sglang.srt.mem_cache.unified_cache.components.swa_component import SWAComponent -from sglang.srt.mem_cache.unified_cache.components.tree_component import ( +from sglang.srt.mem_cache.unified_cache.components.base import ( _NUM_COMPONENT_TYPES, BASE_COMPONENT_TYPE, CacheTransferPhase, @@ -17,6 +14,9 @@ from sglang.srt.mem_cache.unified_cache.components.tree_component import ( get_and_increase_time_counter, next_component_uuid, ) +from sglang.srt.mem_cache.unified_cache.components.full import FullComponent +from sglang.srt.mem_cache.unified_cache.components.mamba import MambaComponent +from sglang.srt.mem_cache.unified_cache.components.swa import SWAComponent __all__ = [ "BASE_COMPONENT_TYPE", diff --git a/python/sglang/srt/mem_cache/unified_cache/components/tree_component.py b/python/sglang/srt/mem_cache/unified_cache/components/base.py similarity index 100% rename from python/sglang/srt/mem_cache/unified_cache/components/tree_component.py rename to python/sglang/srt/mem_cache/unified_cache/components/base.py diff --git a/python/sglang/srt/mem_cache/unified_cache/components/full_component.py b/python/sglang/srt/mem_cache/unified_cache/components/full.py similarity index 99% rename from python/sglang/srt/mem_cache/unified_cache/components/full_component.py rename to python/sglang/srt/mem_cache/unified_cache/components/full.py index f38dd925e..640bceccb 100644 --- a/python/sglang/srt/mem_cache/unified_cache/components/full_component.py +++ b/python/sglang/srt/mem_cache/unified_cache/components/full.py @@ -19,7 +19,7 @@ from sglang.srt.mem_cache.hicache_storage import ( PoolTransferResult, ) from sglang.srt.mem_cache.unified_cache.cache_action import FreeComponentDeviceSlot -from sglang.srt.mem_cache.unified_cache.components.tree_component import ( +from sglang.srt.mem_cache.unified_cache.components.base import ( CacheTransferPhase, ComponentType, EvictLayer, diff --git a/python/sglang/srt/mem_cache/unified_cache/components/mamba_component.py b/python/sglang/srt/mem_cache/unified_cache/components/mamba.py similarity index 99% rename from python/sglang/srt/mem_cache/unified_cache/components/mamba_component.py rename to python/sglang/srt/mem_cache/unified_cache/components/mamba.py index d5d846760..70d7683de 100644 --- a/python/sglang/srt/mem_cache/unified_cache/components/mamba_component.py +++ b/python/sglang/srt/mem_cache/unified_cache/components/mamba.py @@ -25,7 +25,7 @@ from sglang.srt.mem_cache.unified_cache.cache_action import ( FreeComponentHostSlot, MambaEvictExcessPathStates, ) -from sglang.srt.mem_cache.unified_cache.components.tree_component import ( +from sglang.srt.mem_cache.unified_cache.components.base import ( CacheTransferPhase, ComponentType, EvictLayer, diff --git a/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py b/python/sglang/srt/mem_cache/unified_cache/components/swa.py similarity index 99% rename from python/sglang/srt/mem_cache/unified_cache/components/swa_component.py rename to python/sglang/srt/mem_cache/unified_cache/components/swa.py index 49f306458..31f92dd19 100644 --- a/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py +++ b/python/sglang/srt/mem_cache/unified_cache/components/swa.py @@ -30,7 +30,7 @@ from sglang.srt.mem_cache.unified_cache.cache_action import ( RecoverSWAWithLockedFull, SWARebuild, ) -from sglang.srt.mem_cache.unified_cache.components.tree_component import ( +from sglang.srt.mem_cache.unified_cache.components.base import ( BASE_COMPONENT_TYPE, CacheTransferPhase, ComponentType, diff --git a/python/sglang/srt/mem_cache/unified_cache/session_ref_tracker.py b/python/sglang/srt/mem_cache/unified_cache/session_ref_tracker.py index 07cc37248..406069aeb 100644 --- a/python/sglang/srt/mem_cache/unified_cache/session_ref_tracker.py +++ b/python/sglang/srt/mem_cache/unified_cache/session_ref_tracker.py @@ -12,7 +12,7 @@ from typing import TYPE_CHECKING, Optional if TYPE_CHECKING: from sglang.srt.managers.schedule_batch import Req - from sglang.srt.mem_cache.unified_cache.components.tree_component import ( + from sglang.srt.mem_cache.unified_cache.components.base import ( TreeComponent, ) from sglang.srt.mem_cache.unified_cache.unified_tree_core import UnifiedTreeCore diff --git a/rust/sglang-radix-tree/src/components/mod.rs b/rust/sglang-radix-tree/src/components/mod.rs index f098aeacf..cf7154fe7 100644 --- a/rust/sglang-radix-tree/src/components/mod.rs +++ b/rust/sglang-radix-tree/src/components/mod.rs @@ -76,7 +76,7 @@ pub trait TreeComponent { phase: LRURefreshPhase, node_id: NodeIdx_, ) { - // Python reference — tree_component.py::TreeComponent.refresh_lru: + // Python reference — base.py::TreeComponent.refresh_lru: // def refresh_lru( // self, // phase: LRURefreshPhase, @@ -104,7 +104,7 @@ pub trait TreeComponent { /// Return a per-match stateful predicate deciding whether a node is a valid /// match boundary for this component. - // Python reference — tree_component.py::TreeComponent.create_match_validator: + // Python reference — base.py::TreeComponent.create_match_validator: // @abstractmethod // def create_match_validator( // self, match_device_only: bool = False @@ -211,7 +211,7 @@ pub trait TreeComponent { /// Redistribute component data between `new_parent` and `child` when a node is /// split; `new_parent` is the newly created prefix node. - // Python reference — tree_component.py::TreeComponent.redistribute_on_node_split: + // Python reference — base.py::TreeComponent.redistribute_on_node_split: // @abstractmethod // def redistribute_on_node_split( // self, new_parent: UnifiedTreeNode, child: UnifiedTreeNode @@ -234,7 +234,7 @@ pub trait TreeComponent { /// Free this component's KV resources on a node being evicted; returns /// (device_freed, host_freed) token counts. - // Python reference — tree_component.py::TreeComponent.evict_component: + // Python reference — base.py::TreeComponent.evict_component: // @abstractmethod // def evict_component( // self, @@ -289,7 +289,7 @@ pub trait TreeComponent { fn evict_device_end(&self, tree_core: &mut UnifiedTreeCore); /// Increment component lock refs, protecting nodes from eviction. - // Python reference — tree_component.py::TreeComponent.acquire_component_lock: + // Python reference — base.py::TreeComponent.acquire_component_lock: // @abstractmethod // def acquire_component_lock( // self, @@ -321,7 +321,7 @@ pub trait TreeComponent { ) -> IncLockRefResult; /// Decrement component lock refs, un-protecting nodes. - // Python reference — tree_component.py::TreeComponent.release_component_lock: + // Python reference — base.py::TreeComponent.release_component_lock: // @abstractmethod // def release_component_lock( // self, @@ -374,7 +374,7 @@ pub trait TreeComponent { prefetch_tokens: usize, last_hash: Option<&str>, ) -> Result>, TreeCoreRuntimeError> { - // Python reference — tree_component.py::TreeComponent.build_hicache_transfers: + // Python reference — base.py::TreeComponent.build_hicache_transfers: // def build_hicache_transfers( // self, // node: UnifiedTreeNode, @@ -412,7 +412,7 @@ pub trait TreeComponent { insert_result: Option<&mut InsertResult>, pool_storage_result: Option<&PoolTransferResult>, ) { - // Python reference — tree_component.py::TreeComponent.commit_hicache_transfer: + // Python reference — base.py::TreeComponent.commit_hicache_transfer: // def commit_hicache_transfer( // self, // node: UnifiedTreeNode, diff --git a/test/registered/unit/mem_cache/test_buffer_mode_sidecar.py b/test/registered/unit/mem_cache/test_buffer_mode_sidecar.py index f3d0ce6c4..3a9fcae76 100644 --- a/test/registered/unit/mem_cache/test_buffer_mode_sidecar.py +++ b/test/registered/unit/mem_cache/test_buffer_mode_sidecar.py @@ -19,7 +19,7 @@ from sglang.srt.mem_cache.hicache_storage import ( SidecarPoolSpec, ) from sglang.srt.mem_cache.radix_cache import RadixKey -from sglang.srt.mem_cache.unified_cache.components.tree_component import ( +from sglang.srt.mem_cache.unified_cache.components.base import ( ComponentType, ) from sglang.srt.mem_cache.unified_cache.unified_tree_core_interface import ( diff --git a/test/registered/unit/mem_cache/test_mamba_donated_alloc_ratio.py b/test/registered/unit/mem_cache/test_mamba_donated_alloc_ratio.py index 0c0a85ef0..b684d867f 100644 --- a/test/registered/unit/mem_cache/test_mamba_donated_alloc_ratio.py +++ b/test/registered/unit/mem_cache/test_mamba_donated_alloc_ratio.py @@ -20,8 +20,8 @@ from sglang.srt.mem_cache.base_prefix_cache import ( EvictParams, IncLockRefResult, ) -from sglang.srt.mem_cache.unified_cache.components.mamba_component import MambaComponent -from sglang.srt.mem_cache.unified_cache.components.tree_component import ComponentType +from sglang.srt.mem_cache.unified_cache.components.base import ComponentType +from sglang.srt.mem_cache.unified_cache.components.mamba import MambaComponent from sglang.srt.mem_cache.unified_cache.unified_tree_core import UnifiedTreeCore from sglang.srt.mem_cache.unified_radix_cache import UnifiedTreeNode from sglang.test.ci.ci_register import register_cpu_ci diff --git a/test/registered/unit/mem_cache/test_mamba_path_state_cap.py b/test/registered/unit/mem_cache/test_mamba_path_state_cap.py index 50d76c69e..c22b20034 100644 --- a/test/registered/unit/mem_cache/test_mamba_path_state_cap.py +++ b/test/registered/unit/mem_cache/test_mamba_path_state_cap.py @@ -15,12 +15,12 @@ import torch from test_unified_radix_cache_unittest import CacheConfig, UnifiedRadixCacheSuite from sglang.srt.mem_cache.unified_cache.cache_action import MambaEvictExcessPathStates -from sglang.srt.mem_cache.unified_cache.components.mamba_component import ( - MambaComponent, -) -from sglang.srt.mem_cache.unified_cache.components.tree_component import ( +from sglang.srt.mem_cache.unified_cache.components.base import ( ComponentType, ) +from sglang.srt.mem_cache.unified_cache.components.mamba import ( + MambaComponent, +) from sglang.srt.mem_cache.unified_cache.unified_tree_core import UnifiedTreeCore from sglang.srt.mem_cache.unified_radix_cache import UnifiedLRUList, UnifiedTreeNode from sglang.srt.server_args import ServerArgs diff --git a/test/registered/unit/mem_cache/test_swa_locked_full_recover_unified.py b/test/registered/unit/mem_cache/test_swa_locked_full_recover_unified.py index 0fafcb426..2be88d351 100644 --- a/test/registered/unit/mem_cache/test_swa_locked_full_recover_unified.py +++ b/test/registered/unit/mem_cache/test_swa_locked_full_recover_unified.py @@ -37,7 +37,7 @@ from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( ) from sglang.srt.mem_cache.unified_cache.cache_action import RecoverSWAWithLockedFull from sglang.srt.mem_cache.unified_cache.component_type import ComponentType -from sglang.srt.mem_cache.unified_cache.components.swa_component import SWAComponent +from sglang.srt.mem_cache.unified_cache.components.swa import SWAComponent from sglang.srt.mem_cache.unified_memory_pool import MHASubPoolSpec, UnifiedKVPool from sglang.test.ci.ci_register import register_cpu_ci diff --git a/test/registered/unit/mem_cache/test_tree_core_registry.py b/test/registered/unit/mem_cache/test_tree_core_registry.py index 983868e36..11c8b9d59 100644 --- a/test/registered/unit/mem_cache/test_tree_core_registry.py +++ b/test/registered/unit/mem_cache/test_tree_core_registry.py @@ -9,7 +9,7 @@ from sglang.srt.environ import envs from sglang.srt.mem_cache.cache_init_params import CacheInitParams from sglang.srt.mem_cache.memory_pool import ReqToTokenPool from sglang.srt.mem_cache.unified_cache.component_type import ComponentType -from sglang.srt.mem_cache.unified_cache.components.tree_component import ( +from sglang.srt.mem_cache.unified_cache.components.base import ( EvictLayer, TreeComponent, ) diff --git a/test/registered/unit/mem_cache/test_unified_cache_linker.py b/test/registered/unit/mem_cache/test_unified_cache_linker.py index 89c671a6c..fe6ba8ce9 100644 --- a/test/registered/unit/mem_cache/test_unified_cache_linker.py +++ b/test/registered/unit/mem_cache/test_unified_cache_linker.py @@ -36,12 +36,12 @@ from sglang.srt.mem_cache.unified_cache.cache_action import ( ReplaceWriteThroughOnNodeSplit, ) from sglang.srt.mem_cache.unified_cache.component_type import ComponentType -from sglang.srt.mem_cache.unified_cache.components.full_component import FullComponent -from sglang.srt.mem_cache.unified_cache.components.swa_component import SWAComponent -from sglang.srt.mem_cache.unified_cache.components.tree_component import ( +from sglang.srt.mem_cache.unified_cache.components.base import ( ExternalLinkerLoadPhase, LinkerTransferPhase, ) +from sglang.srt.mem_cache.unified_cache.components.full import FullComponent +from sglang.srt.mem_cache.unified_cache.components.swa import SWAComponent from sglang.srt.mem_cache.unified_cache.unified_cache_linker import ( ExternalCacheHitMarker, UnifiedCacheLinker, diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py b/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py index 740e0b0bb..95460dc3c 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py @@ -34,7 +34,7 @@ from sglang.srt.mem_cache.mamba_radix_cache import MambaRadixCache from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool, HybridReqToTokenPool from sglang.srt.mem_cache.radix_cache import RadixKey from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache -from sglang.srt.mem_cache.unified_cache.components.tree_component import ComponentType +from sglang.srt.mem_cache.unified_cache.components.base import ComponentType from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler from sglang.srt.utils import get_device 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 745074fa5..d792be685 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 @@ -65,7 +65,7 @@ from sglang.srt.mem_cache.unified_cache.cache_action import ( ReplaceWriteThroughOnNodeSplit, SWARebuild, ) -from sglang.srt.mem_cache.unified_cache.components.tree_component import ( +from sglang.srt.mem_cache.unified_cache.components.base import ( CacheTransferPhase, ComponentType, EvictLayer,