[mem_cache][11/N] refactor: extract KVCache and BaseSWAKVPool into pool/base.py

python/sglang/srt/disaggregation/decode.py

@@ -90,9 +90,11 @@
 from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
 from sglang.srt.mem_cache.memory_pool import (
     HybridReqToTokenPool,
-    KVCache,
     ReqToTokenPool,
 )
+from sglang.srt.mem_cache.pool.base import (
+    KVCache,
+)
 from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
 from sglang.srt.observability.req_time_stats import (
     set_schedule_time_batch,

python/sglang/srt/disaggregation/prefill.py

@@ -90,7 +90,7 @@
     from torch.distributed import ProcessGroup
 
     from sglang.srt.managers.scheduler import GenerationBatchResult, Scheduler
-    from sglang.srt.mem_cache.memory_pool import KVCache
+    from sglang.srt.mem_cache.pool.base import KVCache
 
 logger = logging.getLogger(__name__)
 

python/sglang/srt/disaggregation/utils.py

@@ -1294,14 +1294,14 @@ def setup_state_kv_args(
     """
     from sglang.srt.disaggregation.base.conn import StateType
     from sglang.srt.hardware_backend.npu.memory_pool_npu import NPUMLATokenToKVPool
-    from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool
     from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
     from sglang.srt.mem_cache.memory_pool import (
         DSATokenToKVPool,
         HybridLinearKVPool,
         MHATokenToKVPoolMXFP8,
         MiniMaxSparseKVPool,
     )
+    from sglang.srt.mem_cache.pool.base import BaseSWAKVPool
     from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
 
     kv_args.state_types = []

python/sglang/srt/hardware_backend/mlx/model_runner_stub.py

@@ -14,7 +14,8 @@
     MlxAuxiliaryStateReqToTokenPool,
 )
 from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator
-from sglang.srt.mem_cache.memory_pool import KVCache, ReqToTokenPool
+from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
+from sglang.srt.mem_cache.pool.base import KVCache
 from sglang.srt.model_executor.model_runner import ModelRunner
 from sglang.srt.model_executor.model_runner_components.layer_setup import (
     ModelLayerInfo,

python/sglang/srt/hardware_backend/musa/attention/flashattention_backend.py

@@ -22,7 +22,7 @@
 from sglang.srt.layers.utils.cp_utils import (
     cp_allgather_and_save_kv_cache,
 )
-from sglang.srt.mem_cache.memory_pool import KVWriteLoc
+from sglang.srt.mem_cache.pool.base import KVWriteLoc
 from sglang.srt.runtime_context import get_schedule
 
 if TYPE_CHECKING:

python/sglang/srt/hardware_backend/npu/allocator_npu.py

@@ -9,7 +9,7 @@
 from sglang.srt.utils import get_num_new_pages, next_power_of_2
 
 if TYPE_CHECKING:
-    from sglang.srt.mem_cache.memory_pool import KVCache
+    from sglang.srt.mem_cache.pool.base import KVCache
 
 
 class NPUPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator):

python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py

@@ -27,7 +27,7 @@
 from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp
 from sglang.srt.layers.radix_attention import AttentionType
 from sglang.srt.layers.utils.cp_utils import cp_all_gather_rerange_kv_cache
-from sglang.srt.mem_cache.memory_pool import KVWriteLoc
+from sglang.srt.mem_cache.pool.base import KVWriteLoc
 from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.srt.runtime_context import (

python/sglang/srt/hardware_backend/npu/memory_pool_npu.py

@@ -10,6 +10,8 @@
     MiniMaxSparseKVPool,
     MLATokenToKVPool,
     get_tensor_size_bytes,
+)
+from sglang.srt.mem_cache.pool.base import (
     unwrap_write_loc,
 )
 from sglang.srt.utils import get_bool_env_var

python/sglang/srt/kv_canary/pool_patcher/api.py

@@ -13,10 +13,12 @@
 from sglang.srt.kv_canary.pool_patcher.buffer_alloc import resolve_real_kv_read_bytes
 from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
 from sglang.srt.mem_cache.memory_pool import (
-    KVCache,
     MHATokenToKVPool,
     MHATokenToKVPoolFP4,
 )
+from sglang.srt.mem_cache.pool.base import (
+    KVCache,
+)
 from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
 
 logger = logging.getLogger(__name__)

python/sglang/srt/layers/attention/aiter_backend.py

@@ -92,7 +92,7 @@
     forward_decode_vectorized_5d,
     forward_extend_vectorized_5d,
 )
-from sglang.srt.mem_cache.memory_pool import KVWriteLoc
+from sglang.srt.mem_cache.pool.base import KVWriteLoc
 from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
 from sglang.srt.utils import get_bool_env_var
 

python/sglang/srt/layers/attention/flashattention_backend.py

@@ -26,7 +26,7 @@
 from sglang.srt.layers.cp.base import CPAttentionBackendKind, get_cp_strategy
 from sglang.srt.layers.cp.utils import is_cp_active
 from sglang.srt.layers.radix_attention import AttentionType
-from sglang.srt.mem_cache.memory_pool import KVWriteLoc
+from sglang.srt.mem_cache.pool.base import KVWriteLoc
 from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.srt.runtime_context import (

python/sglang/srt/layers/attention/flashinfer_backend.py

@@ -34,8 +34,7 @@
     KVCacheAttentionAccessKind,
 )
 from sglang.srt.layers.radix_attention import AttentionType
-from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool
-from sglang.srt.mem_cache.memory_pool import KVWriteLoc
+from sglang.srt.mem_cache.pool.base import BaseSWAKVPool, KVWriteLoc
 from sglang.srt.model_executor.cuda_graph_config import (
     Backend,
     Phase,

python/sglang/srt/layers/attention/hpc_ops_backend.py

@@ -42,7 +42,7 @@
 )
 from sglang.srt.compilation.compilation_config import register_split_op
 from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
-from sglang.srt.mem_cache.memory_pool import KVWriteLoc
+from sglang.srt.mem_cache.pool.base import KVWriteLoc
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_executor.forward_context import get_attn_backend
 from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import (

python/sglang/srt/layers/attention/intel_amx_backend.py

@@ -6,7 +6,7 @@
 
 from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
 from sglang.srt.layers.radix_attention import AttentionType
-from sglang.srt.mem_cache.memory_pool import KVWriteLoc
+from sglang.srt.mem_cache.pool.base import KVWriteLoc
 from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.runtime_context import get_parallel, get_spec

python/sglang/srt/layers/attention/torch_native_backend.py

@@ -7,7 +7,7 @@
 
 from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
 from sglang.srt.layers.radix_attention import AttentionType
-from sglang.srt.mem_cache.memory_pool import KVWriteLoc
+from sglang.srt.mem_cache.pool.base import KVWriteLoc
 from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 

python/sglang/srt/layers/attention/triton_backend.py

@@ -29,8 +29,7 @@
     get_dcp_lens,
 )
 from sglang.srt.layers.radix_attention import AttentionType
-from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool
-from sglang.srt.mem_cache.memory_pool import KVWriteLoc
+from sglang.srt.mem_cache.pool.base import BaseSWAKVPool, KVWriteLoc
 from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
 from sglang.srt.model_executor.cuda_graph_config import (
     Backend,

python/sglang/srt/layers/attention/trtllm_mha_backend.py

@@ -42,7 +42,7 @@
     KVCacheAttentionAccessKind,
 )
 from sglang.srt.layers.radix_attention import AttentionType
-from sglang.srt.mem_cache.memory_pool import KVWriteLoc
+from sglang.srt.mem_cache.pool.base import KVWriteLoc
 from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.srt.runtime_context import (

python/sglang/srt/layers/attention/xpu_backend.py

@@ -12,7 +12,7 @@
     merge_state_v2_wrapper,
     prepare_swa_spec_page_table_triton,
 )
-from sglang.srt.mem_cache.memory_pool import KVWriteLoc
+from sglang.srt.mem_cache.pool.base import KVWriteLoc
 from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.srt.runtime_context import (

python/sglang/srt/layers/cp/zigzag.py

@@ -51,7 +51,7 @@
 from sglang.srt.layers.dp_attention import (
     is_allocation_symmetric,
 )
-from sglang.srt.mem_cache.memory_pool import KVWriteLoc
+from sglang.srt.mem_cache.pool.base import KVWriteLoc
 from sglang.srt.model_executor.forward_context import get_token_to_kv_pool
 from sglang.srt.runtime_context import get_device, get_parallel
 

python/sglang/srt/mem_cache/allocator/base.py

@@ -21,7 +21,7 @@
 import torch
 
 if TYPE_CHECKING:
-    from sglang.srt.mem_cache.memory_pool import KVCache
+    from sglang.srt.mem_cache.pool.base import KVCache
 
 
 class MambaFullCacheDonor(Protocol):

python/sglang/srt/mem_cache/allocator/paged.py

@@ -34,7 +34,7 @@
 _is_hip = is_hip()
 
 if TYPE_CHECKING:
-    from sglang.srt.mem_cache.memory_pool import KVCache
+    from sglang.srt.mem_cache.pool.base import KVCache
 
 
 def alloc_extend_naive(

python/sglang/srt/mem_cache/allocator/swa.py

@@ -5,7 +5,7 @@
 from sglang.srt.mem_cache.allocator.base import BaseTokenToKVPoolAllocator
 from sglang.srt.mem_cache.allocator.paged import PagedTokenToKVPoolAllocator
 from sglang.srt.mem_cache.allocator.token import TokenToKVPoolAllocator
-from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool
+from sglang.srt.mem_cache.pool.base import BaseSWAKVPool
 from sglang.srt.utils import is_npu
 from sglang.srt.utils.common import get_num_new_pages
 from sglang.srt.utils.invariants import Bucket, Invariant, IsTrue, expect

python/sglang/srt/mem_cache/allocator/token.py

@@ -22,7 +22,7 @@
 from sglang.srt.mem_cache.allocator.base import BaseTokenToKVPoolAllocator
 
 if TYPE_CHECKING:
-    from sglang.srt.mem_cache.memory_pool import KVCache
+    from sglang.srt.mem_cache.pool.base import KVCache
 
 
 class TokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):

python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py

@@ -18,9 +18,8 @@
 from sglang.kernels.ops.attention.dsv4.index_buf_accessor import NopeFp8RopeBf16Pack
 from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE
 from sglang.srt.environ import envs
-from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool
 from sglang.srt.mem_cache.deepseek_v4_compress_state import CompressStatePool
-from sglang.srt.mem_cache.memory_pool import KVCache
+from sglang.srt.mem_cache.pool.base import BaseSWAKVPool, KVCache
 from sglang.srt.runtime_context import get_exec, get_spec
 from sglang.srt.utils import ceil_div, is_hip
 

python/sglang/srt/mem_cache/dsa_cache_layer_split.py

@@ -44,6 +44,8 @@
     RadixAttention,
     get_tensor_size_bytes,
     maybe_detect_oob,
+)
+from sglang.srt.mem_cache.pool.base import (
     unwrap_write_loc,
 )
 from sglang.srt.runtime_context import get_parallel

python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py

@@ -1216,7 +1216,7 @@ def build_hicache_draft_sidecars(
     tree_cache: Any,
 ) -> tuple[list[SidecarPoolSpec], list[PoolEntry]]:
     """Compose the full and SWA draft-sidecar paths."""
-    from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool
+    from sglang.srt.mem_cache.pool.base import BaseSWAKVPool
 
     assert len(draft_device_pools) == 1
     draft_kv_pool = draft_device_pools[0]

python/sglang/srt/mem_cache/kv_cache_configurator.py

@@ -62,7 +62,6 @@
     DSATokenToKVPool,
     HybridLinearKVPool,
     HybridReqToTokenPool,
-    KVCache,
     MHATokenToKVPool,
     MHATokenToKVPoolFP4,
     MHATokenToKVPoolMXFP8,
@@ -73,6 +72,9 @@
     PageMajorMHATokenToKVPool,
     ReqToTokenPool,
 )
+from sglang.srt.mem_cache.pool.base import (
+    KVCache,
+)
 from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
 from sglang.srt.platforms import current_platform
 from sglang.srt.runtime_context import (

python/sglang/srt/mem_cache/kv_index_translator.py

@@ -70,7 +70,7 @@
 from sglang.srt.mem_cache.allocator.unified_mamba import (
     UnifiedMambaTokenToKVPoolAllocator,
 )
-from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool
+from sglang.srt.mem_cache.pool.base import BaseSWAKVPool
 from sglang.srt.runtime_context import get_parallel
 
 

python/sglang/srt/mem_cache/kv_vmm_backing.py

@@ -16,7 +16,7 @@
 )
 
 if TYPE_CHECKING:
-    from sglang.srt.mem_cache.memory_pool import KvBufferDesc
+    from sglang.srt.mem_cache.pool.base import KvBufferDesc
 
 logger = logging.getLogger(__name__)