fix(pd): correct DeepSeek V4 layerwise cache handoff

python/tokenspeed/runtime/layers/attention/kv_cache/deepseek_v4.py

@@ -1220,20 +1220,26 @@ def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor) -> None:
                     rows = buf.view(-1, buf.shape[-1])
                     rows[tgt_loc.long()] = rows[src_loc.long()]
 
-    def _all_buffers(self) -> list[torch.Tensor]:
-        out: list[torch.Tensor] = []
-        for layer_id in range(self.layer_num):
-            out.append(self.swa_kv_buffer[layer_id])
-            for buffers in (
-                self.compressed_kv_buffer,
-                self.compressor_state_buffer,
-                self.indexer_kv_buffer,
-                self.indexer_state_buffer,
+    def _all_buffer_entries(self) -> list[tuple[torch.Tensor, str]]:
+        entries: list[tuple[torch.Tensor, str]] = []
+        for layer_id, ratio in enumerate(self.layout.layer_ratio):
+            entries.append((self.swa_kv_buffer[layer_id], V4_SWA_KV_GROUP_ID))
+            for buffers, group_id in (
+                (self.compressed_kv_buffer, v4_compressed_kv_group_id(ratio)),
+                (self.compressor_state_buffer, v4_compressor_state_group_id(ratio)),
+                (self.indexer_kv_buffer, v4_compressed_kv_group_id(ratio)),
+                (
+                    self.indexer_state_buffer,
+                    V4_INDEXER_COMPRESSOR_STATE_GROUP_ID,
+                ),
             ):
                 buf = buffers[layer_id]
                 if buf is not None:
-                    out.append(buf)
-        return out
+                    entries.append((buf, group_id))
+        return entries
+
+    def _all_buffers(self) -> list[torch.Tensor]:
+        return [buffer for buffer, _ in self._all_buffer_entries()]
 
     def get_kv_size_bytes(self) -> int:
         return int(
@@ -1248,6 +1254,9 @@ def get_contiguous_buf_infos(self):
             [buf[0].nbytes for buf in buffers],
         )
 
+    def get_contiguous_buf_group_ids(self) -> list[str]:
+        return [group_id for _, group_id in self._all_buffer_entries()]
+
     def get_layerwise_buf_info_offsets(self, start_idx=0):
         offsets = []
         cursor = start_idx

python/tokenspeed/runtime/models/deepseek_v4.py

@@ -3746,59 +3746,66 @@ def run_compressor() -> None:
         forward_mode = ctx.forward_mode
         if forward_mode is None:
             raise RuntimeError("DeepSeek V4 attention requires forward mode")
-        if forward_mode.is_mixed():
-            with nvtx_range(f"{profile_prefix}_mixed_backend"):
-                attn_output = ctx.attn_backend.forward_deepseek_v4_mixed(
-                    q=q,
-                    positions=positions,
-                    token_to_kv_pool=pool,
-                    layer_id=self.cache_layer_index,
-                    kind=self.attention_kind,
-                    compress_ratio=self.compress_ratio,
-                    num_local_heads=self.num_local_heads,
-                    padded_heads=self.padded_heads,
-                    head_dim=self.head_dim,
-                    window_size=self.swa_window,
-                    softmax_scale=self.scale,
-                    attn_sink=self.attn_sink,
-                    topk_indices=topk_indices,
-                )
-        elif forward_mode.is_decode():
-            with nvtx_range(f"{profile_prefix}_decode_backend"):
-                attn_output = ctx.attn_backend.forward_deepseek_v4_decode(
-                    q=q,
-                    positions=positions,
-                    token_to_kv_pool=pool,
-                    layer_id=self.cache_layer_index,
-                    kind=self.attention_kind,
-                    compress_ratio=self.compress_ratio,
-                    num_local_heads=self.num_local_heads,
-                    padded_heads=self.padded_heads,
-                    head_dim=self.head_dim,
-                    window_size=self.swa_window,
-                    softmax_scale=self.scale,
-                    attn_sink=self.attn_sink,
-                    topk_indices=topk_indices,
-                )
-        elif forward_mode.is_extend_or_mixed():
-            with nvtx_range(f"{profile_prefix}_prefill_backend"):
-                attn_output = ctx.attn_backend.forward_deepseek_v4_prefill(
-                    q=q,
-                    positions=positions,
-                    token_to_kv_pool=pool,
-                    layer_id=self.cache_layer_index,
-                    kind=self.attention_kind,
-                    compress_ratio=self.compress_ratio,
-                    num_local_heads=self.num_local_heads,
-                    padded_heads=self.padded_heads,
-                    head_dim=self.head_dim,
-                    window_size=self.swa_window,
-                    softmax_scale=self.scale,
-                    attn_sink=self.attn_sink,
-                    topk_indices=topk_indices,
+        with ctx.attn_backend.record_pd_cache_step(
+            forward_mode,
+            save_kv_cache=False,
+            record_kv_cache=None,
+        ):
+            if forward_mode.is_mixed():
+                with nvtx_range(f"{profile_prefix}_mixed_backend"):
+                    attn_output = ctx.attn_backend.forward_deepseek_v4_mixed(
+                        q=q,
+                        positions=positions,
+                        token_to_kv_pool=pool,
+                        layer_id=self.cache_layer_index,
+                        kind=self.attention_kind,
+                        compress_ratio=self.compress_ratio,
+                        num_local_heads=self.num_local_heads,
+                        padded_heads=self.padded_heads,
+                        head_dim=self.head_dim,
+                        window_size=self.swa_window,
+                        softmax_scale=self.scale,
+                        attn_sink=self.attn_sink,
+                        topk_indices=topk_indices,
+                    )
+            elif forward_mode.is_decode():
+                with nvtx_range(f"{profile_prefix}_decode_backend"):
+                    attn_output = ctx.attn_backend.forward_deepseek_v4_decode(
+                        q=q,
+                        positions=positions,
+                        token_to_kv_pool=pool,
+                        layer_id=self.cache_layer_index,
+                        kind=self.attention_kind,
+                        compress_ratio=self.compress_ratio,
+                        num_local_heads=self.num_local_heads,
+                        padded_heads=self.padded_heads,
+                        head_dim=self.head_dim,
+                        window_size=self.swa_window,
+                        softmax_scale=self.scale,
+                        attn_sink=self.attn_sink,
+                        topk_indices=topk_indices,
+                    )
+            elif forward_mode.is_extend_or_mixed():
+                with nvtx_range(f"{profile_prefix}_prefill_backend"):
+                    attn_output = ctx.attn_backend.forward_deepseek_v4_prefill(
+                        q=q,
+                        positions=positions,
+                        token_to_kv_pool=pool,
+                        layer_id=self.cache_layer_index,
+                        kind=self.attention_kind,
+                        compress_ratio=self.compress_ratio,
+                        num_local_heads=self.num_local_heads,
+                        padded_heads=self.padded_heads,
+                        head_dim=self.head_dim,
+                        window_size=self.swa_window,
+                        softmax_scale=self.scale,
+                        attn_sink=self.attn_sink,
+                        topk_indices=topk_indices,
+                    )
+            else:
+                raise RuntimeError(
+                    f"Unsupported DeepSeek V4 forward mode: {forward_mode}"
                 )
-        else:
-            raise RuntimeError(f"Unsupported DeepSeek V4 forward mode: {forward_mode}")
         with nvtx_range(f"{profile_prefix}_output_proj"):
             return self._project_attention_output(
                 attn_output,

python/tokenspeed/runtime/pd/decode_executor.py

@@ -26,6 +26,9 @@
 from tokenspeed.runtime.pd.base.bootstrap import BootstrapInfo
 from tokenspeed.runtime.pd.base.status import TransferPoll
 from tokenspeed.runtime.pd.mooncake.decode import MooncakeKVManagerDecode
+from tokenspeed.runtime.pd.mooncake.entities import (
+    paged_cache_pages_from_forward_op,
+)
 from tokenspeed.runtime.pd.mooncake.receiver import MooncakeKVReceiver
 from tokenspeed.runtime.pd.utils import (
     TransferBackend,
@@ -123,13 +126,15 @@ def _prefill(self, op):
             )
             aux_index = op.request_pool_indices[i]
             mamba_indices = self._mamba_transfer_indices(op, i)
+            paged_cache_pages = paged_cache_pages_from_forward_op(op, i)
             self._request_pool_indices[request_id] = aux_index
             self.receivers[request_id].prefill(
                 kv_indices,
                 aux_index,
                 extend_prefix_len,
                 None,  # mla_l1_5_args
                 mamba_indices,
+                paged_cache_pages,
             )
 
     def register(self, request_id: str, bootstrap_info: BootstrapInfo):

python/tokenspeed/runtime/pd/factory.py

@@ -38,6 +38,25 @@ def _get_contiguous_buf_unit_lens(pool, item_lens):
     return unit_lens
 
 
+def _get_contiguous_buf_group_ids(pool, item_lens):
+    getter = getattr(pool, "get_contiguous_buf_group_ids", None)
+    if getter is None:
+        return [None] * len(item_lens)
+    group_ids = list(getter())
+    if len(group_ids) != len(item_lens):
+        raise ValueError(
+            f"contiguous buffer group count mismatch: groups={len(group_ids)}, items={len(item_lens)}"
+        )
+    return group_ids
+
+
+def _get_paged_cache_group_tokens_per_page(pool):
+    return {
+        str(spec.group_id): int(spec.rows_per_page) * int(spec.entry_stride_tokens)
+        for spec in getattr(pool, "paged_cache_group_specs", ())
+    }
+
+
 def get_kv_args(
     engine_rank: int,
     gpu_id,
@@ -50,6 +69,10 @@ def get_kv_args(
         token_to_kv_pool.get_contiguous_buf_infos()
     )
     kv_unit_lens = _get_contiguous_buf_unit_lens(token_to_kv_pool, kv_item_lens)
+    kv_group_ids = _get_contiguous_buf_group_ids(token_to_kv_pool, kv_item_lens)
+    paged_cache_group_tokens_per_page = _get_paged_cache_group_tokens_per_page(
+        token_to_kv_pool
+    )
     # [[layer0buf0, layer0buf1...], [layer1buf0, layer1buf1...], ...]
     offsets = token_to_kv_pool.get_layerwise_buf_info_offsets()
     target_layer_num = token_to_kv_pool.layer_num
@@ -73,6 +96,12 @@ def get_kv_args(
         kv_data_lens += draft_kv_data_lens
         kv_item_lens += draft_kv_item_lens
         kv_unit_lens += draft_kv_unit_lens
+        kv_group_ids += _get_contiguous_buf_group_ids(
+            draft_token_to_kv_pool, draft_kv_item_lens
+        )
+        paged_cache_group_tokens_per_page.update(
+            _get_paged_cache_group_tokens_per_page(draft_token_to_kv_pool)
+        )
         offsets += draft_offsets
         draft_base_layer_id = (
             max(kv_layer_ids) + 1 if kv_layer_ids else target_layer_num
@@ -106,6 +135,8 @@ def get_kv_args(
         draft_layer_num=draft_layer_num,
         kv_layer_ids=kv_layer_ids,
         kv_unit_lens=kv_unit_lens,
+        kv_group_ids=kv_group_ids,
+        paged_cache_group_tokens_per_page=paged_cache_group_tokens_per_page,
         state_data_ptrs=state_data_ptrs,
         state_data_lens=state_data_lens,
         state_item_lens=state_item_lens,

python/tokenspeed/runtime/pd/mooncake/async_conn.py

@@ -30,6 +30,7 @@
 from tokenspeed.runtime.pd.base.status import TransferPoll
 from tokenspeed.runtime.pd.mooncake.entities import (
     KVArgs,
+    PagedCachePages,
     TransferInfo,
     TransferKVChunk,
 )
@@ -54,6 +55,9 @@ class WriteRequest:
     dst_ranks_info: tuple[str, int, int]
     prefill_kv_blocks: npt.NDArray[np.int64]
     dst_kv_blocks: npt.NDArray[np.int64]
+    paged_cache_blocks: dict[str, tuple[list, list]] = dataclasses.field(
+        default_factory=dict
+    )
     submit_bids: list[int] = dataclasses.field(default_factory=list)
 
 
@@ -166,13 +170,13 @@ def add_transfer_request(
         is_last: bool,
         aux_index: int | None = None,
         mla_l1_5_args: tuple[Any, Any] | None = None,
+        paged_cache_pages: dict[str, PagedCachePages] | None = None,
     ):
         logger.debug("async manager add_transfer_request")
         if self.disaggregation_mode != DisaggregationMode.PREFILL:
             raise RuntimeError("Transfer requests can only be added in prefill mode.")
         if is_last and aux_index is None:
             raise ValueError("aux_index must be set for the last transfer chunk.")
-
         if (
             bootstrap_room not in self.request_status
             or self.check_status(bootstrap_room) == TransferPoll.Failed
@@ -202,6 +206,7 @@ def add_transfer_request(
             is_last=is_last,
             prefill_aux_index=aux_index,
             mla_l1_5_args=mla_l1_5_args,
+            prefill_paged_cache_pages=paged_cache_pages or {},
         )
         self.current_transfer_batch.append((kv_chunk, shard_idx))
 
@@ -243,6 +248,7 @@ def submit_layer_cache(
         end_layer_id: int,
         prefill_kv_blocks: npt.NDArray[np.int64],
         dst_kv_blocks: npt.NDArray[np.int64],
+        paged_cache_blocks: dict[str, tuple[list, list]],
     ) -> int:
         dst_kv_ptrs = self.decode_kv_args_table[mooncake_session_id].dst_kv_ptrs
         transfer_blocks = []
@@ -251,7 +257,21 @@ def submit_one_cache(ptr_offset):
             src_ptr = self.kv_args.kv_data_ptrs[ptr_offset]
             dst_ptr = dst_kv_ptrs[ptr_offset]
             item_len = self.kv_args.kv_item_lens[ptr_offset]
-            for prefill_index, decode_index in zip(prefill_kv_blocks, dst_kv_blocks):
+            buffer_group_id = (
+                self.kv_args.kv_group_ids[ptr_offset]
+                if ptr_offset < len(self.kv_args.kv_group_ids)
+                else None
+            )
+            src_blocks = prefill_kv_blocks
+            dst_blocks = dst_kv_blocks
+            if buffer_group_id is not None:
+                group_blocks = paged_cache_blocks.get(buffer_group_id)
+                if group_blocks is None:
+                    raise RuntimeError(
+                        f"Missing transfer blocks for paged-cache group {buffer_group_id!r}"
+                    )
+                src_blocks, dst_blocks = group_blocks
+            for prefill_index, decode_index in zip(src_blocks, dst_blocks):
                 src_addr = src_ptr + int(prefill_index[0]) * item_len
                 dst_addr = dst_ptr + int(decode_index[0]) * item_len
                 length = item_len * len(prefill_index)
@@ -355,6 +375,7 @@ def initialize(tasks: list[LayerWiseTask]) -> list[LayerWiseTask]:
                             break
 
                     resolved = self.resolve_transfer_indices(kv_chunk, req)
+                    paged_cache_blocks = self.resolve_paged_cache_blocks(kv_chunk, req)
 
                     # Group by indices
                     prefill_kv_blocks, dst_kv_blocks = group_concurrent_contiguous(
@@ -366,6 +387,7 @@ def initialize(tasks: list[LayerWiseTask]) -> list[LayerWiseTask]:
                             dst_ranks_info=(req.endpoint, req.dst_port, req.room),
                             prefill_kv_blocks=prefill_kv_blocks,
                             dst_kv_blocks=dst_kv_blocks,
+                            paged_cache_blocks=paged_cache_blocks,
                         )
                     )
 
@@ -395,6 +417,7 @@ def submit_transfer(
                                 task.next_layer_id + 1,
                                 req.prefill_kv_blocks,
                                 req.dst_kv_blocks,
+                                req.paged_cache_blocks,
                             )
 
                             if ret != 0: