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 a3cddf597..cef0ef92b 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 @@ -2,7 +2,7 @@ from __future__ import annotations import logging from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any, Callable, Optional +from typing import TYPE_CHECKING, Any, Callable, NamedTuple, Optional from sglang.srt.mem_cache.hicache_storage import ( PoolHitPolicy, @@ -93,6 +93,44 @@ def _with_mtp_layer_mapping( } +class _DeepSeekV4LayerMappings(NamedTuple): + transfer_layer_num: int + full: dict[int, int] + swa: dict[int, int] + c4: dict[int, int] + c128: dict[int, int] + c4_state: dict[int, int] + c4_state_global_layers: list[int] + + +def _resolve_deepseek_v4_layer_mappings( + kvcache: Any, +) -> _DeepSeekV4LayerMappings: + transfer_layer_num = kvcache.end_layer - kvcache.start_layer + full = {layer: layer for layer in range(transfer_layer_num)} + swa = {} if getattr(kvcache, "_unified_kv", False) else full.copy() + + c4, c128, c4_state_global_layers = {}, {}, [] + for local_layer, item in enumerate( + kvcache.layer_mapping[kvcache.start_layer : kvcache.end_layer] + ): + if item.compress_ratio == 4: + c4[local_layer] = item.compress_layer_id + c4_state_global_layers.append(kvcache.start_layer + local_layer) + elif item.compress_ratio == 128: + c128[local_layer] = item.compress_layer_id + + return _DeepSeekV4LayerMappings( + transfer_layer_num=transfer_layer_num, + full=full, + swa=swa, + c4=c4, + c128=c128, + c4_state={layer: index for index, layer in enumerate(c4)}, + c4_state_global_layers=c4_state_global_layers, + ) + + def build_kv_host_pool( *, kv_pool: Any, @@ -444,10 +482,12 @@ def build_deepseek_v4_hicache_stack( model_name: Optional[str] = None, storage_backend_extra_config: Optional[dict] = None, enable_storage_metrics: bool = False, + layer_mappings: Optional[_DeepSeekV4LayerMappings] = None, ) -> tuple[HostPoolGroup, HybridCacheController]: page_size = params.page_size - transfer_layer_num = kvcache.end_layer - kvcache.start_layer - full_layer_mapping = {layer_id: layer_id for layer_id in range(transfer_layer_num)} + layer_mappings = layer_mappings or _resolve_deepseek_v4_layer_mappings(kvcache) + transfer_layer_num = layer_mappings.transfer_layer_num + full_layer_mapping = layer_mappings.full is_unified_kv = getattr(kvcache, "_unified_kv", False) mtp_swa_device_buffers = [] @@ -462,9 +502,7 @@ def build_deepseek_v4_hicache_stack( f"got {len(kvcache.swa_kv_pool.kv_buffer)} buffers for " f"{transfer_layer_num} local layers" ) - swa_layer_mapping = { - layer_id: layer_id for layer_id in range(transfer_layer_num) - } + swa_layer_mapping = layer_mappings.swa # Keep every uncompressed draft SWA layer after the target SWA layers. # NextN has one layer per pool, while DSpark keeps all stages in one pool. mtp_swa_device_buffers = [ @@ -479,24 +517,10 @@ def build_deepseek_v4_hicache_stack( draft_layer_num=len(mtp_swa_device_buffers), ) - c4_layer_mapping = {} - c128_layer_mapping = {} - c4_state_local_layers = [] - c4_state_global_layers = [] - for local_layer_id, layer_item in enumerate( - kvcache.layer_mapping[kvcache.start_layer : kvcache.end_layer] - ): - global_layer_id = kvcache.start_layer + local_layer_id - if layer_item.compress_ratio == 4: - c4_layer_mapping[local_layer_id] = layer_item.compress_layer_id - c4_state_local_layers.append(local_layer_id) - c4_state_global_layers.append(global_layer_id) - elif layer_item.compress_ratio == 128: - c128_layer_mapping[local_layer_id] = layer_item.compress_layer_id - - c4_state_mapping = { - layer_id: local_id for local_id, layer_id in enumerate(c4_state_local_layers) - } + c4_layer_mapping = layer_mappings.c4 + c128_layer_mapping = layer_mappings.c128 + c4_state_mapping = layer_mappings.c4_state + c4_state_global_layers = layer_mappings.c4_state_global_layers num_host_pages, swa_num_host_pages = _deepseek_v4_num_host_pages( params=params, kvcache=kvcache, @@ -1164,6 +1188,18 @@ class StackStrategy: def matches(self, kvcache: Any, components: set[ComponentType]) -> bool: raise NotImplementedError + def build_direct_linker_pool_group( + self, + *, + kvcache: Any, + params: CacheInitParams, + page_size: int, + ): + raise ValueError( + "The selected hybrid pool strategy does not support the direct " + f"external linker: {type(self).__name__}." + ) + def build( self, *, @@ -1192,6 +1228,13 @@ class _DeepSeekV4Strategy(StackStrategy): ComponentType.SWA, } + def build_direct_linker_pool_group(self, *, kvcache, params, page_size): + from sglang.srt.mem_cache.hybrid_cache.linker_pool_assembler import ( + _build_deepseek_v4_device_pool_group, + ) + + return _build_deepseek_v4_device_pool_group(kvcache, page_size) + def build( self, *, @@ -1206,6 +1249,7 @@ class _DeepSeekV4Strategy(StackStrategy): model_name=None, enable_storage_metrics=False, ): + layer_mappings = _resolve_deepseek_v4_layer_mappings(kvcache) host_pool_group, cache_controller = build_deepseek_v4_hicache_stack( params=params, kvcache=kvcache, @@ -1217,6 +1261,7 @@ class _DeepSeekV4Strategy(StackStrategy): model_name=model_name, storage_backend_extra_config=storage_backend_extra_config, enable_storage_metrics=enable_storage_metrics, + layer_mappings=layer_mappings, ) sidecars = [ SidecarPoolSpec( @@ -1454,6 +1499,13 @@ class _DsaStrategy(StackStrategy): ComponentType.FULL } + def build_direct_linker_pool_group(self, *, kvcache, params, page_size): + from sglang.srt.mem_cache.hybrid_cache.linker_pool_assembler import ( + _build_dsa_device_pool_group, + ) + + return _build_dsa_device_pool_group(kvcache, page_size) + def build( self, *, diff --git a/python/sglang/srt/mem_cache/hybrid_cache/linker_pool_assembler.py b/python/sglang/srt/mem_cache/hybrid_cache/linker_pool_assembler.py new file mode 100644 index 000000000..9fbf01554 --- /dev/null +++ b/python/sglang/srt/mem_cache/hybrid_cache/linker_pool_assembler.py @@ -0,0 +1,380 @@ +"""Materialize hybrid device pools for external cache linkers.""" + +from __future__ import annotations + +from collections.abc import Callable, Sequence +from dataclasses import replace +from typing import Any + +import torch + +from sglang.srt.mem_cache.hicache_storage import ( + PoolHitPolicy, + PoolName, + PoolTransfer, +) +from sglang.srt.mem_cache.unified_cache.component_type import ComponentType + + +class DevicePoolEntry: + """Zero-copy linker view over one physical device pool.""" + + def __init__( + self, + *, + name: PoolName, + indices_from_pool: PoolName, + device_pool: Any, + components: Sequence[Sequence[torch.Tensor]], + layer_mapping: dict[int, int], + page_size: int, + rows_are_pages: bool, + packed: bool = True, + index_mapper: Callable[[torch.Tensor], torch.Tensor] | None = None, + ): + self.name = name + self.indices_from_pool = indices_from_pool + self.device_pool = device_pool + self.components = [list(component) for component in components] + self.layer_mapping = layer_mapping + self.page_size = page_size + self.packed = packed + self._index_mapper = index_mapper + self._page_offsets = torch.arange(page_size) + self._row_span = 1 if rows_are_pages else page_size + + if not self.components or any(not component for component in self.components): + raise ValueError(f"Device pool {name} has no storage buffers.") + self.kv_buffer = [buffer for group in self.components for buffer in group] + self._row_count = min(buffer.shape[0] for buffer in self.kv_buffer) + + self.buffer_meta = [ + [ + ( + buffer.data_ptr(), + buffer.stride(0) * buffer.element_size(), + buffer.nbytes // buffer.shape[0] * self._row_span, + ) + for buffer in component + ] + for component in self.components + ] + + self._component_offsets = [] + offset = 0 + for component in self.buffer_meta: + if not packed: + offset = 0 + offsets = [] + for _, _, size in component: + offsets.append(offset) + offset += size + self._component_offsets.append(offsets) + + def get_hybrid_pool_buffer(self) -> list[torch.Tensor]: + return self.kv_buffer + + def translate_indices(self, indices: torch.Tensor) -> torch.Tensor: + return self._index_mapper(indices) if self._index_mapper else indices + + def _rows(self, indices: torch.Tensor) -> list[int]: + slots = indices.detach().to(device="cpu", dtype=torch.int64).flatten() + if slots.numel() % self.page_size: + raise ValueError( + f"Pool {self.name} got {slots.numel()} indices, expected a " + f"multiple of page_size={self.page_size}." + ) + if not slots.numel(): + return [] + + pages = slots.reshape(-1, self.page_size) + starts = pages[:, 0] + if torch.any(starts.remainder(self.page_size)) or not torch.equal( + pages, starts[:, None] + self._page_offsets + ): + raise ValueError(f"Pool {self.name} requires aligned contiguous pages.") + rows = ( + starts.div(self.page_size, rounding_mode="floor") + if self._row_span == 1 + else starts + ) + first_row = int(rows.min()) + last_row = int(rows.max()) + self._row_span + if first_row < 0 or last_row > self._row_count: + raise ValueError( + f"Pool {self.name} row range [{first_row}, {last_row}) exceeds " + f"buffer shapes {[tuple(buffer.shape) for buffer in self.kv_buffer]}." + ) + return rows.tolist() + + def get_page_buffer_meta(self, indices: torch.Tensor): + rows = self._rows(indices) + ptrs = [ + base_ptr + row * row_stride + for row in rows + for component in self.buffer_meta + for base_ptr, row_stride, _ in component + ] + sizes = [ + size + for _ in rows + for component in self.buffer_meta + for _, _, size in component + ] + return ptrs, sizes + + def prepare_locations(self, indices: torch.Tensor) -> list[int]: + return self._rows(indices) + + def get_prepared_layer_range_meta(self, locations: list[int], layer: int): + buffer_index = self.layer_mapping.get(layer) + if buffer_index is None: + return None + + items = [] + for component, offsets in zip(self.buffer_meta, self._component_offsets): + base_ptr, row_stride, size = component[buffer_index] + items.append((base_ptr, row_stride, size, offsets[buffer_index])) + + ptrs, sizes, offsets = [], [], [] + for row in locations: + row_ptrs = [ + base_ptr + row * row_stride for base_ptr, row_stride, _, _ in items + ] + row_sizes = [size for _, _, size, _ in items] + row_offsets = [offset for _, _, _, offset in items] + if self.packed: + ptrs.append(row_ptrs) + sizes.append(row_sizes) + offsets.append(row_offsets) + else: + ptrs.extend([[value] for value in row_ptrs]) + sizes.extend([[value] for value in row_sizes]) + offsets.extend([[value] for value in row_offsets]) + return ptrs, sizes, offsets + + +class DevicePoolGroup: + """Physical device pools sharing one logical linker layer range.""" + + def __init__( + self, + entries: Sequence[DevicePoolEntry], + num_layers: int, + page_size: int, + *, + rank_replicated: bool = False, + ): + self.entries = list(entries) + self.entry_map = {entry.name: entry for entry in entries} + if len(self.entries) != len(self.entry_map): + raise ValueError("DevicePoolGroup contains duplicate pool names.") + self.sources = {entry.name: entry.indices_from_pool for entry in self.entries} + self.num_layers = num_layers + self.page_size = page_size + self.rank_replicated = rank_replicated + self.kv_buffer = None + + def resolve_transfers( + self, + transfers: list[PoolTransfer], + *, + allow_partial: bool = False, + allow_missing_kv: bool = False, + ) -> list[PoolTransfer]: + """Expand logical component transfers into physical device pools.""" + by_name = {transfer.name: transfer for transfer in transfers} + kv = by_name.get(PoolName.KV) + if not any(transfer.keys for transfer in transfers): + return [] + if not allow_missing_kv and (kv is None or not kv.keys): + return [] + if not allow_partial and not set(self.sources.values()) <= set(by_name): + return [] + + resolved = [] + for name, source_name in self.sources.items(): + source = by_name.get(source_name) + if source is None or not source.keys: + continue + indices = source.device_indices + resolved.append( + replace( + source, + name=name, + host_indices=( + self.entry_map[name].translate_indices(indices) + if indices is not None + else None + ), + keys=list(source.keys), + hit_policy=( + PoolHitPolicy.ALL_PAGES + if source_name == PoolName.KV + else source.hit_policy + ), + indices_from_pool=None, + ) + ) + return resolved + + +def _deepseek_v4_state_views(state_pools: list[Any], global_layers: list[int]): + views = [] + for layer in global_layers: + pool = state_pools[layer] + state = pool.kv_score_buffer.kv_score + ring = int(pool.ring_size) + usable = state.shape[0] // ring * ring + views.append( + state.view(torch.uint8) + .reshape(state.shape[0], -1)[:usable] + .reshape(usable // ring, -1) + ) + return views + + +def _build_deepseek_v4_device_pool_group( + kvcache: Any, page_size: int +) -> DevicePoolGroup: + from sglang.srt.mem_cache.deepseek_v4_memory_pool import HiSparseC4DevicePool + from sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler import ( + _resolve_deepseek_v4_layer_mappings, + ) + + mappings = _resolve_deepseek_v4_layer_mappings(kvcache) + if getattr(kvcache, "_unified_kv", False) or isinstance( + kvcache.c4_kv_pool, HiSparseC4DevicePool + ): + raise ValueError( + "The direct external linker does not support unified-KV or HiSparse." + ) + if kvcache.swa_page_size != page_size: + raise ValueError( + "DeepSeek V4 SWA page size must match the tree page size: " + f"{kvcache.swa_page_size} != {page_size}." + ) + + entries = [ + DevicePoolEntry( + name=PoolName.SWA, + indices_from_pool=PoolName.SWA, + device_pool=kvcache.swa_kv_pool, + components=[kvcache.swa_kv_pool.kv_buffer], + layer_mapping=mappings.swa, + page_size=page_size, + rows_are_pages=True, + ) + ] + + def add(name, source, pool, buffers, layer_mapping): + if layer_mapping: + entries.append( + DevicePoolEntry( + name=name, + indices_from_pool=source, + device_pool=pool, + components=[buffers], + layer_mapping=layer_mapping, + page_size=page_size, + rows_are_pages=True, + ) + ) + + add( + PoolName.DEEPSEEK_V4_C4, + PoolName.KV, + kvcache.c4_kv_pool, + kvcache.c4_kv_pool.kv_buffer, + mappings.c4, + ) + add( + PoolName.DEEPSEEK_V4_C4_INDEXER, + PoolName.KV, + kvcache.c4_indexer_kv_pool, + kvcache.c4_indexer_kv_pool.index_k_with_scale_buffer, + mappings.c4, + ) + add( + PoolName.DEEPSEEK_V4_C128, + PoolName.KV, + kvcache.c128_kv_pool, + kvcache.c128_kv_pool.kv_buffer, + mappings.c128, + ) + add( + PoolName.DEEPSEEK_V4_C4_STATE, + PoolName.SWA, + kvcache.compress_state_pools, + _deepseek_v4_state_views( + kvcache.compress_state_pools, + mappings.c4_state_global_layers, + ), + mappings.c4_state, + ) + add( + PoolName.DEEPSEEK_V4_C4_INDEXER_STATE, + PoolName.SWA, + kvcache.indexer_compress_state_pools, + _deepseek_v4_state_views( + kvcache.indexer_compress_state_pools, + mappings.c4_state_global_layers, + ), + mappings.c4_state, + ) + return DevicePoolGroup( + entries, + mappings.transfer_layer_num, + page_size, + rank_replicated=True, + ) + + +def _build_dsa_device_pool_group(kvcache: Any, page_size: int) -> DevicePoolGroup: + if kvcache.page_size != page_size: + raise ValueError( + "DSA KV page size must match the tree page size: " + f"{kvcache.page_size} != {page_size}." + ) + num_layers = kvcache.layer_num + identity = {layer: layer for layer in range(num_layers)} + entries = [ + DevicePoolEntry( + name=PoolName.KV, + indices_from_pool=PoolName.KV, + device_pool=kvcache, + components=[kvcache.kv_buffer], + layer_mapping=identity, + page_size=page_size, + rows_are_pages=False, + ), + DevicePoolEntry( + name=PoolName.INDEXER, + indices_from_pool=PoolName.KV, + device_pool=kvcache, + components=[kvcache.index_k_with_scale_buffer], + layer_mapping=identity, + page_size=page_size, + rows_are_pages=True, + ), + ] + return DevicePoolGroup(entries, num_layers, page_size, rank_replicated=True) + + +def resolve_hybrid_device_pool_group( + *, + kvcache: Any, + page_size: int, + params: Any, + components: set[ComponentType], +) -> DevicePoolGroup: + """Materialize a direct-linker pool group through the hybrid registry.""" + from sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler import ( + _select_strategy, + ) + + return _select_strategy(kvcache, components).build_direct_linker_pool_group( + kvcache=kvcache, + params=params, + page_size=page_size, + ) diff --git a/test/registered/unit/mem_cache/test_linker_pool_assembler.py b/test/registered/unit/mem_cache/test_linker_pool_assembler.py new file mode 100644 index 000000000..916c300b7 --- /dev/null +++ b/test/registered/unit/mem_cache/test_linker_pool_assembler.py @@ -0,0 +1,291 @@ +"""Unit tests for external-linker device pool assembly.""" + +import unittest +from types import SimpleNamespace + +import torch + +from sglang.srt.mem_cache.hicache_storage import ( + PoolHitPolicy, + PoolName, + PoolTransfer, +) +from sglang.srt.mem_cache.hybrid_cache.linker_pool_assembler import ( + DevicePoolEntry, + DevicePoolGroup, + resolve_hybrid_device_pool_group, +) +from sglang.srt.mem_cache.unified_cache.component_type import ComponentType +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=2, suite="base-a-test-cpu") + + +class TestDevicePoolEntry(CustomTestCase): + def test_sparse_multi_component_layer_ranges(self): + k0 = torch.zeros((8, 3), dtype=torch.uint8) + k2 = torch.zeros((8, 5), dtype=torch.uint8) + v0 = torch.zeros((8, 7), dtype=torch.uint8) + v2 = torch.zeros((8, 11), dtype=torch.uint8) + pool = DevicePoolEntry( + name=PoolName.KV, + indices_from_pool=PoolName.KV, + device_pool=None, + components=[[k0, k2], [v0, v2]], + layer_mapping={0: 0, 2: 1}, + page_size=2, + rows_are_pages=False, + packed=False, + ) + + indices = torch.tensor([0, 1, 4, 5]) + locations = pool.prepare_locations(indices) + self.assertEqual(locations, [0, 4]) + pointers, sizes = pool.get_page_buffer_meta(indices) + self.assertEqual( + pointers, + [ + buffer[row].data_ptr() + for row in locations + for buffer in (k0, k2, v0, v2) + ], + ) + self.assertEqual(sizes, [6, 10, 14, 22] * 2) + self.assertIsNone(pool.get_prepared_layer_range_meta(locations, 1)) + + pointers, sizes, offsets = pool.get_prepared_layer_range_meta(locations, 2) + self.assertEqual( + pointers, + [ + [k2[0].data_ptr()], + [v2[0].data_ptr()], + [k2[4].data_ptr()], + [v2[4].data_ptr()], + ], + ) + self.assertEqual(sizes, [[10], [22], [10], [22]]) + self.assertEqual(offsets, [[6], [14], [6], [14]]) + + def test_rejects_invalid_pages_and_empty_buffers(self): + with self.assertRaisesRegex(ValueError, "has no storage buffers"): + DevicePoolEntry( + name=PoolName.KV, + indices_from_pool=PoolName.KV, + device_pool=None, + components=[], + layer_mapping={}, + page_size=2, + rows_are_pages=False, + ) + + pool = DevicePoolEntry( + name=PoolName.KV, + indices_from_pool=PoolName.KV, + device_pool=None, + components=[[torch.zeros((8, 3), dtype=torch.uint8)]], + layer_mapping={0: 0}, + page_size=2, + rows_are_pages=False, + ) + for indices, error in ( + (torch.tensor([0]), "multiple of page_size"), + (torch.tensor([1, 2]), "aligned contiguous pages"), + (torch.tensor([0, 2]), "aligned contiguous pages"), + (torch.tensor([8, 9]), "exceeds buffer shapes"), + ): + with self.subTest(indices=indices.tolist()): + with self.assertRaisesRegex(ValueError, error): + pool.prepare_locations(indices) + + +class TestDevicePoolGroup(CustomTestCase): + def test_resolve_transfers_expands_physical_pools(self): + entries = [ + SimpleNamespace( + name=PoolName.KV, + indices_from_pool=PoolName.KV, + translate_indices=lambda indices: indices, + ), + SimpleNamespace( + name=PoolName.INDEXER, + indices_from_pool=PoolName.KV, + translate_indices=lambda indices: indices + 100, + ), + ] + group = DevicePoolGroup(entries, num_layers=2, page_size=2) + transfer = PoolTransfer( + name=PoolName.KV, + keys=["a", "b"], + device_indices=torch.tensor([0, 1, 4, 5]), + hit_policy=PoolHitPolicy.TRAILING_PAGES, + ) + + resolved = group.resolve_transfers([transfer]) + + self.assertEqual( + [item.name for item in resolved], [PoolName.KV, PoolName.INDEXER] + ) + self.assertEqual(resolved[0].host_indices.tolist(), [0, 1, 4, 5]) + self.assertEqual(resolved[1].host_indices.tolist(), [100, 101, 104, 105]) + self.assertTrue( + all(item.hit_policy == PoolHitPolicy.ALL_PAGES for item in resolved) + ) + + def test_partial_side_pool_requires_explicit_opt_in(self): + entry = SimpleNamespace( + name=PoolName.SWA, + indices_from_pool=PoolName.SWA, + translate_indices=lambda indices: indices + 100, + ) + group = DevicePoolGroup([entry], num_layers=1, page_size=2) + transfer = PoolTransfer( + name=PoolName.SWA, + keys=["b", "d"], + device_indices=torch.tensor([20, 21, 24, 25]), + hit_policy=PoolHitPolicy.TRAILING_PAGES, + ) + + self.assertEqual(group.resolve_transfers([transfer]), []) + resolved = group.resolve_transfers( + [transfer], allow_partial=True, allow_missing_kv=True + ) + + self.assertEqual(len(resolved), 1) + self.assertEqual(resolved[0].name, PoolName.SWA) + self.assertEqual(resolved[0].keys, ["b", "d"]) + self.assertEqual(resolved[0].host_indices.tolist(), [120, 121, 124, 125]) + self.assertEqual(resolved[0].hit_policy, PoolHitPolicy.TRAILING_PAGES) + + +class TestHybridDevicePoolAssembler(CustomTestCase): + def test_deepseek_v4_maps_sparse_sidecars(self): + from sglang.srt.mem_cache.deepseek_v4_memory_pool import ( + DeepSeekV4LayerItem, + DeepSeekV4TokenToKVPool, + ) + + def state_pool(): + return SimpleNamespace( + ring_size=2, + kv_score_buffer=SimpleNamespace(kv_score=torch.zeros((8, 3))), + ) + + kvcache = DeepSeekV4TokenToKVPool.__new__(DeepSeekV4TokenToKVPool) + kvcache._unified_kv = False + kvcache.start_layer = 1 + kvcache.end_layer = 4 + kvcache.swa_page_size = 2 + kvcache.swa_kv_pool = SimpleNamespace( + kv_buffer=[torch.zeros((8, 3), dtype=torch.uint8) for _ in range(3)] + ) + kvcache.c4_kv_pool = SimpleNamespace( + kv_buffer=[torch.zeros((8, 5), dtype=torch.uint8) for _ in range(2)] + ) + kvcache.c4_indexer_kv_pool = SimpleNamespace( + index_k_with_scale_buffer=[ + torch.zeros((8, 7), dtype=torch.uint8) for _ in range(2) + ] + ) + kvcache.c128_kv_pool = SimpleNamespace( + kv_buffer=[torch.zeros((8, 11), dtype=torch.uint8)] + ) + kvcache.layer_mapping = [ + DeepSeekV4LayerItem(0, -1), + DeepSeekV4LayerItem(4, 0), + DeepSeekV4LayerItem(128, 0), + DeepSeekV4LayerItem(4, 1), + ] + kvcache.compress_state_pools = [None, state_pool(), None, state_pool()] + kvcache.indexer_compress_state_pools = [ + None, + state_pool(), + None, + state_pool(), + ] + + group = resolve_hybrid_device_pool_group( + kvcache=kvcache, + page_size=2, + params=SimpleNamespace(), + components={ComponentType.FULL, ComponentType.SWA}, + ) + + self.assertEqual(group.num_layers, 3) + self.assertTrue(group.rank_replicated) + self.assertEqual( + set(group.entry_map), + { + PoolName.SWA, + PoolName.DEEPSEEK_V4_C4, + PoolName.DEEPSEEK_V4_C4_INDEXER, + PoolName.DEEPSEEK_V4_C128, + PoolName.DEEPSEEK_V4_C4_STATE, + PoolName.DEEPSEEK_V4_C4_INDEXER_STATE, + }, + ) + self.assertEqual(group.sources[PoolName.DEEPSEEK_V4_C4], PoolName.KV) + self.assertEqual(group.sources[PoolName.DEEPSEEK_V4_C4_STATE], PoolName.SWA) + + c4_pool = group.entry_map[PoolName.DEEPSEEK_V4_C4] + pointers, sizes = c4_pool.get_page_buffer_meta(torch.tensor([0, 1])) + self.assertEqual(len(pointers), 2) + self.assertEqual(sizes, [5, 5]) + _, sizes, offsets = c4_pool.get_prepared_layer_range_meta([0], 2) + self.assertEqual(sizes, [[5]]) + self.assertEqual(offsets, [[5]]) + self.assertIsNone(c4_pool.get_prepared_layer_range_meta([0], 1)) + + def test_dsa_uses_hybrid_assembler_strategy(self): + from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool + + kvcache = DSATokenToKVPool.__new__(DSATokenToKVPool) + kvcache.page_size = 2 + kvcache.layer_num = 2 + kvcache.kv_buffer = [ + torch.zeros((8, 3), dtype=torch.uint8), + torch.zeros((8, 5), dtype=torch.uint8), + ] + kvcache.index_key_cache = SimpleNamespace( + buffer=[ + torch.zeros((4, 7), dtype=torch.uint8), + torch.zeros((4, 11), dtype=torch.uint8), + ] + ) + + group = resolve_hybrid_device_pool_group( + kvcache=kvcache, + page_size=2, + params=SimpleNamespace(), + components={ComponentType.FULL}, + ) + + self.assertEqual(group.num_layers, 2) + self.assertTrue(group.rank_replicated) + self.assertEqual(set(group.entry_map), {PoolName.KV, PoolName.INDEXER}) + self.assertEqual( + group.sources, + { + PoolName.KV: PoolName.KV, + PoolName.INDEXER: PoolName.KV, + }, + ) + + def test_unsupported_strategy_fails_with_context(self): + from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool + + kvcache = HybridLinearKVPool.__new__(HybridLinearKVPool) + with self.assertRaisesRegex( + ValueError, + "does not support the direct external linker: _MambaStrategy", + ): + resolve_hybrid_device_pool_group( + kvcache=kvcache, + page_size=2, + params=SimpleNamespace(), + components={ComponentType.FULL, ComponentType.MAMBA}, + ) + + +if __name__ == "__main__": + unittest.main()