Overlap CP HiCache backup without exposing partial host state

CP shared KV with HiCache and EAGLE needs host backup to overlap forward while keeping radix visibility synchronous. The change reserves host slots before forward, drives target and draft backup from explicit layer-end hooks, and commits host visibility only after the final target/draft ack. It also probes the final insertion prefix before early reservation so repeated EAGLE prompts do not prepare duplicate suffix backups that later rollback as insert_miss.

Constraint: CP ranks use independent shared-KV pools, so target/draft host state must remain atomically visible at the radix boundary.

Constraint: Fused MLA and NSA store paths can bypass store-side notifier hooks, so layer end is the safer backup progress boundary.

Rejected: Store-side backup notifier as the primary trigger | fused store and zero-local paths made notifier coverage fragile.

Rejected: Reserve from cache_protected_len alone | EAGLE bigram/page alignment can make final insertion find a longer existing prefix and force duplicate rollback work.

Confidence: medium

Scope-risk: moderate

Directive: Do not add per-layer CP collectives here; keep radix state synchronous and data transfer asynchronous/local-event driven.

Tested: local git diff --check

Tested: local py_compile for touched CP HiCache/cache-controller/deepseek/test files

Tested: remote pytest test/registered/unit/mem_cache/test_cp_hicache_metadata.py test/registered/unit/managers/test_hicache_controller_cp.py -q (115 passed, 5 warnings)

Not-tested: full GLM5 ETE server rerun after this commit

python/sglang/srt/managers/cache_controller.py

@@ -361,7 +361,7 @@ class HiCacheController:
         self.mem_pool_device.register_layer_transfer_counter(self.layer_done_counter)
         if hasattr(self.mem_pool_device, "register_layer_backup_notifier"):
             self.mem_pool_device.register_layer_backup_notifier(
-                lambda layer_id: self.on_layer_kv_stored(layer_id, source="target")
+                lambda layer_id: self.on_layer_end(layer_id, source="target")
             )
         self.draft_mem_pool_host = None
         self.draft_mem_pool_device = None
@@ -431,10 +431,10 @@ class HiCacheController:
             )
         if hasattr(draft_mem_pool_device, "register_layer_backup_notifier"):
             draft_mem_pool_device.register_layer_backup_notifier(
-                lambda layer_id: self.on_layer_kv_stored(layer_id, source="draft")
+                lambda layer_id: self.on_layer_end(layer_id, source="draft")
             )
 
-    def on_layer_kv_stored(self, layer_id: int, source: str = "target") -> None:
+    def on_layer_end(self, layer_id: int, source: str = "target") -> None:
         if not self.pending_layer_writes:
             return
         if source not in ("target", "draft"):
@@ -1236,7 +1236,7 @@ class HiCacheController:
         Current radix insertion calls write_backup after the request KV has already
         been materialized.  For that path, catch up by submitting all layers
         immediately through the per-layer API.  Future early reservations can use
-        catch_up_all_layers=False and rely on on_layer_kv_stored().
+        catch_up_all_layers=False and rely on explicit model layer-end hooks.
         """
 
         state = self._get_or_create_layer_write_state(reservation)

python/sglang/srt/managers/schedule_batch.py

@@ -644,6 +644,7 @@ class Req(ReqDllmMixin):
         self.last_host_backup_node: Any = None
         self.host_hit_length = 0
         self.prefix_match_deferred_by_pending_backup = False
+        self.cp_hicache_prepared_backup = None
         # Tokens loaded from storage backend (L3) during prefetch for this request
         self.storage_hit_length = 0
         # The node to lock until for swa radix tree lock ref

python/sglang/srt/managers/scheduler.py

@@ -2457,6 +2457,11 @@ class Scheduler(
             )
 
         new_batch.prepare_for_extend()
+        if self.enable_hierarchical_cache and hasattr(
+            self.tree_cache, "prepare_write_backup_for_req"
+        ):
+            for req in can_run_list:
+                self.tree_cache.prepare_write_backup_for_req(req)
 
         # Record prefill stats for logging after forward
         new_batch.prefill_stats = PrefillStats.from_adder(

python/sglang/srt/mem_cache/base_prefix_cache.py

@@ -60,6 +60,7 @@ class InsertParams:
     # General
     chunked: bool = False
     priority: int = 0
+    cp_hicache_prepared_backup: Optional[object] = None
 
 
 @dataclasses.dataclass

python/sglang/srt/mem_cache/hiradix_cache.py

@@ -17,6 +17,7 @@ from sglang.srt.environ import envs
 from sglang.srt.managers.cache_controller import (
     HiCacheController,
     HiCacheWriteFailure,
+    HiCacheWriteReservation,
     PrefetchOperation,
 )
 from sglang.srt.mem_cache.allocator import CPSharedPagedTokenToKVPoolAllocator
@@ -44,6 +45,7 @@ from sglang.srt.mem_cache.radix_cache import (
     RadixKey,
     TreeNode,
     compute_node_hash_values,
+    page_align_keys,
     split_node_hash_value,
 )
 from sglang.srt.mem_cache.utils import convert_to_bigram_key
@@ -246,6 +248,15 @@ class PendingHiCacheBackup:
     locked: bool = True
 
 
+@dataclass
+class PreparedCpHiCacheBackup:
+    node_id: int
+    reservation: HiCacheWriteReservation
+    metadata: CpHiCacheNodeMetadata
+    logical_len: int
+    attached: bool = False
+
+
 class HiCachePendingBackupSplit(Exception):
     def __init__(self, node: TreeNode):
         self.node = node
@@ -1010,6 +1021,260 @@ class HiRadixCache(RadixCache):
             self.dec_node_lock_ref(node)
         return node
 
+    def _sync_cp_write_required_host_slots(self, result) -> int:
+        """Return the max host slots any CP rank needs before write backup.
+
+        CP ranks must make the reserve/evict/retry decision collectively.  A
+        local successful reservation is not enough: another CP rank may be full
+        and enter host eviction, whose implementation contains collective
+        all_reduce calls.  If successful ranks skip that branch they can instead
+        enter unrelated collectives (for example disagg polling), causing a
+        process-group mismatch.
+        """
+
+        required_host_slots = (
+            int(result.required_host_slots)
+            if isinstance(result, HiCacheWriteFailure)
+            else 0
+        )
+        if (
+            getattr(self, "tp_group", None) is None
+            or getattr(self, "tp_world_size", 1) <= 1
+        ):
+            return required_host_slots
+
+        required = torch.tensor(required_host_slots, dtype=torch.int64, device="cpu")
+        torch.distributed.all_reduce(
+            required, op=torch.distributed.ReduceOp.MAX, group=self.tp_group
+        )
+        return int(required.item())
+
+    def _release_cp_write_reservation_for_retry(self, result, node_id: int) -> None:
+        if isinstance(result, HiCacheWriteFailure):
+            return
+        logger.info(
+            "[HiCache-write] release local CP reservation before collective retry: node_id=%d owned_positions=%d",
+            node_id,
+            result.metadata.owned_positions.numel(),
+        )
+        self.cache_controller.evict_cp_host(result.metadata)
+
+    def _reserve_write_cp_indices_collectively(
+        self, device_indices: torch.Tensor, node_id: int
+    ):
+        result = self.cache_controller.reserve_write_cp(
+            device_indices=device_indices,
+            node_id=node_id,
+        )
+        required_host_slots = self._sync_cp_write_required_host_slots(result)
+        if required_host_slots <= 0:
+            return result
+
+        self._release_cp_write_reservation_for_retry(result, node_id)
+        self._evict_host_for_physical_slots(
+            required_host_slots,
+            synchronize_across_ranks=getattr(self, "tp_world_size", 1) > 1,
+        )
+        logger.info(
+            "[HiCache-write] write_backup CP retry after host eviction: node_id=%d needed_slots=%d",
+            node_id,
+            required_host_slots,
+        )
+        result = self.cache_controller.reserve_write_cp(
+            device_indices=device_indices,
+            node_id=node_id,
+        )
+        required_host_slots = self._sync_cp_write_required_host_slots(result)
+        if required_host_slots <= 0:
+            return result
+
+        self._release_cp_write_reservation_for_retry(result, node_id)
+        logger.info(
+            "[HiCache-write] write_backup CP FAILED after collective retry: node_id=%d len=%d needed_slots=%d",
+            node_id,
+            len(device_indices) if device_indices is not None else 0,
+            required_host_slots,
+        )
+        return HiCacheWriteFailure(required_host_slots=required_host_slots)
+
+    def _reserve_write_cp_collectively(self, node: TreeNode):
+        return self._reserve_write_cp_indices_collectively(node.value, node.id)
+
+    def _attach_prepared_cp_backup(
+        self, node: TreeNode, prepared: PreparedCpHiCacheBackup
+    ) -> None:
+        if prepared.attached:
+            raise RuntimeError(
+                f"CP HiCache prepared backup node_id={prepared.node_id} was attached twice"
+            )
+        if len(node.value) != prepared.logical_len:
+            raise RuntimeError(
+                f"Prepared CP HiCache backup length mismatch for node_id={prepared.node_id}: "
+                f"node_len={len(node.value)} prepared_len={prepared.logical_len}"
+            )
+        node.id = prepared.node_id
+        self.ongoing_write_through[node.id] = node
+        self.inc_node_lock_ref(node)
+        self.pending_host_backups[node.id] = PendingHiCacheBackup(
+            node=node,
+            metadata=prepared.metadata,
+            logical_len=prepared.logical_len,
+            submitted=True,
+            locked=True,
+        )
+        prepared.attached = True
+        logger.info(
+            "[HiCache-write] attached prepared CP backup: node_id=%d logical_len=%d owned_positions=%d pending_backups=%d",
+            node.id,
+            prepared.logical_len,
+            prepared.metadata.owned_positions.numel(),
+            len(self.pending_host_backups),
+        )
+
+    def _rollback_prepared_cp_backup(
+        self, prepared: Optional[PreparedCpHiCacheBackup], reason: str
+    ) -> None:
+        if prepared is None or prepared.attached:
+            return
+        if prepared.node_id in self.cache_controller.pending_layer_writes:
+            raise RuntimeError(
+                f"Cannot rollback prepared CP HiCache backup node_id={prepared.node_id} "
+                f"while per-layer writes are still pending; reason={reason}"
+            )
+
+        retained_acks = []
+        removed_ack = False
+        for ack in self.cache_controller.ack_write_queue:
+            if prepared.node_id not in ack.node_ids:
+                retained_acks.append(ack)
+                continue
+            ack.finish_event.synchronize()
+            remaining_node_ids = [
+                node_id for node_id in ack.node_ids if node_id != prepared.node_id
+            ]
+            if remaining_node_ids:
+                retained_acks.append(ack._replace(node_ids=remaining_node_ids))
+            removed_ack = True
+        if removed_ack:
+            self.cache_controller.ack_write_queue = retained_acks
+        logger.info(
+            "[HiCache-write] rollback unattached prepared CP backup: node_id=%d logical_len=%d reason=%s",
+            prepared.node_id,
+            prepared.logical_len,
+            reason,
+        )
+        self.cache_controller.evict_cp_host(prepared.metadata)
+
+    def _probe_existing_radix_prefix_len_no_split(self, key: RadixKey) -> int:
+        """Return the already-present radix prefix length without mutating the tree.
+
+        `prepare_write_backup_for_req` runs before forward and must reserve host
+        slots only for KV that will become a *new* radix node at insertion time.
+        The request's `cache_protected_len` is based on scheduler prefix-match
+        semantics, which can be shorter than the prefix that a final insertion
+        will find (for example EAGLE/bigram plus page alignment).  Probing the
+        full insertion key here avoids backing up duplicate pages that insertion
+        will later reject and roll back.
+        """
+
+        if self.disable or len(key) == 0:
+            return 0
+
+        key, _ = self.maybe_bigram_convert(key)
+        if self.page_size != 1:
+            page_aligned_len = len(key) // self.page_size * self.page_size
+            key = key[:page_aligned_len]
+
+        if len(key) == 0:
+            return 0
+
+        node = getattr(self, "root_node", None)
+        if node is None:
+            return 0
+        child_key = self.get_child_key_fn(key)
+        total_prefix_len = 0
+
+        while len(key) > 0 and child_key in node.children:
+            child = node.children[child_key]
+            prefix_len = self.key_match_fn(child.key, key)
+            if prefix_len <= 0:
+                break
+
+            total_prefix_len += prefix_len
+            if prefix_len < len(child.key):
+                break
+
+            node = child
+            key = key[prefix_len:]
+            if len(key):
+                child_key = self.get_child_key_fn(key)
+
+        return total_prefix_len
+
+    def prepare_write_backup_for_req(self, req) -> None:
+        if (
+            self.disable
+            or not self._uses_cp_hicache
+            or self.cache_controller.write_policy == "write_back"
+            or getattr(req, "is_chunked", 0) > 0
+            or getattr(req, "cp_hicache_prepared_backup", None) is not None
+        ):
+            return
+
+        token_ids = req.fill_ids
+        keys = convert_to_bigram_key(token_ids) if self.is_eagle else token_ids
+        keys = page_align_keys(keys, self.page_size)
+        existing_prefix_len = self._probe_existing_radix_prefix_len_no_split(
+            RadixKey(
+                keys,
+                getattr(req, "extra_key", None),
+                is_bigram=self.is_eagle,
+            )
+        )
+        start = min(max(req.cache_protected_len, existing_prefix_len), len(keys))
+        if len(keys) <= start:
+            return
+
+        kv_indices = self.req_to_token_pool.req_to_token[
+            req.req_pool_idx, start : len(keys)
+        ].to(dtype=torch.int64, copy=True)
+        if len(kv_indices) == 0:
+            return
+
+        node_id = TreeNode.counter
+        TreeNode.counter += 1
+        result = self._reserve_write_cp_indices_collectively(kv_indices, node_id)
+        if isinstance(result, HiCacheWriteFailure):
+            logger.info(
+                "[HiCache-write] prepare CP backup skipped after host reservation failure: node_id=%d rid=%s len=%d",
+                node_id,
+                getattr(req, "rid", "<unknown>"),
+                len(kv_indices),
+            )
+            return
+
+        try:
+            self.cache_controller.submit_write_cp_per_layer(
+                result, catch_up_all_layers=False
+            )
+        except Exception:
+            self.cache_controller.evict_cp_host(result.metadata)
+            raise
+
+        req.cp_hicache_prepared_backup = PreparedCpHiCacheBackup(
+            node_id=node_id,
+            reservation=result,
+            metadata=result.metadata,
+            logical_len=len(kv_indices),
+        )
+        logger.info(
+            "[HiCache-write] prepared CP per-layer backup before forward: node_id=%d rid=%s logical_len=%d owned_positions=%d",
+            node_id,
+            getattr(req, "rid", "<unknown>"),
+            len(kv_indices),
+            result.metadata.owned_positions.numel(),
+        )
+
     def _node_host_len(self, node: TreeNode) -> int:
         if self._uses_cp_hicache:
             return node.host_len
@@ -1029,25 +1294,7 @@ class HiRadixCache(RadixCache):
             write_back,
         )
         if self._uses_cp_hicache:
-            result = self.cache_controller.reserve_write_cp(
-                device_indices=node.value,
-                node_id=node.id,
-            )
-            if isinstance(result, HiCacheWriteFailure):
-                required_host_slots = result.required_host_slots
-                self._evict_host_for_physical_slots(
-                    required_host_slots,
-                    synchronize_across_ranks=getattr(self, "tp_world_size", 1) > 1,
-                )
-                logger.info(
-                    "[HiCache-write] write_backup CP retry after host eviction: node_id=%d needed_slots=%d",
-                    node.id,
-                    required_host_slots,
-                )
-                result = self.cache_controller.reserve_write_cp(
-                    device_indices=node.value,
-                    node_id=node.id,
-                )
+            result = self._reserve_write_cp_collectively(node)
             if isinstance(result, HiCacheWriteFailure):
                 logger.info(
                     "[HiCache-write] write_backup CP FAILED (host full): node_id=%d len=%d",
@@ -1191,6 +1438,7 @@ class HiRadixCache(RadixCache):
             if not finish_event.query():
                 break
             finish_count += 1
+        local_finish_count = finish_count
         queue_size = torch.tensor(finish_count, dtype=torch.int, device="cpu")
         if self.tp_world_size > 1:
             # synchronize TP workers to make the same update to radix cache
@@ -1201,11 +1449,11 @@ class HiRadixCache(RadixCache):
             )
 
         finish_count = int(queue_size.item())
-        logger.info(
+        logger.debug(
             "[HiCache-write] writing_check: ongoing=%d ack_queue=%d local_finished=%d sync_finished=%d tp_size=%d",
             len(self.ongoing_write_through),
             ack_queue_len,
-            finish_count if self.tp_world_size <= 1 else queue_size.item(),
+            local_finish_count,
             finish_count,
             getattr(self, "tp_world_size", 1),
         )
@@ -2264,6 +2512,7 @@ class HiRadixCache(RadixCache):
         value = params.value
         chunked = params.chunked
         priority = params.priority
+        prepared_cp_backup = params.cp_hicache_prepared_backup
 
         if priority is None:
             priority = 0
@@ -2325,7 +2574,18 @@ class HiRadixCache(RadixCache):
                 child_key = self.get_child_key_fn(key)
 
         if len(key):
-            new_node = TreeNode(priority=priority)
+            use_prepared_cp_backup = (
+                prepared_cp_backup is not None
+                and getattr(prepared_cp_backup, "logical_len", -1) == len(value)
+            )
+            new_node = TreeNode(
+                id=(
+                    prepared_cp_backup.node_id
+                    if use_prepared_cp_backup
+                    else None
+                ),
+                priority=priority,
+            )
             new_node.parent = node
             new_node.key = key
             new_node.value = value.clone()
@@ -2342,7 +2602,10 @@ class HiRadixCache(RadixCache):
             self._record_store_event(new_node)
 
             if self.cache_controller.write_policy != "write_back":
-                self._inc_hit_count(new_node, chunked)
+                if use_prepared_cp_backup:
+                    self._attach_prepared_cp_backup(new_node, prepared_cp_backup)
+                else:
+                    self._inc_hit_count(new_node, chunked)
         return InsertResult(prefix_len=total_prefix_length)
 
     def release_aborted_request(self, rid: str):

python/sglang/srt/mem_cache/memory_pool.py

@@ -695,7 +695,6 @@ class KVCache(abc.ABC):
         # default state for optional layer-wise transfer control
         self.layer_transfer_counter = None
         self.layer_backup_notifiers = []
-        self.layer_backup_notify_after_indexer = False
 
         # for disagg with nvlink
         self.enable_custom_mem_pool, self.custom_mem_pool, _ = (
@@ -750,11 +749,10 @@ class KVCache(abc.ABC):
     def register_layer_backup_notifier(self, notifier):
         self.layer_backup_notifiers.append(notifier)
 
-    def _notify_layer_kv_stored(self, layer_id: int, source: str = "kv"):
-        if self.layer_backup_notify_after_indexer and source != "indexer":
-            return
+    def notify_layer_end_for_backup(self, layer_id: int):
+        local_layer_id = layer_id - self.start_layer
         for notifier in self.layer_backup_notifiers:
-            notifier(layer_id)
+            notifier(local_layer_id)
 
     def get_cpu_copy(self, indices):
         raise NotImplementedError()
@@ -1047,7 +1045,6 @@ class MHATokenToKVPool(KVCache):
             alt_stream=self.alt_stream,
             same_kv_dim=self.same_kv_dim,
         )
-        self._notify_layer_kv_stored(layer_id - self.start_layer)
 
     def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor):
         if envs.SGLANG_NATIVE_MOVE_KV_CACHE.get():
@@ -1595,7 +1592,6 @@ class MLATokenToKVPool(KVCache):
             )
         else:
             self.kv_buffer[layer_id - self.start_layer][loc] = cache_k
-        self._notify_layer_kv_stored(layer_id - self.start_layer)
 
     def set_mla_kv_buffer(
         self,
@@ -1637,7 +1633,6 @@ class MLATokenToKVPool(KVCache):
                 cache_k_nope,
                 cache_k_rope,
             )
-        self._notify_layer_kv_stored(layer_id - self.start_layer)
 
     def get_mla_kv_buffer(
         self,
@@ -1859,7 +1854,6 @@ class NSATokenToKVPool(MLATokenToKVPool):
             use_nsa=True,
             override_kv_cache_dim=override_dim,
         )
-        self.layer_backup_notify_after_indexer = True
         # self.index_k_dtype = torch.float8_e4m3fn
         # self.index_k_scale_dtype = torch.float32
         self.index_head_dim = index_head_dim
@@ -1966,7 +1960,6 @@ class NSATokenToKVPool(MLATokenToKVPool):
         index_buf_accessor.SetKAndS.execute(
             pool=self, buf=buf, loc=loc, index_k=index_k, index_k_scale=index_k_scale
         )
-        self._notify_layer_kv_stored(layer_id - self.start_layer, source="indexer")
 
     def get_state_buf_infos(self):
         data_ptrs = [

python/sglang/srt/mem_cache/radix_cache.py

@@ -489,22 +489,40 @@ class RadixCache(BasePrefixCache):
         keys = page_align_keys(keys, self.page_size)
         values = kv_indices[: len(keys)].to(dtype=torch.int64, copy=True)
         radix_key = RadixKey(keys, req.extra_key, is_bigram=self.is_eagle)
+        prepared_cp_backup = getattr(req, "cp_hicache_prepared_backup", None)
 
         # Radix Cache takes one ref in memory pool
         if is_insert:
             priority = getattr(req, "priority", 0) or 0
             result = self.insert(
-                InsertParams(key=radix_key, value=values, priority=priority)
+                InsertParams(
+                    key=radix_key,
+                    value=values,
+                    priority=priority,
+                    cp_hicache_prepared_backup=prepared_cp_backup,
+                )
             )
+            if (
+                prepared_cp_backup is not None
+                and not getattr(prepared_cp_backup, "attached", False)
+                and hasattr(self, "_rollback_prepared_cp_backup")
+            ):
+                self._rollback_prepared_cp_backup(prepared_cp_backup, "insert_miss")
             new_prefix_len = result.prefix_len
             # Free the duplicates that were already in the tree
             self.token_to_kv_pool_allocator.free(
                 kv_indices[req.cache_protected_len : new_prefix_len]
             )
         else:
+            if prepared_cp_backup is not None and hasattr(
+                self, "_rollback_prepared_cp_backup"
+            ):
+                self._rollback_prepared_cp_backup(prepared_cp_backup, "no_insert")
             self.token_to_kv_pool_allocator.free(
                 kv_indices[req.cache_protected_len : len(keys)]
             )
+        if prepared_cp_backup is not None:
+            req.cp_hicache_prepared_backup = None
 
         # free the unaligned tail
         self.token_to_kv_pool_allocator.free(kv_indices[len(keys) :])
@@ -527,6 +545,7 @@ class RadixCache(BasePrefixCache):
         keys = page_align_keys(keys, self.page_size)
         values = kv_indices[: len(keys)].to(dtype=torch.int64, copy=True)
         radix_key = RadixKey(keys, req.extra_key, is_bigram=self.is_eagle)
+        prepared_cp_backup = getattr(req, "cp_hicache_prepared_backup", None)
 
         # Radix Cache takes one ref in memory pool
         result = self.insert(
@@ -535,8 +554,17 @@ class RadixCache(BasePrefixCache):
                 value=values,
                 chunked=chunked,
                 priority=getattr(req, "priority", 0) or 0,
+                cp_hicache_prepared_backup=prepared_cp_backup,
             )
         )
+        if (
+            prepared_cp_backup is not None
+            and not getattr(prepared_cp_backup, "attached", False)
+            and hasattr(self, "_rollback_prepared_cp_backup")
+        ):
+            self._rollback_prepared_cp_backup(prepared_cp_backup, "insert_miss")
+        if prepared_cp_backup is not None:
+            req.cp_hicache_prepared_backup = None
         new_prefix_len = result.prefix_len
 
         self.token_to_kv_pool_allocator.free(

python/sglang/srt/models/deepseek_v2.py

@@ -1620,6 +1620,22 @@ class DeepseekV2DecoderLayer(nn.Module):
             and layer_id % self.config.moe_layer_freq == 0
         )
 
+    def _notify_cp_hicache_layer_end(
+        self,
+        forward_batch: ForwardBatch,
+        tbo_subbatch_index: Optional[int] = None,
+    ) -> None:
+        if (
+            getattr(forward_batch, "tbo_parent_token_range", None) is not None
+            and tbo_subbatch_index != 1
+        ):
+            return
+        notifier = getattr(
+            forward_batch.token_to_kv_pool, "notify_layer_end_for_backup", None
+        )
+        if notifier is not None:
+            notifier(self.layer_id)
+
     def forward(
         self,
         positions: torch.Tensor,
@@ -1707,6 +1723,7 @@ class DeepseekV2DecoderLayer(nn.Module):
                 hidden_states, residual, forward_batch
             )
 
+        self._notify_cp_hicache_layer_end(forward_batch)
         return hidden_states, residual
 
     def op_comm_prepare_attn(
@@ -1761,6 +1778,9 @@ class DeepseekV2DecoderLayer(nn.Module):
             state.pop("residual_after_comm_pre_mlp"),
             state.forward_batch,
         )
+        self._notify_cp_hicache_layer_end(
+            state.forward_batch, tbo_subbatch_index=state.tbo_subbatch_index
+        )
 
         output = dict(
             positions=state.positions,