From 60d6914f17ec691edb3258601a3c65b1e4f8de61 Mon Sep 17 00:00:00 2001 From: Zhangheng Date: Tue, 28 Jul 2026 16:28:21 +0800 Subject: [PATCH] [UnifiedTree]: move /mem_cache/unifed_cache_component dir to /mem_cache/unified_cache (#32484) --- .../mlx/kv_cache/auxiliary_state.py | 4 ++-- .../sglang/srt/mem_cache/base_prefix_cache.py | 2 +- .../sglang/srt/mem_cache/cache_init_params.py | 4 ++-- .../hybrid_cache/hybrid_pool_assembler.py | 2 +- python/sglang/srt/mem_cache/registry.py | 2 +- .../srt/mem_cache/unified_cache/__init__.py | 2 +- .../components}/README.md | 0 .../components}/__init__.py | 8 +++---- .../components}/full_component.py | 2 +- .../components}/mamba_component.py | 2 +- .../components}/swa_component.py | 2 +- .../components}/tree_component.py | 0 .../unified_cache/tree_core_registry.py | 2 +- .../unified_cache/unified_tree_core.py | 22 +++++++++---------- .../unified_tree_core_interface.py | 8 +++---- .../srt/mem_cache/unified_radix_cache.py | 16 +++++++------- .../kv_canary/test_self_unit_radix_walker.py | 4 ++-- .../mem_cache/test_mamba_path_state_cap.py | 6 ++--- .../unit/mem_cache/test_registry.py | 12 +++++----- .../unit/mem_cache/test_tree_core_registry.py | 8 +++---- .../test_unified_radix_cache_bench.py | 2 +- .../test_unified_radix_cache_unittest.py | 12 +++++----- .../test_unified_radix_hicache_dispatch.py | 2 +- 23 files changed, 62 insertions(+), 62 deletions(-) rename python/sglang/srt/mem_cache/{unified_cache_components => unified_cache/components}/README.md (100%) rename python/sglang/srt/mem_cache/{unified_cache_components => unified_cache/components}/__init__.py (74%) rename python/sglang/srt/mem_cache/{unified_cache_components => unified_cache/components}/full_component.py (99%) rename python/sglang/srt/mem_cache/{unified_cache_components => unified_cache/components}/mamba_component.py (99%) rename python/sglang/srt/mem_cache/{unified_cache_components => unified_cache/components}/swa_component.py (99%) rename python/sglang/srt/mem_cache/{unified_cache_components => unified_cache/components}/tree_component.py (100%) 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 27284bf08..928d10156 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.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() diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index 7a9a2b8c8..f54a986b1 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -31,7 +31,7 @@ if TYPE_CHECKING: 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, ) diff --git a/python/sglang/srt/mem_cache/cache_init_params.py b/python/sglang/srt/mem_cache/cache_init_params.py index ee749814b..c402cec4a 100644 --- a/python/sglang/srt/mem_cache/cache_init_params.py +++ b/python/sglang/srt/mem_cache/cache_init_params.py @@ -8,8 +8,8 @@ import torch 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, ) diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py index 6986cd219..289f91505 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py @@ -27,7 +27,7 @@ from sglang.srt.mem_cache.pool_host.mha import ( 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 diff --git a/python/sglang/srt/mem_cache/registry.py b/python/sglang/srt/mem_cache/registry.py index 744134e01..e84a01f12 100644 --- a/python/sglang/srt/mem_cache/registry.py +++ b/python/sglang/srt/mem_cache/registry.py @@ -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] diff --git a/python/sglang/srt/mem_cache/unified_cache/__init__.py b/python/sglang/srt/mem_cache/unified_cache/__init__.py index 15c752c54..fd18355c5 100644 --- a/python/sglang/srt/mem_cache/unified_cache/__init__.py +++ b/python/sglang/srt/mem_cache/unified_cache/__init__.py @@ -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 diff --git a/python/sglang/srt/mem_cache/unified_cache_components/README.md b/python/sglang/srt/mem_cache/unified_cache/components/README.md similarity index 100% rename from python/sglang/srt/mem_cache/unified_cache_components/README.md rename to python/sglang/srt/mem_cache/unified_cache/components/README.md diff --git a/python/sglang/srt/mem_cache/unified_cache_components/__init__.py b/python/sglang/srt/mem_cache/unified_cache/components/__init__.py similarity index 74% rename from python/sglang/srt/mem_cache/unified_cache_components/__init__.py rename to python/sglang/srt/mem_cache/unified_cache/components/__init__.py index c8fcbc4ad..89b0665e2 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,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, diff --git a/python/sglang/srt/mem_cache/unified_cache_components/full_component.py b/python/sglang/srt/mem_cache/unified_cache/components/full_component.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_component.py index 93bba1dc5..7558254b7 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/full_component.py +++ b/python/sglang/srt/mem_cache/unified_cache/components/full_component.py @@ -18,7 +18,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.tree_component 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_component.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_component.py index 9d51c6988..9dcb921ba 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 @@ -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.tree_component 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_component.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_component.py index 12621dce7..d586714a9 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 @@ -27,7 +27,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.tree_component import ( BASE_COMPONENT_TYPE, CacheTransferPhase, ComponentType, diff --git a/python/sglang/srt/mem_cache/unified_cache_components/tree_component.py b/python/sglang/srt/mem_cache/unified_cache/components/tree_component.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/tree_component.py diff --git a/python/sglang/srt/mem_cache/unified_cache/tree_core_registry.py b/python/sglang/srt/mem_cache/unified_cache/tree_core_registry.py index f7b9871fb..f8a30bc3f 100644 --- a/python/sglang/srt/mem_cache/unified_cache/tree_core_registry.py +++ b/python/sglang/srt/mem_cache/unified_cache/tree_core_registry.py @@ -13,10 +13,10 @@ from typing import TYPE_CHECKING, Callable, Optional 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]"], 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 6fe91f08e..cf783205f 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 @@ -49,6 +49,17 @@ from sglang.srt.mem_cache.unified_cache.cache_action import ( 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, @@ -61,17 +72,6 @@ from sglang.srt.mem_cache.unified_cache.unified_tree_core_interface import ( 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, 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 f8da32f70..a06b0517a 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 @@ -102,14 +102,14 @@ if TYPE_CHECKING: 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): diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 09d58daad..d699ebd1f 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -41,17 +41,10 @@ from sglang.srt.mem_cache.unified_cache.cache_action import ( 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, @@ -61,6 +54,13 @@ from sglang.srt.mem_cache.unified_cache_components import ( 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, diff --git a/test/registered/kv_canary/test_self_unit_radix_walker.py b/test/registered/kv_canary/test_self_unit_radix_walker.py index 6651bfd49..4e72e41d9 100644 --- a/test/registered/kv_canary/test_self_unit_radix_walker.py +++ b/test/registered/kv_canary/test_self_unit_radix_walker.py @@ -10,11 +10,11 @@ from sglang.srt.kv_canary.radix_cache_walker import walk_radix_cache_for_canary 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 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 283960f18..ebfaa22c5 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 @@ -14,13 +14,13 @@ 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.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 diff --git a/test/registered/unit/mem_cache/test_registry.py b/test/registered/unit/mem_cache/test_registry.py index 1f25923da..adaa68840 100644 --- a/test/registered/unit/mem_cache/test_registry.py +++ b/test/registered/unit/mem_cache/test_registry.py @@ -227,7 +227,7 @@ class TestDefaultRadixCacheFactory(CustomTestCase): 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, }, ), @@ -261,7 +261,7 @@ class TestDefaultRadixCacheFactory(CustomTestCase): 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, }, ): @@ -281,7 +281,7 @@ class TestDefaultRadixCacheFactory(CustomTestCase): 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, }, ): @@ -302,7 +302,7 @@ class TestDefaultRadixCacheFactory(CustomTestCase): 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, }, ): @@ -322,7 +322,7 @@ class TestDefaultRadixCacheFactory(CustomTestCase): 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, }, ): @@ -348,7 +348,7 @@ class TestDefaultRadixCacheFactory(CustomTestCase): 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, }, ): 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 28a37a51e..0769dd8ce 100644 --- a/test/registered/unit/mem_cache/test_tree_core_registry.py +++ b/test/registered/unit/mem_cache/test_tree_core_registry.py @@ -7,6 +7,10 @@ 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 ( + EvictLayer, + TreeComponent, +) from sglang.srt.mem_cache.unified_cache.tree_core_registry import ( _TREE_CORE_REGISTRY, create_tree_core, @@ -14,10 +18,6 @@ from sglang.srt.mem_cache.unified_cache.tree_core_registry import ( 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 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 bd721e671..d8aadd662 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.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 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 c06e34295..9b971612e 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 @@ -59,6 +59,12 @@ from sglang.srt.mem_cache.unified_cache.cache_action import ( 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, @@ -67,12 +73,6 @@ from sglang.srt.mem_cache.unified_cache.unified_tree_core_interface import ( 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, diff --git a/test/registered/unit/mem_cache/test_unified_radix_hicache_dispatch.py b/test/registered/unit/mem_cache/test_unified_radix_hicache_dispatch.py index d39d4e564..f9f492bfa 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_hicache_dispatch.py +++ b/test/registered/unit/mem_cache/test_unified_radix_hicache_dispatch.py @@ -17,7 +17,7 @@ from sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler import ( _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")