Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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.mamba_component import (
MambaComponent,
)
from sglang.srt.mem_cache.unified_cache_components.tree_component import TreeComponent
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/base_prefix_cache.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@
CacheAction,
ComponentAction,
)
from sglang.srt.mem_cache.unified_cache_components.tree_component import (
from sglang.srt.mem_cache.unified_cache.components.tree_component import (
ComponentType,
)

Expand Down
4 changes: 2 additions & 2 deletions python/sglang/srt/mem_cache/cache_init_params.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,8 +8,8 @@
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 import ComponentType
from sglang.srt.mem_cache.unified_cache.components.tree_component import (
TreeComponent,
)

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@
get_mha_host_pool_cls,
)
from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost
from sglang.srt.mem_cache.unified_cache_components import ComponentType
from sglang.srt.mem_cache.unified_cache.components import ComponentType

if TYPE_CHECKING:
import torch
Expand Down
2 changes: 1 addition & 1 deletion python/sglang/srt/mem_cache/registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -166,7 +166,7 @@ def _create_unified_radix_cache(
params: CacheInitParams,
) -> BasePrefixCache:
"""Initialize a UnifiedRadixCache with proper components and optional HiCache."""
from sglang.srt.mem_cache.unified_cache_components import ComponentType
from sglang.srt.mem_cache.unified_cache.components import ComponentType
from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache

tree_components = [ComponentType.FULL]
Expand Down
2 changes: 1 addition & 1 deletion python/sglang/srt/mem_cache/unified_cache/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
events -- the only object that touches that state; no scheduling policy or IO,
and never touches the cache. References the cache-owned component drivers to
drive their tree-level hooks.
- **Components** (``unified_cache_components.TreeComponent`` subclasses): the
- **Components** (``unified_cache.components.TreeComponent`` subclasses): the
per-component (FULL/SWA/MAMBA) drivers. Built and owned by the Controller;
hold the cache for cache-level logic and the TreeCore for tree state.
- **Controller** (``unified_radix_cache.UnifiedRadixCache``): the scheduling brain
Expand Down
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
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.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 (
_NUM_COMPONENT_TYPES,
BASE_COMPONENT_TYPE,
CacheTransferPhase,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,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.tree_component 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.tree_component import (
CacheTransferPhase,
ComponentType,
EvictLayer,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@
RecoverSWAWithLockedFull,
SWARebuild,
)
from sglang.srt.mem_cache.unified_cache_components.tree_component import (
from sglang.srt.mem_cache.unified_cache.components.tree_component import (
BASE_COMPONENT_TYPE,
CacheTransferPhase,
ComponentType,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -13,10 +13,10 @@
if TYPE_CHECKING:
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
from sglang.srt.mem_cache.unified_cache.component_type import ComponentType
from sglang.srt.mem_cache.unified_cache.components import TreeComponent
from sglang.srt.mem_cache.unified_cache.unified_tree_core_interface import (
UnifiedTreeCoreInterface,
)
from sglang.srt.mem_cache.unified_cache_components import TreeComponent

TreeCoreFactory = Callable[
["CacheInitParams", "dict[ComponentType, TreeComponent]"],
Expand Down
22 changes: 11 additions & 11 deletions python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,17 @@
FreeDeviceKV,
ReplaceWriteThroughOnNodeSplit,
)
from sglang.srt.mem_cache.unified_cache.components import (
_NUM_COMPONENT_TYPES,
BASE_COMPONENT_TYPE,
CacheTransferPhase,
ComponentData,
ComponentType,
EvictLayer,
LRURefreshPhase,
TreeComponent,
get_and_increase_time_counter,
)
from sglang.srt.mem_cache.unified_cache.unified_tree_core_interface import (
DecSwaLockOnlyResult,
DemoteResult,
Expand All @@ -61,17 +72,6 @@
RadixCacheWalkResult,
UnifiedTreeCoreInterface,
)
from sglang.srt.mem_cache.unified_cache_components import (
_NUM_COMPONENT_TYPES,
BASE_COMPONENT_TYPE,
CacheTransferPhase,
ComponentData,
ComponentType,
EvictLayer,
LRURefreshPhase,
TreeComponent,
get_and_increase_time_counter,
)
from sglang.srt.mem_cache.utils import (
compute_node_hash_values,
get_eviction_strategy,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -102,14 +102,14 @@ class InsertStepResult(msgspec.Struct, frozen=True):
CacheAction,
ComponentAction,
)
from sglang.srt.mem_cache.unified_cache.components import (
CacheTransferPhase,
ComponentType,
)
from sglang.srt.mem_cache.unified_cache.unified_tree_core import (
StorageBackupSpec,
UnifiedTreeNode,
)
from sglang.srt.mem_cache.unified_cache_components import (
CacheTransferPhase,
ComponentType,
)


class UnifiedTreeCoreInterface(KVCacheEventMixin, ABC):
Expand Down
16 changes: 8 additions & 8 deletions python/sglang/srt/mem_cache/unified_radix_cache.py
Original file line number Diff line number Diff line change
Expand Up @@ -41,17 +41,10 @@
FreeDeviceKV,
ReplaceWriteThroughOnNodeSplit,
)
from sglang.srt.mem_cache.unified_cache.tree_core_registry import create_tree_core
from sglang.srt.mem_cache.unified_cache.unified_tree_core import ( # noqa: F401
NodeId,
UnifiedLRUList,
UnifiedTreeCore,
UnifiedTreeNode,
)

# UnifiedTreeNode / UnifiedLRUList live on the tree core; re-exported here
# because other modules and tests import them from this module.
from sglang.srt.mem_cache.unified_cache_components import (
from sglang.srt.mem_cache.unified_cache.components import (
BASE_COMPONENT_TYPE,
CacheTransferPhase,
ComponentType,
Expand All @@ -61,6 +54,13 @@
SWAComponent,
TreeComponent,
)
from sglang.srt.mem_cache.unified_cache.tree_core_registry import create_tree_core
from sglang.srt.mem_cache.unified_cache.unified_tree_core import ( # noqa: F401
NodeId,
UnifiedLRUList,
UnifiedTreeCore,
UnifiedTreeNode,
)
from sglang.srt.observability.metrics_collector import (
STAT_LOGGER_ROLE_STORAGE,
StorageMetricsCollector,
Expand Down
4 changes: 2 additions & 2 deletions test/registered/kv_canary/test_self_unit_radix_walker.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,11 +10,11 @@
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
from sglang.srt.mem_cache.radix_cache import RadixKey
from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache, TreeNode
from sglang.srt.mem_cache.unified_cache.unified_tree_core import UnifiedTreeCore
from sglang.srt.mem_cache.unified_cache_components import (
from sglang.srt.mem_cache.unified_cache.components import (
BASE_COMPONENT_TYPE,
ComponentType,
)
from sglang.srt.mem_cache.unified_cache.unified_tree_core import UnifiedTreeCore
from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache, UnifiedTreeNode
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.kv_canary.fixtures import DEFAULT_DEVICE, make_radix_cache
Expand Down
6 changes: 3 additions & 3 deletions test/registered/unit/mem_cache/test_mamba_path_state_cap.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,13 +14,13 @@
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.unified_tree_core import UnifiedTreeCore
from sglang.srt.mem_cache.unified_cache_components.mamba_component import (
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.tree_component import (
ComponentType,
)
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
from sglang.test.test_utils import CustomTestCase
Expand Down
12 changes: 6 additions & 6 deletions test/registered/unit/mem_cache/test_registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -227,7 +227,7 @@ def test_unified_radix_cache_when_env_flag_set(self):
patch.dict(
"sys.modules",
{
"sglang.srt.mem_cache.unified_cache_components": fake_components,
"sglang.srt.mem_cache.unified_cache.components": fake_components,
"sglang.srt.mem_cache.unified_radix_cache": fake_radix,
},
),
Expand Down Expand Up @@ -261,7 +261,7 @@ def test_unified_radix_cache_when_hierarchical_and_hybrid_ssm(self):
with patch.dict(
"sys.modules",
{
"sglang.srt.mem_cache.unified_cache_components": fake_components,
"sglang.srt.mem_cache.unified_cache.components": fake_components,
"sglang.srt.mem_cache.unified_radix_cache": fake_radix,
},
):
Expand All @@ -281,7 +281,7 @@ def test_unified_radix_cache_when_hierarchical_and_hybrid_swa(self):
with patch.dict(
"sys.modules",
{
"sglang.srt.mem_cache.unified_cache_components": fake_components,
"sglang.srt.mem_cache.unified_cache.components": fake_components,
"sglang.srt.mem_cache.unified_radix_cache": fake_radix,
},
):
Expand All @@ -302,7 +302,7 @@ def test_unified_radix_cache_when_hierarchical_and_dsa(self):
with patch.dict(
"sys.modules",
{
"sglang.srt.mem_cache.unified_cache_components": fake_components,
"sglang.srt.mem_cache.unified_cache.components": fake_components,
"sglang.srt.mem_cache.unified_radix_cache": fake_radix,
},
):
Expand All @@ -322,7 +322,7 @@ def test_swa_radix_cache_when_hybrid_swa(self):
with patch.dict(
"sys.modules",
{
"sglang.srt.mem_cache.unified_cache_components": fake_components,
"sglang.srt.mem_cache.unified_cache.components": fake_components,
"sglang.srt.mem_cache.unified_radix_cache": fake_radix,
},
):
Expand All @@ -348,7 +348,7 @@ def test_mamba_radix_cache_when_hybrid_ssm(self):
with patch.dict(
"sys.modules",
{
"sglang.srt.mem_cache.unified_cache_components": fake_components,
"sglang.srt.mem_cache.unified_cache.components": fake_components,
"sglang.srt.mem_cache.unified_radix_cache": fake_radix,
},
):
Expand Down
8 changes: 4 additions & 4 deletions test/registered/unit/mem_cache/test_tree_core_registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,17 +7,17 @@
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 (
EvictLayer,
TreeComponent,
)
from sglang.srt.mem_cache.unified_cache.tree_core_registry import (
_TREE_CORE_REGISTRY,
create_tree_core,
register_tree_core_backend,
registered_tree_core_backends,
)
from sglang.srt.mem_cache.unified_cache.unified_tree_core import UnifiedTreeCore
from sglang.srt.mem_cache.unified_cache_components.tree_component import (
EvictLayer,
TreeComponent,
)
from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,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.tree_component 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 @@ -59,6 +59,12 @@
ReplaceWriteThroughOnNodeSplit,
SWARebuild,
)
from sglang.srt.mem_cache.unified_cache.components.tree_component import (
CacheTransferPhase,
ComponentType,
EvictLayer,
TreeComponent,
)
from sglang.srt.mem_cache.unified_cache.unified_tree_core_interface import (
DecSwaLockOnlyResult,
DemoteResult,
Expand All @@ -67,12 +73,6 @@
EvictDeviceLeafResult,
EvictDeviceNextNodeResult,
)
from sglang.srt.mem_cache.unified_cache_components.tree_component import (
CacheTransferPhase,
ComponentType,
EvictLayer,
TreeComponent,
)
from sglang.srt.mem_cache.unified_radix_cache import (
COMPONENT_REGISTRY,
UnifiedLRUList,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@
_SwaStrategy,
register_stack_strategy,
)
from sglang.srt.mem_cache.unified_cache_components import ComponentType
from sglang.srt.mem_cache.unified_cache.components import ComponentType
from sglang.test.ci.ci_register import register_cpu_ci

register_cpu_ci(est_time=2, suite="base-a-test-cpu")
Expand Down
Loading