[mem_cache][10/N] refactor: drop the redundant _component suffix in unified_cache/components (#35644)
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 5
parent
14a131ad5b
commit
7f1f8c706a
@@ -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()
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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",
|
||||
|
||||
+1
-1
@@ -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,
|
||||
+1
-1
@@ -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,
|
||||
+1
-1
@@ -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,
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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,
|
||||
@@ -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,
|
||||
@@ -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,
|
||||
@@ -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,
|
||||
@@ -412,7 +412,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,
|
||||
|
||||
@@ -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 (
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user