[NPU] Support DSV4 host memory cache management (#37382)

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

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 = (

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

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]]:

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:

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
 

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,

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}"

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(

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