diff --git a/python/tokenspeed/runtime/layers/attention/kv_cache/deepseek_v4.py b/python/tokenspeed/runtime/layers/attention/kv_cache/deepseek_v4.py index dd4785814..4e75e4245 100644 --- a/python/tokenspeed/runtime/layers/attention/kv_cache/deepseek_v4.py +++ b/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 diff --git a/python/tokenspeed/runtime/models/deepseek_v4.py b/python/tokenspeed/runtime/models/deepseek_v4.py index 18b09b5dd..ad35eabd3 100644 --- a/python/tokenspeed/runtime/models/deepseek_v4.py +++ b/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, diff --git a/python/tokenspeed/runtime/pd/decode_executor.py b/python/tokenspeed/runtime/pd/decode_executor.py index cdc198aca..6e93dd844 100644 --- a/python/tokenspeed/runtime/pd/decode_executor.py +++ b/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,6 +126,7 @@ 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, @@ -130,6 +134,7 @@ def _prefill(self, op): extend_prefix_len, None, # mla_l1_5_args mamba_indices, + paged_cache_pages, ) def register(self, request_id: str, bootstrap_info: BootstrapInfo): diff --git a/python/tokenspeed/runtime/pd/factory.py b/python/tokenspeed/runtime/pd/factory.py index 5b8d1ca7c..cc1870171 100644 --- a/python/tokenspeed/runtime/pd/factory.py +++ b/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, diff --git a/python/tokenspeed/runtime/pd/mooncake/async_conn.py b/python/tokenspeed/runtime/pd/mooncake/async_conn.py index c1bc84291..668af5e01 100644 --- a/python/tokenspeed/runtime/pd/mooncake/async_conn.py +++ b/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: diff --git a/python/tokenspeed/runtime/pd/mooncake/entities.py b/python/tokenspeed/runtime/pd/mooncake/entities.py index d34f5ebc2..bf80486fb 100644 --- a/python/tokenspeed/runtime/pd/mooncake/entities.py +++ b/python/tokenspeed/runtime/pd/mooncake/entities.py @@ -19,6 +19,7 @@ # SOFTWARE. import dataclasses +import json import struct import numpy as np @@ -30,6 +31,78 @@ ) from tokenspeed.runtime.pd.utils import PageTransferMetadata +_PAGED_CACHE_PAGES_PREFIX = b"tokenspeed-paged-cache-pages-v1:" + + +@dataclasses.dataclass(frozen=True) +class PagedCachePages: + page_ids: npt.NDArray[np.int64] + base_logical_page: int = 0 + + +def paged_cache_pages_from_forward_op( + op, + index: int, + tokens_per_page: dict[str, int] | None = None, + token_begin: int | None = None, + token_end: int | None = None, + include_partial_page: bool = True, +) -> dict[str, PagedCachePages]: + tables = dict(getattr(op, "paged_cache_block_tables", {})) + base_offsets = dict(getattr(op, "paged_cache_block_table_base_offsets", {})) + pages: dict[str, PagedCachePages] = {} + for group_id, table in tables.items(): + row = np.asarray(table[index], dtype=np.int64) + row = row[row >= 0] + offsets = base_offsets.get(group_id) + base = int(offsets[index]) if offsets is not None else 0 + group_id = str(group_id) + page_tokens = (tokens_per_page or {}).get(group_id) + if ( + page_tokens is not None + and token_begin is not None + and token_end is not None + ): + logical_begin = token_begin // page_tokens + if include_partial_page: + logical_end = (token_end + page_tokens - 1) // page_tokens + else: + logical_end = token_end // page_tokens + selected_begin = max(logical_begin, base) + selected_end = min(logical_end, base + len(row)) + if selected_begin >= selected_end: + pages[group_id] = PagedCachePages( + np.array([], dtype=np.int64), selected_begin + ) + continue + row_begin = selected_begin - base + row = row[row_begin : row_begin + selected_end - selected_begin] + base = selected_begin + pages[group_id] = PagedCachePages(row, base) + return pages + + +def encode_paged_cache_pages(pages: dict[str, PagedCachePages]) -> bytes: + payload = { + group_id: [entry.base_logical_page, entry.page_ids.tolist()] + for group_id, entry in pages.items() + } + return _PAGED_CACHE_PAGES_PREFIX + json.dumps( + payload, separators=(",", ":"), sort_keys=True + ).encode("utf-8") + + +def decode_paged_cache_pages(frame: bytes) -> dict[str, PagedCachePages]: + if not frame.startswith(_PAGED_CACHE_PAGES_PREFIX): + return {} + payload = json.loads(frame[len(_PAGED_CACHE_PAGES_PREFIX) :].decode("utf-8")) + return { + str(group_id): PagedCachePages( + np.asarray(entry[1], dtype=np.int64), int(entry[0]) + ) + for group_id, entry in payload.items() + } + @dataclasses.dataclass class KVArgs: @@ -47,6 +120,10 @@ class KVArgs: draft_layer_num: int kv_layer_ids: list[int] = dataclasses.field(default_factory=list) kv_unit_lens: list[int] = dataclasses.field(default_factory=list) + kv_group_ids: list[str | None] = dataclasses.field(default_factory=list) + paged_cache_group_tokens_per_page: dict[str, int] = dataclasses.field( + default_factory=dict + ) state_data_ptrs: list[int] = dataclasses.field(default_factory=list) state_data_lens: list[int] = dataclasses.field(default_factory=list) state_item_lens: list[int] = dataclasses.field(default_factory=list) @@ -88,6 +165,9 @@ class TransferKVChunk: layerwise_interval: int = 1 wait_for_bootstrap_token: bool = False spec_candidate_ids: list[int] | None = None + prefill_paged_cache_pages: dict[str, PagedCachePages] = dataclasses.field( + default_factory=dict + ) @dataclasses.dataclass @@ -114,10 +194,14 @@ class TransferInfo: dst_mamba_indices: npt.NDArray[np.int64] | None is_dummy: bool transfer_fragments: tuple[TransferFragment, ...] = () + dst_paged_cache_pages: dict[str, PagedCachePages] = dataclasses.field( + default_factory=dict + ) @classmethod def from_zmq(cls, msg: list[bytes]): transfer_fragments = () + dst_paged_cache_pages = decode_paged_cache_pages(msg[-1]) if msg else {} if msg[4] == b"" and msg[5] == b"": dst_kv_indices = np.array([], dtype=np.int64) dst_aux_index = None @@ -172,6 +256,7 @@ def from_zmq(cls, msg: list[bytes]): dst_mamba_indices=dst_mamba_indices, is_dummy=is_dummy, transfer_fragments=transfer_fragments, + dst_paged_cache_pages=dst_paged_cache_pages, ) diff --git a/python/tokenspeed/runtime/pd/mooncake/prefill.py b/python/tokenspeed/runtime/pd/mooncake/prefill.py index 8dd480d0b..233f4daa2 100644 --- a/python/tokenspeed/runtime/pd/mooncake/prefill.py +++ b/python/tokenspeed/runtime/pd/mooncake/prefill.py @@ -35,6 +35,7 @@ KVArgs, KVArgsRegisterInfo, KVManagerArgs, + PagedCachePages, TransferIndexResolution, TransferInfo, TransferKVChunk, @@ -342,6 +343,43 @@ def resolve_transfer_indices( dst_indices=dst_local, ) + def resolve_paged_cache_blocks( + self, + kv_chunk: TransferKVChunk, + req: TransferInfo, + ) -> dict[str, tuple[list, list]]: + required_group_ids = { + group_id for group_id in self.kv_args.kv_group_ids if group_id is not None + } + missing_src = required_group_ids - kv_chunk.prefill_paged_cache_pages.keys() + missing_dst = required_group_ids - req.dst_paged_cache_pages.keys() + if missing_src or missing_dst: + raise RuntimeError( + "Missing paged-cache transfer metadata: " + f"source={sorted(missing_src)}, destination={sorted(missing_dst)}" + ) + blocks: dict[str, tuple[list, list]] = {} + for group_id, src in kv_chunk.prefill_paged_cache_pages.items(): + dst = req.dst_paged_cache_pages.get(group_id) + if dst is None: + continue + logical_begin = max(src.base_logical_page, dst.base_logical_page) + logical_end = min( + src.base_logical_page + len(src.page_ids), + dst.base_logical_page + len(dst.page_ids), + ) + if logical_begin >= logical_end: + blocks[group_id] = ([], []) + continue + src_begin = logical_begin - src.base_logical_page + dst_begin = logical_begin - dst.base_logical_page + page_count = logical_end - logical_begin + blocks[group_id] = group_concurrent_contiguous( + src.page_ids[src_begin : src_begin + page_count], + dst.page_ids[dst_begin : dst_begin + page_count], + ) + return blocks + def _transfer_data(self, mooncake_session_id, transfer_blocks): if not transfer_blocks: return 0 @@ -359,6 +397,7 @@ def send_kvcache( dst_kv_indices: npt.NDArray[np.int64], executor: concurrent.futures.ThreadPoolExecutor, transfer_fragments: tuple[TransferFragment, ...] = (), + paged_cache_blocks: dict[str, tuple[list, list]] | None = None, ): # Group by indices prefill_kv_blocks, dst_kv_blocks = group_concurrent_contiguous( @@ -380,6 +419,7 @@ def send_kvcache( begin_layer_id=0, end_layer_id=len(self.kv_args.offsets), transfer_fragments=transfer_fragments, + paged_cache_blocks=paged_cache_blocks, ) return self._transfer_data(mooncake_session_id, transfer_blocks) @@ -391,6 +431,7 @@ def _layer_transfer_blocks( begin_layer_id: int, end_layer_id: int, transfer_fragments: tuple[TransferFragment, ...] = (), + paged_cache_blocks: dict[str, tuple[list, list]] | None = None, ) -> list[tuple[int, int, int]]: transfer_blocks = [] fragments_by_buffer: dict[int, list[TransferFragment]] = defaultdict(list) @@ -405,9 +446,25 @@ def _layer_transfer_blocks( dst_ptr = dst_ptrs[ptr_offset] item_len = self.kv_args.kv_item_lens[ptr_offset] buffer_fragments = fragments_by_buffer.get(ptr_offset, ()) + buffer_group_id = ( + self.kv_args.kv_group_ids[ptr_offset] + if ptr_offset < len(self.kv_args.kv_group_ids) + else None + ) + buffer_src_blocks = src_blocks + buffer_dst_blocks = dst_blocks + if buffer_group_id is not None: + group_blocks = (paged_cache_blocks or {}).get(buffer_group_id) + if group_blocks is None: + raise RuntimeError( + f"Missing transfer blocks for paged-cache group {buffer_group_id!r}" + ) + buffer_src_blocks, buffer_dst_blocks = group_blocks if has_fragment_plan: for fragment in buffer_fragments: - for prefill_index, decode_index in zip(src_blocks, dst_blocks): + for prefill_index, decode_index in zip( + buffer_src_blocks, buffer_dst_blocks + ): page_count = min(len(prefill_index), len(decode_index)) if fragment.page_count is not None: page_count = min(page_count, fragment.page_count) @@ -429,7 +486,9 @@ def _layer_transfer_blocks( ) continue - for prefill_index, decode_index in zip(src_blocks, dst_blocks): + for prefill_index, decode_index in zip( + buffer_src_blocks, buffer_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) @@ -591,6 +650,7 @@ def send_kvcache_layerwise( prefill_mamba_indices: npt.NDArray[np.int64] | None = None, dst_mamba_indices: npt.NDArray[np.int64] | None = None, transfer_fragments: tuple[TransferFragment, ...] = (), + paged_cache_blocks: dict[str, tuple[list, list]] | None = None, ) -> int: prefill_kv_blocks, dst_kv_blocks = group_concurrent_contiguous( prefill_kv_indices, dst_kv_indices @@ -613,7 +673,7 @@ def send_kvcache_layerwise( self._wait_until_cache_step(target_step) transfer_blocks = [] - if prefill_kv_blocks: + if prefill_kv_blocks or paged_cache_blocks: for global_layer_id in range(begin_layer_id, end_layer_id): kv_layer_index = self._kv_layer_to_index.get(global_layer_id) if kv_layer_index is None: @@ -626,6 +686,7 @@ def send_kvcache_layerwise( begin_layer_id=kv_layer_index, end_layer_id=kv_layer_index + 1, transfer_fragments=transfer_fragments, + paged_cache_blocks=paged_cache_blocks, ) ) if transfer_blocks: @@ -758,6 +819,9 @@ def transfer_worker( ) break resolved = self.resolve_transfer_indices(kv_chunk, req) + paged_cache_blocks = self.resolve_paged_cache_blocks( + kv_chunk, req + ) logger.debug( "[TRANSFER_WORKER] Calling send_kvcache for room %s, session %s", @@ -780,6 +844,7 @@ def transfer_worker( resolved.dst_indices, executor, req.transfer_fragments, + paged_cache_blocks, ) else: ret = self.send_kvcache_layerwise( @@ -793,6 +858,7 @@ def transfer_worker( kv_chunk.prefill_mamba_indices, req.dst_mamba_indices, req.transfer_fragments, + paged_cache_blocks, ) if ret == 0 and kv_chunk.is_last: if kv_chunk.wait_for_bootstrap_token: @@ -993,6 +1059,7 @@ def add_transfer_request( wait_for_bootstrap_token: bool = False, mamba_indices: npt.NDArray[np.int64] | None = None, spec_candidate_ids: list[int] | None = None, + paged_cache_pages: dict[str, PagedCachePages] | None = None, ): if self.disaggregation_mode != DisaggregationMode.PREFILL: raise RuntimeError("Transfer requests can only be added in prefill mode.") @@ -1034,6 +1101,7 @@ def add_transfer_request( wait_for_bootstrap_token=wait_for_bootstrap_token, prefill_mamba_indices=mamba_indices, spec_candidate_ids=spec_candidate_ids, + prefill_paged_cache_pages=paged_cache_pages or {}, ) ) diff --git a/python/tokenspeed/runtime/pd/mooncake/receiver.py b/python/tokenspeed/runtime/pd/mooncake/receiver.py index d9472eeed..c2c9c79ce 100644 --- a/python/tokenspeed/runtime/pd/mooncake/receiver.py +++ b/python/tokenspeed/runtime/pd/mooncake/receiver.py @@ -29,7 +29,11 @@ import zmq from tokenspeed.runtime.pd.base.status import TransferPoll -from tokenspeed.runtime.pd.mooncake.entities import KVTransferError +from tokenspeed.runtime.pd.mooncake.entities import ( + KVTransferError, + PagedCachePages, + encode_paged_cache_pages, +) from tokenspeed.runtime.pd.transfer_plan import ( BufferKind, BufferLayout, @@ -601,6 +605,7 @@ def prefill( decode_prefix_len: int | None = 0, mla_l1_5_args: PageTransferMetadata | None = None, mamba_indices: npt.NDArray[np.int64] | None = None, + paged_cache_pages: dict[str, PagedCachePages] | None = None, ): logger.info( "[MooncakeKVReceiver.init] bootstrap_room=%s kv_indices_len=%d aux_index=%s decode_prefix_len=%s", @@ -673,6 +678,8 @@ def prefill( transfer_fragments = bootstrap_info.get("transfer_fragments", ()) if not is_dummy and transfer_fragments: message_parts.extend(encode_transfer_fragments(transfer_fragments)) + if not is_dummy and paged_cache_pages: + message_parts.append(encode_paged_cache_pages(paged_cache_pages)) sock.send_multipart(message_parts) self.init_time = time.time() diff --git a/python/tokenspeed/runtime/pd/mooncake/sender.py b/python/tokenspeed/runtime/pd/mooncake/sender.py index afb8e2057..38d888b22 100644 --- a/python/tokenspeed/runtime/pd/mooncake/sender.py +++ b/python/tokenspeed/runtime/pd/mooncake/sender.py @@ -24,7 +24,10 @@ import numpy.typing as npt from tokenspeed.runtime.pd.base.status import TransferPoll -from tokenspeed.runtime.pd.mooncake.entities import KVTransferError +from tokenspeed.runtime.pd.mooncake.entities import ( + KVTransferError, + PagedCachePages, +) from tokenspeed.runtime.pd.utils import PageTransferMetadata from tokenspeed.runtime.utils import get_colorful_logger @@ -66,6 +69,7 @@ def send( bootstrap_token: int = -1, spec_candidate_ids: list[int] | None = None, mamba_indices: npt.NDArray[np.int64] | None = None, + paged_cache_pages: dict[str, PagedCachePages] | None = None, ): """ Send the kv cache at the given kv indices to the decoder server @@ -94,6 +98,7 @@ def send( False, mla_l1_5_args=mla_l1_5_args, mamba_indices=mamba_indices, + paged_cache_pages=paged_cache_pages, ) else: self.kv_mgr.add_transfer_request( @@ -106,6 +111,7 @@ def send( bootstrap_token=bootstrap_token, spec_candidate_ids=spec_candidate_ids, mamba_indices=mamba_indices, + paged_cache_pages=paged_cache_pages, ) def send_layerwise( @@ -121,11 +127,15 @@ def send_layerwise( wait_for_bootstrap_token: bool = False, spec_candidate_ids: list[int] | None = None, mamba_indices: npt.NDArray[np.int64] | None = None, + paged_cache_pages: dict[str, PagedCachePages] | None = None, ): self._layerwise_transfer_started = True self.curr_idx = max(self.curr_idx, index_slice.stop) - if len(kv_indices) == 0 and not is_last: + has_paged_cache_pages = any( + entry.page_ids.size > 0 for entry in (paged_cache_pages or {}).values() + ) + if len(kv_indices) == 0 and not is_last and not has_paged_cache_pages: return logger.info( @@ -152,6 +162,7 @@ def send_layerwise( wait_for_bootstrap_token=wait_for_bootstrap_token, spec_candidate_ids=spec_candidate_ids, mamba_indices=mamba_indices, + paged_cache_pages=paged_cache_pages, ) def poll(self) -> TransferPoll: diff --git a/python/tokenspeed/runtime/pd/prefill_executor.py b/python/tokenspeed/runtime/pd/prefill_executor.py index edcd2beb6..ea7278b32 100644 --- a/python/tokenspeed/runtime/pd/prefill_executor.py +++ b/python/tokenspeed/runtime/pd/prefill_executor.py @@ -24,6 +24,9 @@ from tokenspeed.runtime.pd.base.bootstrap import BootstrapInfo from tokenspeed.runtime.pd.base.status import TransferPoll +from tokenspeed.runtime.pd.mooncake.entities import ( + paged_cache_pages_from_forward_op, +) from tokenspeed.runtime.pd.mooncake.prefill import ( MooncakeKVManagerPrefill, MooncakeKVSender, @@ -54,6 +57,9 @@ def __init__( ) self.senders: dict[int, MooncakeKVSender] = {} self.kv_manager = MooncakeKVManagerPrefill(args, kv_args) + self._paged_cache_group_tokens_per_page = ( + kv_args.paged_cache_group_tokens_per_page + ) self.gloo_group = gloo_group self._local_states = {} self._layerwise_enabled = False @@ -207,7 +213,23 @@ def prepare_prefill(self, op) -> None: self._drop_request_state(request_id) continue kv_indices, index_slice, is_last = self._prefill_page_window(op, i, sender) - if len(kv_indices) == 0 and not is_last: + chunk_begin = int(op.extend_prefix_lens[i]) + chunk_end = chunk_begin + int(op.input_lengths[i]) + grouped_begin = self._decode_prefix_len(sender.bootstrap_room) + ( + sender.curr_idx * self.page_size + ) + paged_cache_pages = paged_cache_pages_from_forward_op( + op, + i, + self._paged_cache_group_tokens_per_page, + grouped_begin, + chunk_end, + include_partial_page=is_last, + ) + has_paged_cache_pages = any( + entry.page_ids.size > 0 for entry in paged_cache_pages.values() + ) + if len(kv_indices) == 0 and not is_last and not has_paged_cache_pages: continue mamba_indices = self._mamba_transfer_indices(op, i) if is_last else None sender.send_layerwise( @@ -219,6 +241,7 @@ def prepare_prefill(self, op) -> None: layerwise_interval=self._layerwise_interval, wait_for_bootstrap_token=is_last, mamba_indices=mamba_indices, + paged_cache_pages=paged_cache_pages, ) def _decode(self, op): @@ -258,6 +281,13 @@ def _decode(self, op): dtype=np.int64, ) mamba_indices = self._mamba_transfer_indices(op, i) + paged_cache_pages = paged_cache_pages_from_forward_op( + op, + i, + self._paged_cache_group_tokens_per_page, + decode_prefix_len, + int(op.prefill_lengths[i]), + ) logger.debug( "[prefill][_decode] rid=%s aux_index=%d kv_indices(len=%d)=%s bootstrap_token=%d", request_id, @@ -273,6 +303,7 @@ def _decode(self, op): bootstrap_token=bootstrap_token, spec_candidate_ids=spec_candidate_ids, mamba_indices=mamba_indices, + paged_cache_pages=paged_cache_pages, ) def register(self, request_id: str, bootstrap_info: BootstrapInfo):