From ca8ecc6a6f06426a8226ffc40403872f647ca55f Mon Sep 17 00:00:00 2001 From: Cheng Xin <77869622+cx22757@users.noreply.github.com> Date: Mon, 14 Sep 2026 10:38:50 +0800 Subject: [PATCH] [NPU] Support DSV4 host memory cache management (#37382) --- .../npu/dsv4/c128_sidecar_component.py | 259 ++++++++++++-- .../npu/dsv4/dsv4_allocator.py | 82 ++++- .../npu/dsv4/dsv4_memory_pool.py | 8 + python/sglang/srt/managers/schedule_batch.py | 15 +- python/sglang/srt/mem_cache/allocation.py | 6 + python/sglang/srt/mem_cache/allocator/base.py | 13 +- .../hybrid_cache/hybrid_pool_assembler.py | 158 ++++++--- .../sglang/srt/mem_cache/memory_pool_host.py | 191 ++++++++++- .../storage/mooncake_store/mooncake_store.py | 4 +- .../srt/mem_cache/unified_radix_cache.py | 42 ++- .../unit/mem_cache/test_dsv4_hicache_l2.py | 322 ++++++++++++++++++ .../mem_cache/test_hicache_host_register.py | 4 +- 12 files changed, 1010 insertions(+), 94 deletions(-) create mode 100644 test/registered/unit/mem_cache/test_dsv4_hicache_l2.py diff --git a/python/sglang/srt/hardware_backend/npu/dsv4/c128_sidecar_component.py b/python/sglang/srt/hardware_backend/npu/dsv4/c128_sidecar_component.py index a7b338d98..9fdb0684d 100644 --- a/python/sglang/srt/hardware_backend/npu/dsv4/c128_sidecar_component.py +++ b/python/sglang/srt/hardware_backend/npu/dsv4/c128_sidecar_component.py @@ -6,7 +6,7 @@ radix tree; partial tail pages remain request-owned. from __future__ import annotations -from typing import TYPE_CHECKING, Callable, Optional +from typing import TYPE_CHECKING, Callable, Optional, Sequence import torch @@ -16,14 +16,21 @@ from sglang.srt.mem_cache.base_prefix_cache import ( MatchPrefixParams, MatchResult, ) +from sglang.srt.mem_cache.hicache_storage import ( + PoolName, + PoolTransfer, + PoolTransferResult, +) from sglang.srt.mem_cache.unified_cache.cache_action import ( FreeComponentDeviceSlot, SWARebuild, ) from sglang.srt.mem_cache.unified_cache.components import ( BASE_COMPONENT_TYPE, + CacheTransferPhase, ComponentType, EvictLayer, + PrepareLoadBackResult, TreeComponent, ) @@ -33,14 +40,16 @@ if TYPE_CHECKING: CacheAction, ComponentAction, ) - from sglang.srt.mem_cache.unified_radix_cache import ( - UnifiedTreeNode, - ) + from sglang.srt.mem_cache.unified_radix_cache import UnifiedTreeNode class C128SidecarComponent(TreeComponent): component_type = ComponentType.C128 + # Bound by _apply_stack_result (hybrid_pool_assembler) on the NPU path: + # C128 is an independent-index pool whose host values live in this pool. + _c128_kv_pool_host = None + @property def allocator(self): return self.cache.token_to_kv_pool_allocator @@ -97,8 +106,31 @@ class C128SidecarComponent(TreeComponent): def create_match_validator( self, match_device_only: bool = False ) -> Callable[[UnifiedTreeNode], bool]: - # A page is attached only to the node ending its full physical group. - return lambda node: node.component_data[self.component_type].value is not None + # Pages attach only to full-group endpoints. A host-backed endpoint is a + # valid match unless match_device_only requires device residency. + def _valid(node: UnifiedTreeNode) -> bool: + cd = node.component_data[self.component_type] + if match_device_only: + return cd.value is not None + return cd.value is not None or cd.host_value is not None + + return _valid + + def _collect_device_pages(self, node_id: int) -> torch.Tensor: + chunks = [] + node = self.tree_core.node_by_id(node_id) + root = self.tree_core.root_node + while node is not root: + value = node.component_data[self.component_type].value + if value is not None: + chunks.append(value) + node = node.parent + chunks.reverse() + return ( + torch.cat(chunks) + if chunks + else self.allocator.c128_attn_allocator.free_pages.new_empty((0,)) + ) def finalize_match_result_in_cache( self, params: MatchPrefixParams, result: MatchResult @@ -107,31 +139,52 @@ class C128SidecarComponent(TreeComponent): if req is None: return result - chunks = [] - node = self.tree_core.node_by_id(result.best_match_node) - root = self.tree_core.root_node - while node is not root: - value = node.component_data[self.component_type].value - if value is not None: - chunks.append(value) - node = node.parent - chunks.reverse() - pages = ( - torch.cat(chunks) - if chunks - else self.allocator.c128_attn_allocator.free_pages.new_empty((0,)) - ) - group_tokens = 128 * self.allocator.c128_attn_allocator.page_size - assert pages.numel() == len(result.device_indices) // group_tokens + pages = self._collect_device_pages(result.best_match_node) self.cache.req_to_token_pool.set_c128_prefix_pages(req, pages) return result + def finalize_match_result_in_tree_core( + self, + result: MatchResult, + params: MatchPrefixParams, + value_chunks: list[torch.Tensor], + best_value_len: int, + ) -> MatchResult: + branching_seqlen = result.swa_branching_seqlen + if branching_seqlen is None: + return result + + # SWA initially derives its branching point from the Full-KV hit and the + # radix page size. C128 can only resume from a complete compression-group + # endpoint, so constrain that candidate to the nearest C128 boundary. + group_tokens = 128 * self.allocator.c128_attn_allocator.page_size + branching_seqlen = branching_seqlen // group_tokens * group_tokens + current_boundary = len(result.device_indices) + result.host_hit_length + return result._replace( + swa_branching_seqlen=( + branching_seqlen if branching_seqlen > current_boundary else None + ) + ) + + def finalize_load_back( + self, req: Optional[Req], prep: PrepareLoadBackResult, success: bool + ) -> None: + if not success or req is None: + return + + # match_prefix runs before load-back, so a host-only C128 endpoint is + # absent from the request-local page table built by the match finalizer. + # Refresh it after commit_load_back attaches the restored page to the tree. + pages = self._collect_device_pages(req.best_match_node) + self.cache.req_to_token_pool.set_c128_prefix_pages(req, pages) + def recover_after_unevict( self, node: UnifiedTreeNode, prefix_len: int, total_prefix_len: int, params: InsertParams, + result: InsertResult, cache_actions: list[CacheAction | ComponentAction], ) -> None: pages = params.c128_value @@ -236,13 +289,26 @@ class C128SidecarComponent(TreeComponent): ) -> tuple[int, int]: cd = node.component_data[self.component_type] if EvictLayer.DEVICE in target and cd.value is not None: + # Device pages use retain/release_c128_pages refcounts. + # _drain_device_frees converts these IDs into release actions. device_frees[self.component_type].append(cd.value) self.tree_core.component_evictable_size_[self.component_type] -= len( cd.value ) cd.value = None - # C128 pages are auxiliary to Full tokens and must not inflate the - # public token-eviction count. + # A device tombstone with a host copy makes the node host-only. + # Promote it so every host-only node remains in host_lru. + if cd.host_value is not None: + host_lru = self.tree_core.host_lru_lists[self.component_type] + if not host_lru.in_list(node): + host_lru.insert_mru(node) + if EvictLayer.HOST in target and cd.host_value is not None: + # Host values have no refcount; free_host_values returns them directly. + host_frees[self.component_type].append(cd.host_value) + cd.host_value = None + # C128 pages are auxiliary to Full tokens and must not inflate the public + # token-eviction count; the host return also stays 0 so a FULL host-leaf + # eviction's tracker only counts FULL tokens (C128 is a required payload). return 0, 0 def prepare_for_caching_req( @@ -285,10 +351,151 @@ class C128SidecarComponent(TreeComponent): pass def acquire_component_lock(self, node, result, lock_host=False): + # Device path is a no-op: C128 device pages are owned via refcount, not + # the FULL path-lock. Host path mirrors FULL's single-node host lock. + if lock_host: + cd = node.component_data[self.component_type] + # write_back mode: the anchor may be device-only (no host_value); + # pin it anyway. + if cd.host_value is None and not self.tree_core.is_write_back: + return result + cd.host_lock_ref += 1 + self.tree_core._update_evictable_leaf_sets(node) return result def release_component_lock(self, node, params, lock_host=False) -> None: - pass + if lock_host: + cd = node.component_data[self.component_type] + if cd.host_lock_ref == 0: + return + # Mirror of `acquire`. write_back uses a pure counter. + if cd.host_value is None and not self.tree_core.is_write_back: + return + cd.host_lock_ref -= 1 + self.tree_core._update_evictable_leaf_sets(node) def free_host_values(self, host_values) -> None: - pass + if self._c128_kv_pool_host is None: + return + for host_value in host_values: + self._c128_kv_pool_host.free(host_value) + + # ---- HiCache Hooks ---- + + @staticmethod + def _expand_page_indices(page_ids: torch.Tensor, page_size: int) -> torch.Tensor: + """Expand each page ID to ``page_id * page_size + arange(page_size)``.""" + page_ids = page_ids.view(-1) + if page_ids.numel() == 0: + return page_ids.new_empty((0,), dtype=torch.int64) + return ( + page_ids[:, None] * page_size + + torch.arange(page_size, device=page_ids.device) + ).flatten() + + def build_hicache_transfers( + self, + node: UnifiedTreeNode, + phase: CacheTransferPhase, + *, + mamba_pool_idx: Optional[torch.Tensor] = None, + host_indices: Optional[torch.Tensor] = None, + token_ids: Optional[Sequence[int]] = None, + prefetch_tokens: int = 0, + last_hash: Optional[str] = None, + ) -> Optional[list[PoolTransfer]]: + ct = self.component_type + page_size = self.allocator.c128_attn_allocator.page_size + + if phase == CacheTransferPhase.BACKUP_HOST: + # Back up the C128 pages attached to this node (its group endpoints). + # The transfer is independent (indices_from_pool=None): the controller + # allocates C128 host slots of len(device_indices) = groups * P. + page_ids = node.component_data[ct].value + if page_ids is None or page_ids.numel() == 0: + return None + return [ + PoolTransfer( + name=PoolName.DEEPSEEK_V4_C128, + indices_from_pool=None, + device_indices=self._expand_page_indices(page_ids, page_size), + nodes_to_load=[node.id], + ) + ] + + if phase == CacheTransferPhase.LOAD_BACK: + # Collect host values from complete-group endpoints on the evicted path. + # For example, G groups produce G * page_size host and device indices. + backed_up: list[torch.Tensor] = [] + nodes: list[UnifiedTreeNode] = [] + cur = node + while cur is not self.tree_core.root_node and cur.evicted: + cd = cur.component_data[ct] + if cd.host_value is not None: + backed_up.append(cd.host_value) + nodes.append(cur) + cur = cur.parent + if not backed_up: + return None + backed_up.reverse() + nodes.reverse() + return [ + PoolTransfer( + name=PoolName.DEEPSEEK_V4_C128, + indices_from_pool=None, + host_indices=torch.cat(backed_up), + device_indices=None, + nodes_to_load=[n.id for n in nodes], + ) + ] + + return None + + def commit_hicache_transfer( + self, + node: UnifiedTreeNode, + phase: CacheTransferPhase, + transfers: list[PoolTransfer] = (), + *, + cache_actions: list[CacheAction | ComponentAction], + insert_result: Optional[InsertResult] = None, + pool_storage_result: Optional[PoolTransferResult] = None, + ) -> None: + ct = self.component_type + page_size = self.allocator.c128_attn_allocator.page_size + + if phase == CacheTransferPhase.BACKUP_HOST: + # Publish the controller-allocated host slots as this node's C128 + # host residency (L2). The device value stays until demote/evict. + if transfers and transfers[0].host_indices is not None: + node.component_data[ct].host_value = transfers[0].host_indices.clone() + # An evict-then-backup race can make this node host-only. + # Insert it now so every host-only node remains in host_lru. + if node.component_data[ct].value is None: + host_lru = self.tree_core.host_lru_lists[ct] + if not host_lru.in_list(node): + host_lru.insert_mru(node) + return + + if phase == CacheTransferPhase.LOAD_BACK: + if not transfers or transfers[0].device_indices is None: + return + xfer = transfers[0] + device_indices = xfer.device_indices + offset = 0 + for nid in xfer.nodes_to_load or []: + n = self.tree_core.node_by_id(nid) + cd = n.component_data[ct] + n_len = len(cd.host_value) + # Each page occupies P consecutive expanded slots; ``// P`` yields + # P copies of the page id, so unique() recovers the distinct page + # ids (the allocator's retain_c128_pages does NOT dedup). + page_ids = torch.unique( + device_indices[offset : offset + n_len] // page_size + ) + # Once retained, tree eviction releases these pages. + # The controller frees them only on rollback before commit. + self.allocator.retain_c128_pages(page_ids) + self.tree_core.set_component_device_value(nid, ct, page_ids.clone()) + offset += n_len + return diff --git a/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_allocator.py b/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_allocator.py index 6b9d8913d..1069bbd3f 100644 --- a/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_allocator.py +++ b/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_allocator.py @@ -30,7 +30,12 @@ from sglang.srt.hardware_backend.npu.dsv4.dsv4_common_hooks import ( maybe_write_dsv4_extend, ) from sglang.srt.mem_cache.allocation import alloc_paged_token_slots_extend +from sglang.srt.mem_cache.allocation_sizing import ( + get_alloc_reserve_per_decode, + page_aligned_decode_alloc_lens, +) from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator +from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache, EvictParams from sglang.srt.model_executor.forward_batch_info import DSV4OutCacheLoc @@ -262,14 +267,80 @@ class DSV4NPUTokenToKVPoolAllocator(SWATokenToKVPoolAllocator): ) return result + def c128_num_pages_needed( + self, prefix_lens_cpu: torch.Tensor, seq_lens_cpu: torch.Tensor + ) -> int: + """Return the number of new C128 physical pages needed by this step.""" + if prefix_lens_cpu is None or seq_lens_cpu is None: + return 0 + ratio = 128 + page_size = self.c128_attn_allocator.page_size + prefix_pages = (prefix_lens_cpu // ratio + page_size - 1) // page_size + seq_pages = (seq_lens_cpu // ratio + page_size - 1) // page_size + return int((seq_pages - prefix_pages).clamp(min=0).sum().item()) + + def ensure_c128_capacity( + self, tree_cache: Optional[BasePrefixCache], num_pages: int + ) -> bool: + """Reclaim C128 pages through FULL leaf eviction. + + Evicting FULL also reclaims its C128 payload without interior holes.""" + if num_pages <= 0: + return True + + page_size = self.c128_attn_allocator.page_size + + def available_pages() -> int: + return self.c128_attn_allocator.available_size() // page_size + + if available_pages() >= num_pages: + return True + if tree_cache is None or tree_cache.is_chunk_cache(): + return False + + while available_pages() < num_pages: + result = tree_cache.evict(EvictParams(num_tokens=self.page_size)) + if result.num_tokens_evicted == 0: + break + return available_pages() >= num_pages + + def check_decode_capacity( + self, + *, + num_tokens: int, + tree_cache, + requests=None, + spec_algorithm=None, + ) -> bool: + """Check FULL/SWA and exact per-request C128 demand for one decode.""" + if requests is None: + return super().check_decode_capacity( + num_tokens=num_tokens, tree_cache=tree_cache + ) + + self.evict_to_free_tokens(tree_cache, num_tokens) + if spec_algorithm is not None and spec_algorithm.is_some(): + prefix_lens, seq_lens, _ = page_aligned_decode_alloc_lens( + requests, + reserve=get_alloc_reserve_per_decode(), + page_size=self.page_size, + ) + else: + prefix_lens = [req.kv.kv_allocated_len for req in requests] + seq_lens = [length + 1 for length in prefix_lens] + + prefix_lens_cpu = torch.tensor(prefix_lens, dtype=torch.int64) + seq_lens_cpu = torch.tensor(seq_lens, dtype=torch.int64) + c128_num_pages = self.c128_num_pages_needed(prefix_lens_cpu, seq_lens_cpu) + c128_ok = self.ensure_c128_capacity(tree_cache, c128_num_pages) + full_swa_ok = self.full_swa_available_size() >= num_tokens + return full_swa_ok and c128_ok + def _has_c128_sidecar_capacity( self, prefix_lens_cpu: torch.Tensor, seq_lens_cpu: torch.Tensor ) -> bool: - ratio = 128 + need = self.c128_num_pages_needed(prefix_lens_cpu, seq_lens_cpu) page_size = self.c128_attn_allocator.page_size - prefix_groups = (prefix_lens_cpu // ratio + page_size - 1) // page_size - seq_groups = (seq_lens_cpu // ratio + page_size - 1) // page_size - need = int((seq_groups - prefix_groups).clamp(min=0).sum().item()) return need <= self.c128_attn_allocator.available_size() // page_size def _alloc_compressed_kv( @@ -481,6 +552,9 @@ class DSV4NPUTokenToKVPoolAllocator(SWATokenToKVPoolAllocator): self.c128_attn_allocator.available_size() * 128, ) + def full_swa_available_size(self): + return super().available_size() + def resize(self, config) -> None: self.c128_attn_allocator.size = int(config.c128_max_total_num_tokens) self.c128_attn_allocator.num_pages = ( diff --git a/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_memory_pool.py b/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_memory_pool.py index 884516b49..b4491cf5a 100644 --- a/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_memory_pool.py +++ b/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_memory_pool.py @@ -76,6 +76,11 @@ class NPUDeepSeekV4SingleKVPool(DeepSeekV4SingleKVPool): kv_dim = self.qk_nope_head_dim + self.qk_rope_head_dim kv_dtype = torch.bfloat16 self.kv_cache_total_dim = kv_dim + # The HiCache assembler uses bytes_per_page_padded as host item_bytes. + # For example: kernel_page_size * kv_dim * sizeof(bf16). + self.bytes_per_page_padded = ( + self.kernel_page_size * kv_dim * torch.bfloat16.itemsize + ) # Writes are flat-indexed by loc; kernel_page_size controls the physical # page layout exposed to the NPU operators. npu_num_pages = (self.size + self.kernel_page_size + 1) // self.kernel_page_size @@ -464,6 +469,7 @@ class DSV4NPUTokenToKVPool(DeepSeekV4TokenToKVPool): # ------------------------------------------------------------------ def get_key_buffer(self, layer_id: int) -> torch.Tensor: + self.wait_layer_transfer(layer_id) item = self.layer_mapping[layer_id] ratio = item.compress_ratio if ratio == 0: @@ -490,6 +496,7 @@ class DSV4NPUTokenToKVPool(DeepSeekV4TokenToKVPool): flatten across (num_pages, page_size) and gather the matching tokens — shape becomes (num_tokens, 1, dim). """ + self.wait_layer_transfer(layer_id) # Index by RAW layer_id, not compress_layer_id (a per-bucket counter that # would collide across ratios). swa_kv_pool is sized layer_num=total_layers. kv = self.swa_kv_pool.kv_buffer[layer_id] @@ -510,6 +517,7 @@ class DSV4NPUTokenToKVPool(DeepSeekV4TokenToKVPool): from_indexer=True branch returns the dedicated quantized K buffer that ``torch.ops.custom.npu_quant_lightning_indexer`` consumes. """ + self.wait_layer_transfer(layer_id) item = self.layer_mapping[layer_id] if item.compress_ratio == 0: return None diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index f699df9b9..178399f72 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -99,10 +99,7 @@ from sglang.srt.managers.embed_types import PositionalEmbeds from sglang.srt.managers.scheduler_components.new_token_ratio_tracker import ( NewTokenRatioTracker, ) -from sglang.srt.mem_cache.allocation import ( - alloc_for_decode, - alloc_for_extend, -) +from sglang.srt.mem_cache.allocation import alloc_for_decode, alloc_for_extend from sglang.srt.mem_cache.allocation_sizing import get_alloc_reserve_per_decode from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import ( @@ -3097,8 +3094,16 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): shortfalls retract gracefully instead of tripping fail-loud alloc errors.""" num_tokens = self.new_tokens_required_next_decode(selected_indices) + requests = ( + self.reqs + if selected_indices is None + else [self.reqs[i] for i in selected_indices] + ) return self.token_to_kv_pool_allocator.check_decode_capacity( - num_tokens=num_tokens, tree_cache=self.tree_cache + num_tokens=num_tokens, + tree_cache=self.tree_cache, + requests=requests, + spec_algorithm=self.spec_algorithm, ) def retract_decode(self) -> Tuple[List[Req], float, List[Req]]: diff --git a/python/sglang/srt/mem_cache/allocation.py b/python/sglang/srt/mem_cache/allocation.py index 75d33454e..0032a8ccc 100644 --- a/python/sglang/srt/mem_cache/allocation.py +++ b/python/sglang/srt/mem_cache/allocation.py @@ -201,6 +201,8 @@ def alloc_paged_token_slots_extend( ) extra_alloc_kwargs["rotation_bases"] = kv_shard_rotation_bases if is_dsv4: + c128_num_pages = allocator.c128_num_pages_needed(prefix_lens_cpu, seq_lens_cpu) + allocator.ensure_c128_capacity(tree_cache, c128_num_pages) extra_alloc_kwargs["req_pool_indices"] = req_pool_indices # Per-call per-req table for the C128 KV last_loc lookup. if batch is not None: @@ -549,6 +551,10 @@ def alloc_paged_token_slots_decode( is_dsv4 = req_pool_indices is not None and hasattr(allocator, "c128_attn_allocator") extra_alloc_kwargs = {} if is_dsv4: + c128_num_pages = allocator.c128_num_pages_needed( + (seq_lens_cpu - 1).clamp(min=0), seq_lens_cpu + ) + allocator.ensure_c128_capacity(tree_cache, c128_num_pages) extra_alloc_kwargs["req_pool_indices"] = req_pool_indices # Per-call per-req C128 table for the last_loc lookup. if batch is not None: diff --git a/python/sglang/srt/mem_cache/allocator/base.py b/python/sglang/srt/mem_cache/allocator/base.py index fa640b650..68ecb5e18 100644 --- a/python/sglang/srt/mem_cache/allocator/base.py +++ b/python/sglang/srt/mem_cache/allocator/base.py @@ -76,10 +76,19 @@ class BaseTokenToKVPoolAllocator(abc.ABC): evict_from_tree_cache(tree_cache, num_tokens) - def check_decode_capacity(self, *, num_tokens: int, tree_cache) -> bool: + def check_decode_capacity( + self, + *, + num_tokens: int, + tree_cache, + requests=None, + spec_algorithm=None, + ) -> bool: """Whether the next decode step's ``num_tokens`` allocation fits after evicting reclaimable cache. The retract loop converges on this same - check, so a shortfall here retracts instead of failing in alloc.""" + check, so a shortfall here retracts instead of failing in alloc. + ``requests`` and ``spec_algorithm`` provide optional request-level context for allocators + whose demand cannot be represented by a single token count.""" self.evict_to_free_tokens(tree_cache, num_tokens) return self.available_size() >= num_tokens 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 3cb29e016..006b52a72 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,11 +27,7 @@ from sglang.srt.mem_cache.pool_host.mha import ( ) from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost from sglang.srt.mem_cache.unified_cache.component_type import ComponentType -from sglang.srt.runtime_context import ( - get_memory, - get_parallel, - get_serving, -) +from sglang.srt.runtime_context import get_memory, get_parallel, get_serving if TYPE_CHECKING: import torch @@ -441,7 +437,7 @@ def _deepseek_v4_num_host_pages( kvcache: Any, page_size: int, swa_page_size: int, -) -> tuple[int, int]: +) -> tuple[int, int, int]: allocator = params.token_to_kv_pool_allocator device_full_size = getattr(allocator, "size_full", kvcache.size) device_full_pages = (device_full_size + page_size - 1) // page_size @@ -456,7 +452,15 @@ def _deepseek_v4_num_host_pages( ratio = get_memory().hicache_ratio full_host_pages = int(device_full_pages * ratio) swa_host_pages = int(device_swa_pages * ratio) - return full_host_pages, swa_host_pages + + # NPU sizes the independent C128 host pool from its device page count. + # For example, host pages = device pages * hicache_ratio. + # GPU keeps the FULL page count because C128 is a KV-derived sidecar. + c128_host_pages = full_host_pages + c128_attn_allocator = getattr(allocator, "c128_attn_allocator", None) + if c128_attn_allocator is not None: + c128_host_pages = int(c128_attn_allocator.num_pages * ratio) + return full_host_pages, swa_host_pages, c128_host_pages def _dsv4_compressed_region_buffers(kvcache: Any, ratio: int) -> tuple[list, int]: @@ -477,6 +481,7 @@ class _IndexerRegion: name: PoolName device_buffers: list item_bytes: int + slot_page_size: int # FP4 page rows group their slots instead of laying tokens out flat, so the # fused-row token-granular copy does not apply and transfers must be whole # pages. The fused FP8 row has no such restriction. @@ -492,6 +497,31 @@ def _dsv4_indexer_regions(kvcache: Any, page_size: int) -> list[_IndexerRegion]: import torch pool = kvcache.c4_indexer_kv_pool + if getattr(pool, "has_npu_storage", False): + # NPU stores C4 int8 K and fp16 scale in separate PA_ND buffers. Keep + # the FULL logical page geometry for HiCache indices. The transfer path + # derives the native C4 page geometry from these device buffers. + k_buffers = pool.index_k_buffer + scale_buffers = pool.index_scale_buffer + return [ + _IndexerRegion( + name=PoolName.DEEPSEEK_V4_C4_INDEXER, + device_buffers=k_buffers, + item_bytes=int(k_buffers[0][0].numel() * k_buffers[0].element_size()), + slot_page_size=page_size, + page_aligned_only=False, + ), + _IndexerRegion( + name=PoolName.DEEPSEEK_V4_C4_INDEXER_SCALE, + device_buffers=scale_buffers, + item_bytes=int( + scale_buffers[0][0].numel() * scale_buffers[0].element_size() + ), + slot_page_size=page_size, + page_aligned_only=False, + ), + ] + fused = pool.index_k_with_scale_buffer if fused is not None: return [ @@ -499,6 +529,7 @@ def _dsv4_indexer_regions(kvcache: Any, page_size: int) -> list[_IndexerRegion]: name=PoolName.DEEPSEEK_V4_C4_INDEXER, device_buffers=fused, item_bytes=fused[0].shape[1] * fused[0].element_size(), + slot_page_size=page_size, page_aligned_only=False, ) ] @@ -525,12 +556,14 @@ def _dsv4_indexer_regions(kvcache: Any, page_size: int) -> list[_IndexerRegion]: name=PoolName.DEEPSEEK_V4_C4_INDEXER, device_buffers=payload, item_bytes=payload[0].shape[1], + slot_page_size=page_size, page_aligned_only=True, ), _IndexerRegion( name=PoolName.DEEPSEEK_V4_C4_INDEXER_SCALE, device_buffers=scale, item_bytes=scale[0].shape[1], + slot_page_size=page_size, page_aligned_only=True, ), ] @@ -544,6 +577,7 @@ def build_deepseek_v4_hicache_stack( storage_backend: Optional[str], host_swa_evict_fn: Optional[Callable[[int], Any]] = None, device_swa_evict_fn: Optional[Callable[[int], Any]] = None, + host_c128_evict_fn: Optional[Callable[[int], Any]] = None, prefetch_threshold: int = 256, model_name: Optional[str] = None, storage_backend_extra_config: Optional[dict] = None, @@ -587,7 +621,11 @@ def build_deepseek_v4_hicache_stack( 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( + ( + num_host_pages, + swa_num_host_pages, + c128_num_host_pages, + ) = _deepseek_v4_num_host_pages( params=params, kvcache=kvcache, page_size=page_size, @@ -639,6 +677,9 @@ def build_deepseek_v4_hicache_stack( if c4_layer_mapping: c4_device_buffers, c4_item_bytes = _dsv4_compressed_region_buffers(kvcache, 4) + # C4 is KV-derived, so its HiCache indices stay in the FULL logical + # coordinate space. NPU transfer code derives the native C4 page size + # (for example 32 rather than 128) from the device-buffer shape. c4_host_pool = DeepSeekV4PagedHostPool( pool_name=str(PoolName.DEEPSEEK_V4_C4), device_buffers=c4_device_buffers, @@ -666,7 +707,7 @@ def build_deepseek_v4_hicache_stack( device_buffers=region.device_buffers, item_bytes=region.item_bytes, num_host_pages=num_host_pages, - slot_page_size=page_size, + slot_page_size=region.slot_page_size, layout=get_memory().hicache_mem_layout, allocator_type=_get_allocator_type(), page_aligned_only=region.page_aligned_only, @@ -723,12 +764,20 @@ def build_deepseek_v4_hicache_stack( c128_device_buffers, c128_item_bytes = _dsv4_compressed_region_buffers( kvcache, 128 ) + # NPU C128 host views use kernel_page_size, for example 16 instead of 128. + # GPU pools lack this attribute and fall back to the global page_size. + _c128_pool = getattr(kvcache, "c128_kv_pool", None) + c128_slot_page_size = getattr(_c128_pool, "kernel_page_size", page_size) + # NPU derives the independent C128 host budget from its device pool. + c128_attn_allocator = getattr( + params.token_to_kv_pool_allocator, "c128_attn_allocator", None + ) c128_host_pool = DeepSeekV4PagedHostPool( pool_name=str(PoolName.DEEPSEEK_V4_C128), device_buffers=c128_device_buffers, item_bytes=c128_item_bytes, - num_host_pages=num_host_pages, - slot_page_size=page_size, + num_host_pages=c128_num_host_pages, + slot_page_size=c128_slot_page_size, layout=get_memory().hicache_mem_layout, allocator_type=_get_allocator_type(), ) @@ -742,6 +791,20 @@ def build_deepseek_v4_hicache_stack( device_pool=kvcache.c128_kv_pool, layer_mapping=c128_layer_mapping, transfer_layer_num=transfer_layer_num, + # NPU C128 uses bare allocator callbacks for independent indices. + # FULL leaf eviction also frees each attached C128 host value. + # GPU keeps its KV-derived sidecar callbacks unset. + host_evict_fn=host_c128_evict_fn, + device_alloc_fn=( + c128_attn_allocator.alloc + if c128_attn_allocator is not None + else None + ), + device_free_fn=( + c128_attn_allocator.free + if c128_attn_allocator is not None + else None + ), ), ] ) @@ -1088,10 +1151,7 @@ def build_full_draft_pools( tree_cache: Any, ) -> tuple[list[SidecarPoolSpec], list[PoolEntry]]: """Build draft KV/DSA sidecars whose indices follow target full KV.""" - from sglang.srt.mem_cache.memory_pool import ( - DSATokenToKVPool, - HybridLinearKVPool, - ) + from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool, HybridLinearKVPool pool = draft_kv_pool if isinstance(pool, HybridLinearKVPool): @@ -1235,6 +1295,7 @@ _COMPONENT_HOST_ATTR: dict[ComponentType, tuple[str, str]] = { ComponentType.FULL: ("full_kv_pool_host", "_full_kv_pool_host"), ComponentType.SWA: ("swa_kv_pool_host", "_swa_kv_pool_host"), ComponentType.MAMBA: ("mamba_pool_host", "_mamba_pool_host"), + ComponentType.C128: ("c128_kv_pool_host", "_c128_kv_pool_host"), } @@ -1284,16 +1345,21 @@ class StackStrategy: raise NotImplementedError +def _delegate_c128_host_evict(cache, n: int) -> int: + """Delegate C128 host pressure to FULL leaf eviction. + + For example, ``n`` C128 slots request ``n * 128`` FULL tokens.""" + return cache.evict_host(n * 128, ComponentType.FULL) + + class _DeepSeekV4Strategy(StackStrategy): def matches(self, kvcache, components): - from sglang.srt.mem_cache.deepseek_v4_memory_pool import ( - DeepSeekV4TokenToKVPool, - ) + from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool - return isinstance(kvcache, DeepSeekV4TokenToKVPool) and components == { - ComponentType.FULL, - ComponentType.SWA, - } + return isinstance(kvcache, DeepSeekV4TokenToKVPool) and components in ( + {ComponentType.FULL, ComponentType.SWA}, + {ComponentType.FULL, ComponentType.SWA, ComponentType.C128}, + ) def build_direct_linker_pool_group(self, *, kvcache, params, page_size): from sglang.srt.mem_cache.hybrid_cache.linker_pool_assembler import ( @@ -1326,12 +1392,31 @@ class _DeepSeekV4Strategy(StackStrategy): storage_backend=storage_backend, host_swa_evict_fn=lambda n: cache.evict_host(n, ComponentType.SWA), device_swa_evict_fn=lambda n: _evict_swa_for_device_alloc(cache, n), + # NPU delegates C128 host pressure to FULL host-leaf eviction. + # _delegate_c128_host_evict converts C128 slots to FULL tokens. + host_c128_evict_fn=( + (lambda n: _delegate_c128_host_evict(cache, n)) + if ComponentType.C128 in cache.components + else None + ), prefetch_threshold=prefetch_threshold, model_name=model_name, storage_backend_extra_config=storage_backend_extra_config, enable_storage_metrics=enable_storage_metrics, layer_mappings=layer_mappings, ) + # NPU drives C128 as an independent tree component, so adding a KV-derived + # sidecar would duplicate transfers. Add that sidecar only on GPU. + _sidecar_srcs = [ + (PoolName.DEEPSEEK_V4_C4, PoolName.KV), + (PoolName.DEEPSEEK_V4_C4_INDEXER, PoolName.KV), + (PoolName.DEEPSEEK_V4_C4_INDEXER_SCALE, PoolName.KV), + (PoolName.DEEPSEEK_V4_C4_STATE, PoolName.SWA), + (PoolName.DEEPSEEK_V4_C4_INDEXER_STATE, PoolName.SWA), + (PoolName.DEEPSEEK_V4_C128_STATE, PoolName.SWA), + ] + if ComponentType.C128 not in cache.components: + _sidecar_srcs.append((PoolName.DEEPSEEK_V4_C128, PoolName.KV)) sidecars = [ SidecarPoolSpec( pool_name=name, @@ -1342,15 +1427,7 @@ class _DeepSeekV4Strategy(StackStrategy): else PoolHitPolicy.ALL_PAGES ), ) - for name, src in ( - (PoolName.DEEPSEEK_V4_C4, PoolName.KV), - (PoolName.DEEPSEEK_V4_C4_INDEXER, PoolName.KV), - (PoolName.DEEPSEEK_V4_C4_INDEXER_SCALE, PoolName.KV), - (PoolName.DEEPSEEK_V4_C128, PoolName.KV), - (PoolName.DEEPSEEK_V4_C4_STATE, PoolName.SWA), - (PoolName.DEEPSEEK_V4_C4_INDEXER_STATE, PoolName.SWA), - (PoolName.DEEPSEEK_V4_C128_STATE, PoolName.SWA), - ) + for name, src in _sidecar_srcs if name in host_pool_group.entry_map ] component_host_pools = { @@ -1360,6 +1437,13 @@ class _DeepSeekV4Strategy(StackStrategy): component_host_pools[ComponentType.SWA] = host_pool_group.get_pool( PoolName.SWA ) + if ( + ComponentType.C128 in cache.components + and PoolName.DEEPSEEK_V4_C128 in host_pool_group.entry_map + ): + component_host_pools[ComponentType.C128] = host_pool_group.get_pool( + PoolName.DEEPSEEK_V4_C128 + ) return StackBuildResult( host_pool_group=host_pool_group, @@ -1435,9 +1519,7 @@ def _swa_layer_mappings(kvcache) -> tuple[dict[int, int], dict[int, int]]: class _SwaStrategy(StackStrategy): def matches(self, kvcache, components): - from sglang.srt.mem_cache.deepseek_v4_memory_pool import ( - DeepSeekV4TokenToKVPool, - ) + from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool return ( @@ -1491,9 +1573,7 @@ class _SwaStrategy(StackStrategy): class _MambaSwaStrategy(StackStrategy): def matches(self, kvcache, components): - from sglang.srt.mem_cache.deepseek_v4_memory_pool import ( - DeepSeekV4TokenToKVPool, - ) + from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool return ( @@ -1690,9 +1770,7 @@ class _MiniMaxSparseStrategy(StackStrategy): class _PlainKvStrategy(StackStrategy): def matches(self, kvcache, components): - from sglang.srt.mem_cache.deepseek_v4_memory_pool import ( - DeepSeekV4TokenToKVPool, - ) + from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool from sglang.srt.mem_cache.memory_pool import ( DSATokenToKVPool, HybridLinearKVPool, diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index 006d0dbde..dafbe1e2c 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -31,6 +31,9 @@ if _is_cuda or _is_hip: transfer_kv_per_layer_mla_pf_lf, ) +if _is_npu: + from sgl_kernel_npu.kvcacheio import TransferDirection, transfer_kv_dim_exchange + logger = logging.getLogger(__name__) @@ -220,7 +223,15 @@ class DeepSeekV4PagedHostPool(HiSparseHostPoolMixin, HostKVCache): f"{available_bytes / 1e9:.2f} GB free." ) - alloc_func = ALLOC_MEMORY_FUNCS[self.gpu_device] + # ALLOC_MEMORY_FUNCS is keyed by device *type* string ("npu"/"musa"/...), + # not torch.device objects; a torch.device key silently falls back to + # cudaHostRegister, which fails on NPU. Resolve the alloc func by type str. + _alloc_key = ( + self.gpu_device.type + if isinstance(self.gpu_device, torch.device) + else str(self.gpu_device) + ) + alloc_func = ALLOC_MEMORY_FUNCS[_alloc_key] self.data_refs = [] if self.layout == "layer_first": self.kv_buffer = [ @@ -300,6 +311,29 @@ class DeepSeekV4PagedHostPool(HiSparseHostPoolMixin, HostKVCache): device=self.gpu_device, ) + def _host_page_view(self, l: int) -> torch.Tensor: + """View a host layer as ``[pages, 1, P, 1, dim]``.""" + if self.layout == "layer_first": + layer_buffer = self.kv_buffer[l] + elif self.layout == "page_first": + layer_buffer = self.kv_buffer[:, l, :] + elif self.layout == "page_first_direct": + layer_buffer = self.kv_buffer[:, l, 0, :] + else: + raise ValueError( + f"{self.pool_name} _host_page_view: kernel_ascend requires " + "layer_first/page_first/page_first_direct layout, " + f"got {self.layout!r}" + ) + device_buffer = self.device_buffers[l] + return layer_buffer.view(device_buffer.dtype).view( + self.num_host_pages, + 1, + device_buffer.shape[1], + 1, + device_buffer.shape[-1], + ) + def get_contiguous_buf_infos(self): """Return per-layer page-row buffers for PD direct-to-host transfer.""" data_ptrs = [int(self.data_ptrs[i].item()) for i in range(self.layer_num)] @@ -310,6 +344,14 @@ class DeepSeekV4PagedHostPool(HiSparseHostPoolMixin, HostKVCache): def _to_page_indices(self, indices: torch.Tensor) -> torch.Tensor: return indices.reshape(-1, self.slot_page_size)[:, 0] // self.slot_page_size + def _to_native_page_indices( + self, indices: torch.Tensor, native_page_size: int + ) -> torch.Tensor: + """Expand logical page indices into a device buffer's native slots.""" + rows = self._to_page_indices(indices) + offsets = torch.arange(native_page_size, dtype=rows.dtype, device=rows.device) + return (rows[:, None] * native_page_size + offsets).reshape(-1) + def _unaligned_transfer_error( self, host_indices: torch.Tensor, device_indices: torch.Tensor ) -> ValueError: @@ -446,6 +488,33 @@ class DeepSeekV4PagedHostPool(HiSparseHostPoolMixin, HostKVCache): dst_indices=host_rows, page_size=1, ) + elif io_backend == "kernel_ascend": + # HiCache keeps KV-derived pools in the FULL logical coordinate + # space. Ascend consumes the native slots of the shaped device + # buffer, so translate the already resolved page rows only at this + # backend boundary. + native_page_size = int(self.device_buffers[0].shape[1]) + native_device_indices = self._to_native_page_indices( + device_indices, native_page_size + ) + native_host_indices = self._to_native_page_indices( + host_indices, native_page_size + ) + for l in range(self.layer_num): + dev_view = self.device_buffers[l].unsqueeze(0) + # dev_view: [1, dev_pages, native_page_size, 1, kv_dim] + host_view = self._host_page_view(l) + # host_view: [num_host_pages, 1, native_page_size, 1, kv_dim] + transfer_kv_dim_exchange( + device_k=dev_view, + host_k=host_view, + device_v=torch.empty(0, device=dev_view.device), + host_v=torch.empty(0, device="cpu"), + device_indices=native_device_indices, + host_indices=native_host_indices, + page_size=native_page_size, + direction=TransferDirection.D2H, + ) else: raise ValueError( f"Unsupported V4 paged host layout/backend: {self.layout}/{io_backend}" @@ -516,6 +585,29 @@ class DeepSeekV4PagedHostPool(HiSparseHostPoolMixin, HostKVCache): layer_id=layer_id, page_size=1, ) + elif io_backend == "kernel_ascend": + # NPU whole-page H2D via Ascend dim-exchange op, for layer_id only. + native_page_size = int(self.device_buffers[layer_id].shape[1]) + native_device_indices = self._to_native_page_indices( + device_indices, native_page_size + ) + native_host_indices = self._to_native_page_indices( + host_indices, native_page_size + ) + dev_view = self.device_buffers[layer_id].unsqueeze(0) + # dev_view: [1, dev_pages, native_page_size, 1, kv_dim] + host_view = self._host_page_view(layer_id) + # host_view: [num_host_pages, 1, native_page_size, 1, kv_dim] + transfer_kv_dim_exchange( + device_k=dev_view, + host_k=host_view, + device_v=torch.empty(0, device=dev_view.device), + host_v=torch.empty(0, device="cpu"), + device_indices=native_device_indices, + host_indices=native_host_indices, + page_size=native_page_size, + direction=TransferDirection.H2D, + ) else: raise ValueError( f"Unsupported V4 paged host layout/backend: {self.layout}/{io_backend}" @@ -638,7 +730,14 @@ class DeepSeekV4StateHostPool(HostKVCache): f"{available_bytes / 1e9:.2f} GB free." ) - alloc_func = ALLOC_MEMORY_FUNCS[self.gpu_device] + # ALLOC_MEMORY_FUNCS is keyed by device *type* string ("npu"/"musa"/...), + # not torch.device objects; resolve the key the same way PagedHostPool does. + _state_alloc_key = ( + self.gpu_device.type + if isinstance(self.gpu_device, torch.device) + else str(self.gpu_device) + ) + alloc_func = ALLOC_MEMORY_FUNCS[_state_alloc_key] self.data_refs = [] if self.layout == "layer_first": self.kv_buffer = [ @@ -756,6 +855,35 @@ class DeepSeekV4StateHostPool(HostKVCache): ) return indices.reshape(-1, self.swa_page_size)[:, 0] // self.swa_page_size + def _ring_op_indices(self, rows: torch.Tensor) -> torch.Tensor: + """Expand each SWA page row into ``ring_size`` operator indices. + + For example, row ``r`` maps to ``r * ring_size + arange(ring_size)``.""" + ar = torch.arange(self.ring_size, device=rows.device, dtype=rows.dtype) + return (rows.reshape(-1, 1) * self.ring_size + ar).reshape(-1) + + def _state_host_page_view(self, l: int) -> torch.Tensor: + """View host state layer ``l`` as ``[pages, 1, R, 1, last_dim]``. + + ``last_dim = state_page_bytes // R // state_dtype.itemsize``.""" + state_dtype = self.state_pools[l].kv_score_buffer.kv_score.dtype + last_dim = self.state_page_bytes // self.ring_size // state_dtype.itemsize + if self.layout == "layer_first": + layer_buffer = self.kv_buffer[l] + elif self.layout == "page_first": + layer_buffer = self.kv_buffer[:, l, :] + elif self.layout == "page_first_direct": + layer_buffer = self.kv_buffer[:, l, 0, :] + else: + raise ValueError( + f"{self.pool_name} _state_host_page_view: kernel_ascend requires " + "layer_first/page_first/page_first_direct layout, " + f"got {self.layout!r}" + ) + return layer_buffer.view(state_dtype).view( + self.num_host_pages, 1, self.ring_size, 1, last_dim + ) + def get_size_per_token(self): return self.state_page_bytes @@ -842,6 +970,36 @@ class DeepSeekV4StateHostPool(HostKVCache): dst_indices=host_rows, page_size=1, ) + elif io_backend == "kernel_ascend": + # Ascend copies ring_size state slots for each page-aligned SWA row. + # _ring_op_indices expands device and host rows into operator indices. + dev_op_indices = self._ring_op_indices(device_rows) + host_op_indices = self._ring_op_indices(host_rows) + for l in range(self.layer_num): + state_dtype = self.state_pools[l].kv_score_buffer.kv_score.dtype + last_dim = ( + self.state_page_bytes // self.ring_size // state_dtype.itemsize + ) + num_dev_pages = self.device_page_views[l].shape[0] + dev_view = ( + self.device_page_views[l] + .view(state_dtype) + .view(num_dev_pages, self.ring_size, 1, last_dim) + .unsqueeze(0) + ) + # dev_view: [1, num_dev_pages, R, 1, last_dim] + host_view = self._state_host_page_view(l) + # host_view: [num_host_pages, 1, R, 1, last_dim] + transfer_kv_dim_exchange( + device_k=dev_view, + host_k=host_view, + device_v=torch.empty(0, device=dev_view.device), + host_v=torch.empty(0, device="cpu"), + device_indices=dev_op_indices, + host_indices=host_op_indices, + page_size=self.ring_size, + direction=TransferDirection.D2H, + ) else: raise ValueError( f"Unsupported V4 state host layout/backend: {self.layout}/{io_backend}" @@ -896,6 +1054,35 @@ class DeepSeekV4StateHostPool(HostKVCache): layer_id=layer_id, page_size=1, ) + elif io_backend == "kernel_ascend": + # NPU whole-page H2D via Ascend dim-exchange op, for layer_id only. + # See backup_from_device_all_layer: indices must be ring-row indices + # (ring_size entries per SWA page), not raw SWA locs. + R = self.ring_size + dev_op_indices = self._ring_op_indices(device_rows) + host_op_indices = self._ring_op_indices(host_rows) + state_dtype = self.state_pools[layer_id].kv_score_buffer.kv_score.dtype + last_dim = self.state_page_bytes // R // state_dtype.itemsize + num_dev_pages = self.device_page_views[layer_id].shape[0] + dev_view = ( + self.device_page_views[layer_id] + .view(state_dtype) + .view(num_dev_pages, R, 1, last_dim) + .unsqueeze(0) + ) + # dev_view: [1, num_dev_pages, R, 1, last_dim] + host_view = self._state_host_page_view(layer_id) + # host_view: [num_host_pages, 1, R, 1, last_dim] + transfer_kv_dim_exchange( + device_k=dev_view, + host_k=host_view, + device_v=torch.empty(0, device=dev_view.device), + host_v=torch.empty(0, device="cpu"), + device_indices=dev_op_indices, + host_indices=host_op_indices, + page_size=R, + direction=TransferDirection.H2D, + ) else: raise ValueError( f"Unsupported V4 state host layout/backend: {self.layout}/{io_backend}" diff --git a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py index ec6c606a5..097532d9f 100644 --- a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py +++ b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py @@ -778,9 +778,7 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore): f"_{self.mha_suffix}_{PoolName.DRAFT}_v", ] elif pool_name == PoolName.DRAFT_SWA: - from sglang.srt.mem_cache.memory_pool_host import ( - DeepSeekV4PagedHostPool, - ) + from sglang.srt.mem_cache.memory_pool_host import DeepSeekV4PagedHostPool from sglang.srt.mem_cache.pool_host.mha import MHATokenToKVPoolHost if isinstance( diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 2f44b0568..9e906e8c7 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -37,11 +37,7 @@ from sglang.srt.mem_cache.buffer_mode.storage_existence_cache import ( StorageExistenceCache, ) from sglang.srt.mem_cache.common import RetractionBackup -from sglang.srt.mem_cache.hicache_storage import ( - PoolName, - PoolTransfer, - SidecarPoolSpec, -) +from sglang.srt.mem_cache.hicache_storage import PoolName, PoolTransfer, SidecarPoolSpec from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import ( HybridCacheController, ) @@ -89,11 +85,7 @@ from sglang.srt.observability.metrics_collector import ( StorageMetrics, StorageMetricsCollector, ) -from sglang.srt.runtime_context import ( - get_memory, - get_model, - get_observability, -) +from sglang.srt.runtime_context import get_memory, get_model, get_observability from sglang.srt.session.streaming_session import StreamingSession from sglang.srt.utils.common import ceil_align @@ -111,6 +103,21 @@ from sglang.srt.utils.rank_consensus_checker import rank_consensus T = TypeVar("T") + +def _c128_transfer_num_pages(transfers: Sequence[PoolTransfer], page_size: int) -> int: + num_pages = 0 + for transfer in transfers: + if transfer.host_indices is None: + continue + num_slots = len(transfer.host_indices) + assert num_slots % page_size == 0, ( + f"C128 load-back transfers must contain complete physical pages: " + f"{num_slots=}, {page_size=}" + ) + num_pages += num_slots // page_size + return num_pages + + COMPONENT_REGISTRY: dict[ComponentType, type[TreeComponent]] = { ComponentType.FULL: FullComponent, ComponentType.MAMBA: MambaComponent, @@ -1686,6 +1693,21 @@ class UnifiedRadixCache(BasePrefixCache): self.dec_host_lock_ref(node_id, host_anchor_params) return False + c128_allocator = getattr( + self.token_to_kv_pool_allocator, "c128_attn_allocator", None + ) + if c128_allocator is not None: + c128_num_pages = _c128_transfer_num_pages( + comp_xfers.get(ComponentType.C128, ()), + c128_allocator.page_size, + ) + if not self.token_to_kv_pool_allocator.ensure_c128_capacity( + self, c128_num_pages + ): + self.dec_lock_ref(node_id, ancestor_lock_params) + self.dec_host_lock_ref(node_id, host_anchor_params) + return False + avail = self._component_available_size(ComponentType.FULL) if avail < kv_tokens: needed = kv_tokens - avail diff --git a/test/registered/unit/mem_cache/test_dsv4_hicache_l2.py b/test/registered/unit/mem_cache/test_dsv4_hicache_l2.py new file mode 100644 index 000000000..9046a6941 --- /dev/null +++ b/test/registered/unit/mem_cache/test_dsv4_hicache_l2.py @@ -0,0 +1,322 @@ +"""Focused CPU contracts for DeepSeek V4 NPU HiCache L2.""" + +import unittest +from types import MethodType, SimpleNamespace +from unittest.mock import MagicMock, patch + +import torch + +from sglang.srt.hardware_backend.npu.dsv4.c128_sidecar_component import ( + C128SidecarComponent, +) +from sglang.srt.hardware_backend.npu.dsv4.dsv4_allocator import ( + DSV4NPUTokenToKVPoolAllocator, +) +from sglang.srt.mem_cache.hicache_storage import PoolName, PoolTransfer +from sglang.srt.mem_cache.hybrid_cache import hybrid_pool_assembler as assembler +from sglang.srt.mem_cache.memory_pool_host import ( + DeepSeekV4PagedHostPool, + DeepSeekV4StateHostPool, +) +from sglang.srt.mem_cache.unified_cache.components import ( + CacheTransferPhase, + 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 TestDSV4DecodeCapacity(CustomTestCase): + def test_decode_capacity_uses_exact_c128_page_demand(self): + allocator = SimpleNamespace( + page_size=128, + c128_attn_allocator=SimpleNamespace(page_size=16), + evict_to_free_tokens=MagicMock(), + ensure_c128_capacity=MagicMock(return_value=True), + full_swa_available_size=MagicMock(return_value=384), + ) + allocator.c128_num_pages_needed = MethodType( + DSV4NPUTokenToKVPoolAllocator.c128_num_pages_needed, allocator + ) + requests = [ + SimpleNamespace(kv=SimpleNamespace(kv_allocated_len=length)) + for length in (127, 2175, 2047) + ] + + result = DSV4NPUTokenToKVPoolAllocator.check_decode_capacity( + allocator, + num_tokens=384, + tree_cache=MagicMock(), + requests=requests, + spec_algorithm=SimpleNamespace(is_some=lambda: False), + ) + + self.assertTrue(result) + # 127 -> 128 and 2175 -> 2176 open new physical C128 pages; + # 2047 -> 2048 only fills the current page. + allocator.ensure_c128_capacity.assert_called_once() + self.assertEqual(allocator.ensure_c128_capacity.call_args.args[1], 2) + allocator.evict_to_free_tokens.assert_called_once() + + def test_c128_capacity_evicts_full_leaves_until_enough_pages_exist(self): + available_pages = 1 + c128_page_size = 16 + + def available_size(): + return available_pages * c128_page_size + + def evict(_params): + nonlocal available_pages + available_pages += 1 + return SimpleNamespace(num_tokens_evicted=128) + + allocator = SimpleNamespace( + page_size=128, + c128_attn_allocator=SimpleNamespace( + page_size=c128_page_size, + available_size=available_size, + ), + ) + tree_cache = SimpleNamespace( + is_chunk_cache=lambda: False, + evict=MagicMock(side_effect=evict), + ) + + result = DSV4NPUTokenToKVPoolAllocator.ensure_c128_capacity( + allocator, tree_cache, num_pages=3 + ) + + self.assertTrue(result) + self.assertEqual(tree_cache.evict.call_count, 2) + self.assertEqual( + [call.args[0].num_tokens for call in tree_cache.evict.call_args_list], + [128, 128], + ) + + +class TestDSV4HostViews(CustomTestCase): + @staticmethod + def _layout_buffer(raw, layout): + if layout == "layer_first": + return [raw[:, layer, :].clone() for layer in range(raw.shape[1])] + if layout == "page_first": + return raw.clone() + return raw.unsqueeze(2).clone() + + @staticmethod + def _selected_layer(buffer, layout, layer): + if layout == "layer_first": + return buffer[layer] + if layout == "page_first": + return buffer[:, layer, :] + return buffer[:, layer, 0, :] + + def test_paged_and_state_host_views_support_every_layout(self): + pages, layers, page_slots, last_dim = 2, 2, 2, 3 + item_bytes = page_slots * last_dim + raw = torch.arange(pages * layers * item_bytes, dtype=torch.uint8).reshape( + pages, layers, item_bytes + ) + + state_dtype = torch.float16 + ring_size = 2 + state_last_dim = 3 + state_page_bytes = ring_size * state_last_dim * state_dtype.itemsize + state_raw = torch.arange( + pages * layers * state_page_bytes, dtype=torch.uint8 + ).reshape(pages, layers, state_page_bytes) + + for layout in ("layer_first", "page_first", "page_first_direct"): + with self.subTest(pool="paged", layout=layout): + host = DeepSeekV4PagedHostPool.__new__(DeepSeekV4PagedHostPool) + host.pool_name = "c4" + host.layout = layout + host.num_host_pages = pages + host.slot_page_size = page_slots + host.device_buffers = [ + torch.empty((4, page_slots, 1, last_dim), dtype=torch.int8) + for _ in range(layers) + ] + host.kv_buffer = self._layout_buffer(raw, layout) + + view = host._host_page_view(1) + expected = self._selected_layer(host.kv_buffer, layout, 1).view( + torch.int8 + ) + self.assertEqual(view.shape, (pages, 1, page_slots, 1, last_dim)) + self.assertTrue( + torch.equal(view.reshape(pages, -1), expected.reshape(pages, -1)) + ) + + with self.subTest(pool="state", layout=layout): + host = DeepSeekV4StateHostPool.__new__(DeepSeekV4StateHostPool) + host.pool_name = "c4_state" + host.layout = layout + host.num_host_pages = pages + host.ring_size = ring_size + host.state_page_bytes = state_page_bytes + host.state_pools = [ + SimpleNamespace( + kv_score_buffer=SimpleNamespace( + kv_score=torch.empty(1, dtype=state_dtype) + ) + ) + for _ in range(layers) + ] + host.kv_buffer = self._layout_buffer(state_raw, layout) + + view = host._state_host_page_view(1) + expected = self._selected_layer(host.kv_buffer, layout, 1).view( + state_dtype + ) + self.assertEqual(view.shape, (pages, 1, ring_size, 1, state_last_dim)) + self.assertTrue( + torch.equal(view.reshape(pages, -1), expected.reshape(pages, -1)) + ) + + +class TestC128L2Ownership(CustomTestCase): + def test_load_back_retains_and_attaches_each_c128_page(self): + allocator = SimpleNamespace( + c128_attn_allocator=SimpleNamespace(page_size=4), + retain_c128_pages=MagicMock(), + ) + nodes = { + node_id: SimpleNamespace( + component_data={ + ComponentType.C128: SimpleNamespace( + host_value=torch.arange(4), value=None + ) + } + ) + for node_id in (1, 2) + } + tree_core = SimpleNamespace( + node_by_id=lambda node_id: nodes[node_id], + set_component_device_value=MagicMock(), + ) + component = C128SidecarComponent.__new__(C128SidecarComponent) + component.cache = SimpleNamespace(token_to_kv_pool_allocator=allocator) + component.tree_core = tree_core + transfer = PoolTransfer( + name=PoolName.DEEPSEEK_V4_C128, + device_indices=torch.arange(28, 36), + nodes_to_load=[1, 2], + ) + + component.commit_hicache_transfer( + node=None, + phase=CacheTransferPhase.LOAD_BACK, + transfers=[transfer], + cache_actions=[], + ) + + retained = [ + call.args[0].tolist() for call in allocator.retain_c128_pages.call_args_list + ] + attached = [ + (call.args[0], call.args[2].tolist()) + for call in tree_core.set_component_device_value.call_args_list + ] + self.assertEqual(retained, [[7], [8]]) + self.assertEqual(attached, [(1, [7]), (2, [8])]) + + +class TestDSV4PoolAssembly(CustomTestCase): + def test_npu_indexer_registers_k_and_scale_as_separate_host_pools(self): + k_buffer = [torch.empty((2, 32, 1, 4), dtype=torch.int8)] + scale_buffer = [torch.empty((2, 32, 1, 1), dtype=torch.float16)] + c4_buffer = [torch.empty((2, 32, 1, 8), dtype=torch.int8)] + indexer_pool = SimpleNamespace( + has_npu_storage=True, + index_k_buffer=k_buffer, + index_scale_buffer=scale_buffer, + ) + kvcache = SimpleNamespace( + _unified_kv=True, + size=256, + swa_size=0, + swa_page_size=256, + c4_kv_pool=SimpleNamespace(kernel_page_size=32), + c4_indexer_kv_pool=indexer_pool, + unified_region_buffers=lambda ratio: (c4_buffer, 256), + ) + params = SimpleNamespace( + page_size=128, + token_to_kv_pool_allocator=SimpleNamespace(size_full=256), + mtp_draft_device_pools=[], + tp_cache_group=None, + attn_cp_cache_group=None, + attn_tp_cache_group=None, + pp_cache_group=None, + ) + mappings = assembler._DeepSeekV4LayerMappings( + transfer_layer_num=1, + full={0: 0}, + swa={}, + c4={0: 0}, + c128={}, + c4_state={}, + c4_state_global_layers=[], + ) + memory = SimpleNamespace( + hicache_mem_layout="page_first_direct", + hicache_write_policy="write_through", + hicache_io_backend="kernel_ascend", + hicache_host_memory_mode="cache", + ) + logical_pool = SimpleNamespace( + layout=memory.hicache_mem_layout, + page_size=128, + device="cpu", + size=256, + logical_size=256, + can_use_write_back_jit=False, + ) + + def make_host_pool(**kwargs): + return SimpleNamespace( + **kwargs, + page_size=kwargs["slot_page_size"], + device="cpu", + size=kwargs["num_host_pages"] * kwargs["slot_page_size"], + logical_size=kwargs["num_host_pages"] * kwargs["slot_page_size"], + can_use_write_back_jit=False, + ) + + controller = object() + with ( + patch.object(assembler, "get_memory", return_value=memory), + patch.object( + assembler, "_deepseek_v4_num_host_pages", return_value=(2, 0, 0) + ), + patch.object(assembler, "_get_allocator_type", return_value="default"), + patch.object(assembler, "LogicalHostPool", return_value=logical_pool), + patch.object( + assembler, + "DeepSeekV4PagedHostPool", + side_effect=make_host_pool, + ), + patch.object(assembler, "HybridCacheController", return_value=controller), + ): + group, built_controller = assembler.build_deepseek_v4_hicache_stack( + params=params, + kvcache=kvcache, + load_cache_event=None, + storage_backend=None, + layer_mappings=mappings, + ) + + k_host = group.entry_map[PoolName.DEEPSEEK_V4_C4_INDEXER].host_pool + scale_host = group.entry_map[PoolName.DEEPSEEK_V4_C4_INDEXER_SCALE].host_pool + self.assertIs(k_host.device_buffers, k_buffer) + self.assertIs(scale_host.device_buffers, scale_buffer) + self.assertEqual((k_host.item_bytes, scale_host.item_bytes), (128, 64)) + self.assertEqual((k_host.slot_page_size, scale_host.slot_page_size), (128, 128)) + self.assertIs(built_controller, controller) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/mem_cache/test_hicache_host_register.py b/test/registered/unit/mem_cache/test_hicache_host_register.py index 2de076e8d..e12f30cd1 100644 --- a/test/registered/unit/mem_cache/test_hicache_host_register.py +++ b/test/registered/unit/mem_cache/test_hicache_host_register.py @@ -196,7 +196,7 @@ class TestHiCacheHostRegister(unittest.TestCase): "host_memory_budget_bytes", return_value=1024**3, ), - mock.patch.dict(ALLOC_MEMORY_FUNCS, {torch.device("cpu"): alloc}), + mock.patch.dict(ALLOC_MEMORY_FUNCS, {"cpu": alloc}), ): DeepSeekV4PagedHostPool( pool_name="test", @@ -229,7 +229,7 @@ class TestHiCacheHostRegister(unittest.TestCase): "host_memory_budget_bytes", return_value=1024**3, ), - mock.patch.dict(ALLOC_MEMORY_FUNCS, {torch.device("cpu"): alloc}), + mock.patch.dict(ALLOC_MEMORY_FUNCS, {"cpu": alloc}), ): DeepSeekV4StateHostPool( pool_name="test",