Skip to content
Original file line number Diff line number Diff line change
Expand Up @@ -17,10 +17,10 @@

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()
Expand Down
2 changes: 1 addition & 1 deletion python/sglang/srt/mem_cache/cache_init_params.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
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,
)

Expand Down
10 changes: 5 additions & 5 deletions python/sglang/srt/mem_cache/unified_cache/components/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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` |

Expand Down Expand Up @@ -259,7 +259,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

Expand Down
Original file line number Diff line number Diff line change
@@ -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,
Expand All @@ -17,6 +14,9 @@
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",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
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,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@
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,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@
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,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@

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
Expand Down
16 changes: 8 additions & 8 deletions rust/sglang-radix-tree/src/components/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,7 @@ pub trait TreeComponent<K: ChildKeyType> {
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,
Expand Down Expand Up @@ -104,7 +104,7 @@ pub trait TreeComponent<K: ChildKeyType> {

/// 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
Expand Down Expand Up @@ -211,7 +211,7 @@ pub trait TreeComponent<K: ChildKeyType> {

/// 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
Expand All @@ -234,7 +234,7 @@ pub trait TreeComponent<K: ChildKeyType> {

/// 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,
Expand Down Expand Up @@ -289,7 +289,7 @@ pub trait TreeComponent<K: ChildKeyType> {
fn evict_device_end(&self, tree_core: &mut UnifiedTreeCore<K>);

/// 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,
Expand Down Expand Up @@ -321,7 +321,7 @@ pub trait TreeComponent<K: ChildKeyType> {
) -> 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,
Expand Down Expand Up @@ -374,7 +374,7 @@ pub trait TreeComponent<K: ChildKeyType> {
prefetch_tokens: usize,
last_hash: Option<&str>,
) -> Result<Option<Vec<PoolTransfer>>, 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,
Expand Down Expand Up @@ -403,7 +403,7 @@ pub trait TreeComponent<K: ChildKeyType> {
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,
Expand Down
2 changes: 1 addition & 1 deletion test/registered/unit/mem_cache/test_buffer_mode_sidecar.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
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 (
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,8 +19,8 @@
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
Expand Down
8 changes: 4 additions & 4 deletions test/registered/unit/mem_cache/test_mamba_path_state_cap.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,12 +15,12 @@
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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@
)
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

Expand Down
2 changes: 1 addition & 1 deletion test/registered/unit/mem_cache/test_tree_core_registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
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,
)
Expand Down
6 changes: 3 additions & 3 deletions test/registered/unit/mem_cache/test_unified_cache_linker.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,12 +9,12 @@
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 (
UnifiedCacheLinker,
UnifiedCacheLinkerWrapper,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@
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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -64,7 +64,7 @@
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,
Expand Down
Loading