diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index ac69c3e73..3d4905e50 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -22,13 +22,12 @@ from typing import TYPE_CHECKING, List, NamedTuple, Optional import torch -from sglang.srt.mem_cache.hicache_storage import ( - HiCacheStorageConfig, - HiCacheStorageExtraInfo, -) - if TYPE_CHECKING: from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator + from sglang.srt.mem_cache.hicache_storage import ( + HiCacheStorageConfig, + HiCacheStorageExtraInfo, + ) from sglang.srt.mem_cache.memory_pool_host import HostKVCache from sglang.srt.distributed import ( @@ -926,9 +925,7 @@ class HiCacheController: for page in pages: self.host_mem_release_queue.put(page) - def _page_get_zero_copy( - self, operation, hash_values, host_indices, extra_info=None - ): + def _page_get_zero_copy(self, operation, hash_values, host_indices, extra_info=None): results = self.storage_backend.batch_get_v1( hash_values, host_indices, extra_info ) @@ -975,6 +972,8 @@ class HiCacheController: ] prev_completed_tokens = operation.completed_tokens # Get one batch token, and update the completed_tokens if succeed + from sglang.srt.mem_cache.hicache_storage import HiCacheStorageExtraInfo + extra_info = HiCacheStorageExtraInfo(prefix_keys=prefix_keys) self.page_get_func(operation, batch_hashes, batch_host_indices, extra_info) # Check termination @@ -1036,6 +1035,8 @@ class HiCacheController: batch_tokens[i : i + self.page_size], last_hash ) batch_hashes.append(last_hash) + from sglang.srt.mem_cache.hicache_storage import HiCacheStorageExtraInfo + extra_info = HiCacheStorageExtraInfo(prefix_keys=prefix_keys) hit_page_num = self.storage_backend.batch_exists(batch_hashes, extra_info) hash_value.extend(batch_hashes[:hit_page_num]) @@ -1137,6 +1138,8 @@ class HiCacheController: ] # Set one batch token, and record if success. # todo: allow partial success + from sglang.srt.mem_cache.hicache_storage import HiCacheStorageExtraInfo + extra_info = HiCacheStorageExtraInfo(prefix_keys=prefix_keys) success = self.page_set_func(batch_hashes, batch_host_indices, extra_info) if not success: diff --git a/python/sglang/srt/mem_cache/hicache_storage.py b/python/sglang/srt/mem_cache/hicache_storage.py index e76cfebd4..43644bd7c 100644 --- a/python/sglang/srt/mem_cache/hicache_storage.py +++ b/python/sglang/srt/mem_cache/hicache_storage.py @@ -1,14 +1,18 @@ +from __future__ import annotations + import hashlib import logging import os from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import Any, List, Optional +from typing import TYPE_CHECKING, Any, List, Optional import torch from sglang.srt.environ import envs -from sglang.srt.mem_cache.memory_pool_host import HostKVCache + +if TYPE_CHECKING: + from sglang.srt.mem_cache.memory_pool_host import HostKVCache logger = logging.getLogger(__name__) diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index a48e49094..aa4d1241f 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -34,11 +34,6 @@ from sglang.srt.mem_cache.memory_pool import ( MLATokenToKVPool, NSATokenToKVPool, ) -from sglang.srt.mem_cache.memory_pool_host import ( - MHATokenToKVPoolHost, - MLATokenToKVPoolHost, - NSATokenToKVPoolHost, -) from sglang.srt.mem_cache.radix_cache import ( RadixCache, RadixKey, @@ -136,6 +131,12 @@ class HiRadixCache(RadixCache): ) self.kv_cache = params.token_to_kv_pool_allocator.get_kvcache() + from sglang.srt.mem_cache.memory_pool_host import ( + MHATokenToKVPoolHost, + MLATokenToKVPoolHost, + NSATokenToKVPoolHost, + ) + if isinstance(self.kv_cache, MHATokenToKVPool): self.token_to_kv_pool_host = MHATokenToKVPoolHost( self.kv_cache, diff --git a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py index 49fa4409e..0f61eaa1d 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py +++ b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py @@ -19,6 +19,24 @@ from sglang.test.test_utils import CustomTestCase register_cpu_ci(est_time=2, suite="stage-a-test-cpu") +class TestCpHiCacheImports(CustomTestCase): + def test_cp_hicache_public_imports_without_sgl_kernel(self): + import subprocess + + subprocess.run( + [ + sys.executable, + "-c", + "from sglang.srt.mem_cache.hiradix_cache import HiRadixCache, CpHiCacheNodeMetadata; " + "from sglang.srt.managers.cache_controller import HiCacheController; " + "print('OK')", + ], + check=True, + capture_output=True, + text=True, + ) + + class TestCpHiCacheNodeMetadata(CustomTestCase): def test_split_zero_len_moves_all_positions_to_child(self): metadata = CpHiCacheNodeMetadata(