From e0c4a7a414a8be371f76016a51f10f73db7f0cd1 Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Thu, 3 Sep 2026 20:42:08 +0800 Subject: [PATCH] [None][refactor] Anchor disagg KV block lists with explicit per-group block ordinals Replace the implicit convention that a layer group's block-ID array is the contiguous resident suffix of [0, slice_end) - with token positions inferred from len() - by explicit per-group block-ordinal anchors (first_ordinals) on KVSlice and RecvReqInfo. Producer: CacheReuseAdapter.get_transfer_span derives (pages, anchor) from manager facts (V1: get_num_front_blocks_removed + scratch strip + reconciliation check; V2: ordinal-indexed pages + backed-run anchor + hole-consistency check). SWA/reuse trims become explicit head-slices with the anchor advanced. Consumer: _build_kv_write_meta collapses to an interval intersection via _align_kv_blocks on both sides' anchors. Deleted: TokenRange, KVSlice.token_range, project_blocks_to_global_chunk, _trim_receiver_window_head, _trim_packed_beam_block_ids, Sender._beam0_block_count, RecvReqInfo.dst_start_token (wire change, no fallback). Beam tails are split exactly from the anchored beam-0 length instead of guessed from beam_width; V1 dangling detached block IDs are stripped before pool translation instead of being passed through. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- .../_torch/disaggregation/base/transfer.py | 43 +- .../_torch/disaggregation/native/transfer.py | 195 +-- .../disaggregation/resource/cache_reuse.py | 210 ++- .../_torch/disaggregation/transceiver.py | 212 ++- .../defs/disaggregated/test_disaggregated.py | 2 +- .../executor/test_mamba_cache_manager.py | 2 + .../disaggregated/test_cache_reuse_adapter.py | 1144 ++++++++++------- .../disaggregated/test_chunked_transfer.py | 696 ++++++++-- .../disaggregated/test_kv_transfer.py | 204 ++- .../disaggregated/test_kv_transfer_mp.py | 2 + 10 files changed, 1767 insertions(+), 943 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/base/transfer.py b/tensorrt_llm/_torch/disaggregation/base/transfer.py index c87147576754..4080c4aa4b41 100644 --- a/tensorrt_llm/_torch/disaggregation/base/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/base/transfer.py @@ -11,23 +11,6 @@ from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest -@dataclass -class TokenRange: - """Half-open token range [start, end) within one request. - - ``KVSlice`` ranges are block-aligned. Empty ranges are valid. - """ - - start: int - end: int # exclusive - - def __post_init__(self): - if self.start < 0 or self.end < 0: - raise ValueError("Token indices must be non-negative") - if self.start > self.end: - raise ValueError(f"Invalid range: [{self.start}, {self.end})") - - @dataclass class LayerRange: """Range of layers to transfer.""" @@ -46,10 +29,16 @@ def __post_init__(self): class KVSlice: """KV-cache blocks for one request transfer slice. - Monolithic transfers omit ``token_range`` and cover ``prompt_len``. - Pipelined transfers use a block-aligned ``token_range`` and mark only the - final chunk with ``is_last_slice``. Block lists may omit cached or evicted - prefixes. + Each layer group's block list is anchored explicitly: ``first_ordinals[i]`` + is the block ordinal (0-based sequence block index) of element 0 of group + ``i``'s beam-0 list, so positions never have to be inferred from list + lengths. The anchor describes only the beam-0 prefix; packed beam tails + (``beam_width > 1`` appends per-beam tail blocks after beam-0) stay outside + the anchored region. STATE groups and empty lists use anchor 0. + + Monolithic transfers cover ``prompt_len``. Pipelined transfers send one + slice per prefill chunk (anchored at the chunk's first block ordinal) and + mark only the final chunk with ``is_last_slice``. """ layer_range: Optional[LayerRange] = None @@ -57,7 +46,9 @@ class KVSlice: default_factory=list ) # Physical block IDs per layer group, each np.ndarray(dtype=np.int64) is_last_slice: bool = False - token_range: Optional[TokenRange] = None + # Block ordinal of element 0 of each group's beam-0 list; parallel to + # block_ids_per_layer_groups. + first_ordinals: List[int] = field(default_factory=list) class SessionStatus(Enum): @@ -96,7 +87,8 @@ class SessionArgsBase: """Base arguments for transfer sessions.""" params: DisaggregatedParams - # Captured from LlmRequest.prompt_len; needed for SWA stale_end derivation. + # Captured from LlmRequest.prompt_len; bounds the beam-0 span + # (ceil(prompt_len / tokens_per_block) block ordinals). prompt_len: int beam_width: int = 1 @@ -156,9 +148,8 @@ def send(self, slice: KVSlice) -> None: Args: slice: The KV slice describing which source blocks to send. - For pipelined chunks, ``token_range`` is the shared sender-side - chunk cursor; each layer group projects it into its own - resident/windowed source and destination block ranges. + For pipelined chunks, each layer group's list is anchored at + the chunk's first block ordinal via ``first_ordinals``. """ ... diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index 4d9d9495d35d..1bab3f8b0fee 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -69,6 +69,7 @@ from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo from tensorrt_llm._torch.disaggregation.native.utils import get_local_ip from tensorrt_llm._torch.disaggregation.nixl.agent import NixlTransferAgent +from tensorrt_llm._torch.disaggregation.resource.cache_reuse import split_packed_beam_block_ids from tensorrt_llm._torch.disaggregation.resource.kv_extractor import KVRegionExtractorV1 from tensorrt_llm._torch.disaggregation.resource.page import CacheKind, KVCachePageTable, MapperKind from tensorrt_llm._torch.disaggregation.resource.utils import get_unique_pool_memory_descs @@ -104,9 +105,10 @@ class RecvReqInfo: np.ndarray ] # Block IDs per layer group, each np.ndarray(dtype=np.int64) unique_rid: int - # Block-aligned token offset where the receiver's block list starts. - # None means "end-of-range suffix" — sender derives it from len(blocks). - dst_start_token: Optional[int] = None + # Block ordinal (0-based sequence block index) of element 0 of each + # group's beam-0 list; parallel to block_ids_per_layer_groups. 0 for + # STATE groups and empty lists (see KVSlice.first_ordinals). + first_ordinals: list[int] aux_slot: Optional[int] = None slice_id: Optional[int] = None bounce_dst_base: Optional[int] = None @@ -121,7 +123,7 @@ def to_bytes(self) -> bytes: arr.tobytes() for arr in self.block_ids_per_layer_groups ], "unique_rid": self.unique_rid, - "dst_start_token": self.dst_start_token, + "first_ordinals": self.first_ordinals, "aux_slot": self.aux_slot, "slice_id": self.slice_id, "bounce_dst_base": self.bounce_dst_base, @@ -130,6 +132,8 @@ def to_bytes(self) -> bytes: @classmethod def from_bytes(cls, data: bytes) -> "RecvReqInfo": + # No wire compatibility: a peer that still sends the pre-anchor layout + # (dst_start_token) fails here loudly instead of misaligning KV. d = msgpack.unpackb(data, raw=False) d["block_ids_per_layer_groups"] = [ np.frombuffer(b, dtype=np.int64).copy() for b in d["block_ids_per_layer_groups"] @@ -149,36 +153,6 @@ class WriteMetaType(Enum): AUX = "AUX" -def project_blocks_to_global_chunk( - block_ids: np.ndarray, - chunk_block_offset: int, - chunk_block_count: int, - resident_block_end: int, -) -> np.ndarray: - """Project a global block chunk into a suffix-resident block list. - - ``block_ids`` represents the resident suffix of the logical range - ``[0, resident_block_end)``. ``chunk_block_offset`` and - ``chunk_block_count`` describe a chunk in that global coordinate space. - """ - if chunk_block_count <= 0 or len(block_ids) == 0: - return block_ids[:0] - - resident_start = max(0, resident_block_end - len(block_ids)) - resident_end = resident_block_end - chunk_start = chunk_block_offset - chunk_end = chunk_start + chunk_block_count - - overlap_start = max(chunk_start, resident_start) - overlap_end = min(chunk_end, resident_end) - if overlap_start >= overlap_end: - return block_ids[:0] - - local_start = overlap_start - resident_start - local_end = overlap_end - resident_start - return block_ids[local_start:local_end] - - @dataclass class WriteMeta: task: Union[SendTaskBase, "KVRecvTask"] @@ -797,7 +771,7 @@ def _align_kv_blocks( 1. No prefix cache on either side → identity (start_token == 0 both) 2. Context prefix cache (src starts later than 0) → trim dst head 3. Generation prefix cache (dst starts later than 0) → trim src head - 4. Chunked context (each slice has its own token_range) → correct + 4. Chunked context (each slice anchored at its chunk start) → correct overlap even when the slice is entirely before dst_token_start """ overlap_start = max(src_token_start, dst_token_start) @@ -811,47 +785,15 @@ def _align_kv_blocks( dst_block_ids[dst_skip : dst_skip + n_transfer], ) - @staticmethod - def _beam0_block_count(block_ids: np.ndarray, total_blocks: int, beam_width: int) -> int: - """Return the number of beam-0 blocks in a packed 1-D beam layout.""" - if beam_width <= 1 or block_ids.size <= total_blocks: - return block_ids.size - return max(0, block_ids.size - (beam_width - 1)) - - @staticmethod - def _trim_receiver_window_head( - src_block_ids: np.ndarray, - dst_block_ids: np.ndarray, - peer_window_size: Optional[int], - beam_width: int, - ) -> np.ndarray: - """Drop the receiver's extra leading blocks for a windowed layer group. - - A windowed receiver keeps a larger window when only it runs speculative - decoding, so its suffix starts earlier and the extra blocks are at the - head. Both starts are derived from list length, so trimming the tail - instead would shift every block one position early. - - Non-windowed lists are both trimmed to ceil(prompt_len / tpb) in - _create_kv_slice, so there dst must not exceed src. A smaller dst - (generation prefix-cache reuse) is handled via dst_start. - """ - block_diff = dst_block_ids.size - src_block_ids.size - if block_diff <= 0: - return dst_block_ids - if peer_window_size is None or beam_width > 1: - raise ValueError( - f"src/dst block count mismatch: {src_block_ids.size} vs " - f"{dst_block_ids.size} (dst must not exceed src)" - ) - return dst_block_ids[block_diff:] - @nvtx_range("_build_kv_write_meta") def _build_kv_write_meta(self, task: KVSendTask, req_info: RecvReqInfo) -> WriteMeta: peer_ri = self._registrar.get_peer_rank_info(req_info.instance_name, req_info.instance_rank) self_ri = self._registrar.self_rank_info - token_range = task._slice.token_range - if token_range is not None and (self_ri.cp_size > 1 or peer_ri.cp_size > 1): + # A partial chunk (any non-final slice, or a final slice after earlier + # ones) addresses a global block sub-range, which the helix striding + # below cannot represent. + is_partial_chunk = task.slice_id > 0 or not task._slice.is_last_slice + if is_partial_chunk and (self_ri.cp_size > 1 or peer_ri.cp_size > 1): raise ValueError( "enable_pipelined_transfer is not supported with context parallelism " f"(sender cp_size={self_ri.cp_size}, receiver cp_size={peer_ri.cp_size})" @@ -902,74 +844,53 @@ def _build_kv_write_meta(self, task: KVSendTask, req_info: RecvReqInfo) -> Write dst_region = peer_extractor.extract_slot(int(dst_block_ids[0]), peer_lg, peer_pi) else: tpb = extractor.page_table.tokens_per_block + src_first = task._slice.first_ordinals[self_lg] + dst_first = req_info.first_ordinals[peer_lg] if peer_ri.cp_size > 1 and self_ri.cp_size == 1: # Helix: the receiver owns global blocks [cp_rank::cp_size] - # (same protocol as partition_context_for_helix). The strided - # subset has exactly the receiver's block count, so the - # suffix alignment below degenerates to identity; block - # reuse is rejected under helix. + # (same protocol as partition_context_for_helix). Block + # reuse and SWA trims are rejected under helix, so both + # anchors must be 0; the strided subset is then in the + # receiver's local ordinal space and aligns identically. + if src_first != 0 or dst_first != 0: + raise ValueError( + f"helix CP requires zero block anchors, got src={src_first} " + f"dst={dst_first} for unique_rid={task._unique_rid}" + ) src_block_ids = src_block_ids[peer_ri.cp_rank :: peer_ri.cp_size] - window_size = getattr(lg_info, "sliding_window_size", None) - # Block lists are the suffix of [..., slice_end); cached prefix - # is implicit in their size. token_start = (total_blocks - n) * tpb. - slice_end = token_range.end if token_range is not None else task._prompt_len - total_blocks = (slice_end + tpb - 1) // tpb - - # Project the peer's whole-prompt list only for partial chunks. - # A final SWA slice carries the complete active window rather than - # only the last context chunk, so its peer list must remain whole. - prompt_blocks = (task._prompt_len + tpb - 1) // tpb - is_windowed = window_size is not None and window_size < task._prompt_len - if ( - token_range is not None - and (token_range.start > 0 or total_blocks < prompt_blocks) - and not (is_windowed and task._slice.is_last_slice) - ): - dst_block_ids = project_blocks_to_global_chunk( - dst_block_ids, - chunk_block_offset=token_range.start // tpb, - chunk_block_count=total_blocks - token_range.start // tpb, - resident_block_end=prompt_blocks, - ) - - peer_lg_info = peer_extractor.page_table.layer_groups[peer_lg] - dst_block_ids = Sender._trim_receiver_window_head( - src_block_ids, - dst_block_ids, - peer_window_size=getattr(peer_lg_info, "sliding_window_size", None), - beam_width=task._beam_width, + assert task._prompt_len is not None, ( + "paged KV transfer requires session.prompt_len; " + "set TxSession(prompt_len=request.prompt_len)." ) - src_beam0 = Sender._beam0_block_count(src_block_ids, total_blocks, task._beam_width) - dst_beam0 = Sender._beam0_block_count(dst_block_ids, total_blocks, task._beam_width) - assert src_beam0 <= total_blocks, ( - f"src beam-0 block list ({src_beam0}) exceeds total slice " - f"blocks ({total_blocks}); slice_end={slice_end}, tpb={tpb}" + prompt_blocks = (task._prompt_len + tpb - 1) // tpb + beam_width = task._beam_width + # Anchors make the beam-0 length known (prompt_blocks − first + # ordinal), so tail counts are exact, not beam_width guesses. + src_beam0, src_tail = split_packed_beam_block_ids( + src_block_ids, beam_width, max(0, prompt_blocks - src_first) ) - assert dst_beam0 <= total_blocks, ( - f"dst beam-0 block list ({dst_beam0}) exceeds total slice " - f"blocks ({total_blocks}); slice_end={slice_end}, tpb={tpb}" + dst_beam0, dst_tail = split_packed_beam_block_ids( + dst_block_ids, beam_width, max(0, prompt_blocks - dst_first) ) - src_start = (total_blocks - src_beam0) * tpb - dst_start = (total_blocks - dst_beam0) * tpb - if req_info.dst_start_token is not None: - dst_start = max(dst_start, req_info.dst_start_token) - if window_size is not None: - # SWA eviction is based on the full prompt, not this slice. - assert task._prompt_len is not None, ( - "SWA layer requires session.prompt_len; " - "set TxSession(prompt_len=request.prompt_len)." - ) - stale_end = max(0, (task._prompt_len + 1 - window_size) // tpb) - src_start = max(stale_end * tpb, src_start) - dst_start = max(stale_end * tpb, dst_start) + # Pure interval intersection of the two anchored beam-0 spans. src_block_ids, dst_block_ids = Sender._align_kv_blocks( - src_block_ids, - dst_block_ids, - src_token_start=src_start, - dst_token_start=dst_start, + src_beam0, + dst_beam0, + src_token_start=src_first * tpb, + dst_token_start=dst_first * tpb, tokens_per_block=tpb, ) + if src_block_ids.size > 0 and (src_tail.size > 0 or dst_tail.size > 0): + # Beam tails pair up positionally after the beam-0 spans. + if src_tail.size != dst_tail.size: + raise ValueError( + f"packed beam-tail count mismatch for unique_rid=" + f"{task._unique_rid}: src has {src_tail.size} tail " + f"block(s), dst has {dst_tail.size}" + ) + src_block_ids = np.concatenate([src_block_ids, src_tail]) + dst_block_ids = np.concatenate([dst_block_ids, dst_tail]) src_region = extractor.extract(src_block_ids, self_lg, self_pi) dst_region = peer_extractor.extract(dst_block_ids, peer_lg, peer_pi) @@ -1727,14 +1648,22 @@ def _build_recv_req_info(self, task: KVRecvTask) -> RecvReqInfo: "both ctx_request_id and disagg_request_id are None for task " f"unique_rid={task._unique_rid}" ) - # Receiver's cached prefix is implicit in block_ids size; sender derives dst_start. + kv_slice = task._kv_slice + if len(kv_slice.first_ordinals) != len(kv_slice.block_ids_per_layer_groups): + # Not an assert: must survive python -O so an unanchored slice never + # reaches the wire. + raise ValueError( + f"KVSlice for unique_rid={task._unique_rid} carries " + f"{len(kv_slice.first_ordinals)} first_ordinals for " + f"{len(kv_slice.block_ids_per_layer_groups)} layer groups" + ) return RecvReqInfo( sender_req_id=sender_req_id, instance_name=self_ri.instance_name, instance_rank=self_ri.instance_rank, - block_ids_per_layer_groups=task._kv_slice.block_ids_per_layer_groups, + block_ids_per_layer_groups=kv_slice.block_ids_per_layer_groups, unique_rid=task._unique_rid, - dst_start_token=None, + first_ordinals=kv_slice.first_ordinals, aux_slot=task._aux_slot, slice_id=task.slice_id, ) diff --git a/tensorrt_llm/_torch/disaggregation/resource/cache_reuse.py b/tensorrt_llm/_torch/disaggregation/resource/cache_reuse.py index 142161093f12..cf3734e2ef2c 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/cache_reuse.py +++ b/tensorrt_llm/_torch/disaggregation/resource/cache_reuse.py @@ -22,11 +22,32 @@ from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager +from tensorrt_llm.runtime.kv_cache_manager_v2 import BAD_PAGE_INDEX from .page import AttentionLayerGroup from .utils import get_global_layer_ids +def split_packed_beam_block_ids( + block_ids: np.ndarray, + beam_width: int, + beam0_len: int, +) -> tuple[np.ndarray, np.ndarray]: + """Split 1-D block IDs into beam-0 prefix and appended beam-tail blocks. + + ``beam0_len`` is the *known* beam-0 span length — with explicit anchors it + is span_end − first_ordinal (e.g. prompt_blocks − first_ordinal for a + whole-prompt span). The packer appends only the UNSHARED final block of + each non-zero beam (``resource_manager._pack_beam_cache_indices``), so a + list carries up to ``beam_width − 1`` tails but possibly fewer; the tail + count must therefore be derived as ``size − beam0_len``, never guessed + from ``beam_width``. ``beam_width <= 1`` lists never carry tails. + """ + if beam_width <= 1 or block_ids.size <= beam0_len: + return block_ids, block_ids[:0] + return block_ids[:beam0_len], block_ids[beam0_len:] + + class CacheReuseAdapter(ABC): """Uniform prefix-reuse API over KVCacheManager V1/V2.""" @@ -58,20 +79,29 @@ def get_cached_token_count_per_layer_group( return [scalar] * len(layer_groups) @abstractmethod - def get_block_ids( + def get_transfer_span( self, req: LlmRequest, group_idx: int, lg: AttentionLayerGroup, - ) -> np.ndarray: - """Per-layer-group block identifiers for *req* (dtype ``int64``). - - Returned values are **primary memory-pool slot indices**, not raw block IDs: - ``KVRegionExtractorV1.extract`` and downstream transfer code do - ``base_ptr + slot_idx * slot_bytes`` and require the value to be a current - primary-pool offset. With host offload enabled, a block's logical ID can - diverge from its primary slot index after offload/onboard, so each backend - must translate before returning. + ) -> tuple[np.ndarray, int]: + """Anchored transferable block span for one attention layer group. + + Returns ``(pages, first_block_ordinal)``: ``pages`` is the beam-0 block + list (dtype ``int64``) with manager-evicted/stale prefixes and + speculative scratch handled, followed by packed beam-tail blocks when + ``beam_width > 1``; ``first_block_ordinal`` is the 0-based sequence + block ordinal of ``pages[0]`` (beam tails stay outside the anchored + region). Empty spans use anchor 0. The anchor reflects manager facts + only; transfer-policy trims (SWA bandwidth trim, gen-side reuse skip) + belong to the caller, which must advance the anchor with any head-slice. + + Returned values are **primary memory-pool slot indices**, not raw block + IDs: ``KVRegionExtractorV1.extract`` and downstream transfer code do + ``base_ptr + slot_idx * slot_bytes`` and require the value to be a + current primary-pool offset. With host offload enabled, a block's + logical ID can diverge from its primary slot index after + offload/onboard, so each backend must translate before returning. """ @abstractmethod @@ -102,26 +132,78 @@ def _global_cached_token_count(self, req: LlmRequest) -> int: tpb = self.tokens_per_block return (req.prepopulated_prompt_len // tpb) * tpb - def get_block_ids(self, req, group_idx, lg): # noqa: ARG002 + def _to_pool_indices(self, block_ids: np.ndarray, window_size: int) -> np.ndarray: + """Translate logical block IDs to primary-pool slot indices. + + block_id != primary-pool slot index once host offload kicks in; translate + so the cache transceiver's pointer arithmetic is correct. The manager aborts + if any referenced block is currently offloaded — disagg transfer cannot read + from the secondary pool, and a held block can never be offloaded. + """ + if block_ids.size == 0: + return np.array([], dtype=np.int64) + pool_indices = self._mgr.get_memory_pool_block_indices( + block_ids.tolist(), window_size=window_size + ) + return np.asarray(pool_indices, dtype=np.int64) + + def get_transfer_span(self, req, group_idx, lg): # noqa: ARG002 + window_size = lg.sliding_window_size + # V1 layer groups carry the manager's window key (full-attention layers get the + # max window), so this is always set; see kv_extractor.build_page_table. + assert window_size is not None first_layer = get_global_layer_ids(lg)[0] beam_width = req.py_beam_width raw_ids = self._mgr.get_batch_cache_indices( [req.py_request_id], layer_idx=first_layer, beam_width=beam_width )[0] if not raw_ids: - return np.array([], dtype=np.int64) - # block_id != primary-pool slot index once host offload kicks in; translate - # so the cache transceiver's pointer arithmetic is correct. The manager aborts - # if any referenced block is currently offloaded — disagg transfer cannot read - # from the secondary pool, and a held block can never be offloaded. - window_size = lg.sliding_window_size - # V1 layer groups carry the manager's window key (full-attention layers get the - # max window), so this is always set; see kv_extractor.build_page_table. - assert window_size is not None - pool_indices = self._mgr.get_memory_pool_block_indices( - list(raw_ids), window_size=window_size - ) - return np.asarray(pool_indices, dtype=np.int64) + return np.array([], dtype=np.int64), 0 + block_ids = np.asarray(raw_ids, dtype=np.int64) + if self._mgr.mapping.cp_size > 1: + # Helix CP: the list holds this rank's strided local blocks, so + # ordinals are local and scratch/eviction bookkeeping (which is + # global) does not apply. + return self._to_pool_indices(block_ids, window_size), 0 + + tpb = self.tokens_per_block + prompt_blocks = (req.prompt_len + tpb - 1) // tpb + allocated_blocks = (req.prompt_len + self._mgr.num_extra_kv_tokens + tpb - 1) // tpb + beam0, tails = split_packed_beam_block_ids(block_ids, beam_width, allocated_blocks) + # Draft-token allocation can extend past the speculative scratch bound. + if beam0.size > allocated_blocks: + beam0 = beam0[:allocated_blocks] + # Only prompt KV is transferred; drop the speculative scratch tail. + scratch_blocks = allocated_blocks - prompt_blocks + if scratch_blocks > 0: + if beam_width != 1: + raise ValueError("speculative scratch blocks require beam_width == 1") + beam0 = beam0[:-scratch_blocks] if scratch_blocks < beam0.size else beam0[:0] + + # The anchor reflects what the manager actually evicted. Detached front + # blocks remain in the C++ cache-block-id list (detachFrontBlock only + # advances a counter), so the leading `anchor` entries are dangling ids + # and must be dropped before pool translation — a detached block may + # have been reused or offloaded by now. + anchor = self._mgr.get_num_front_blocks_removed(req.py_request_id, window_size) + if anchor > 0: + # detachFrontBlock asserts beam_width == 1 in C++. + assert beam_width == 1, "front-block eviction requires beam_width == 1" + beam0 = beam0[anchor:] if anchor < beam0.size else beam0[:0] + + expected = max(0, prompt_blocks - anchor) + if beam0.size != expected: + raise RuntimeError( + f"request {req.py_request_id} window={window_size}: beam-0 block list " + f"holds {beam0.size} blocks after stripping {anchor} evicted front " + f"block(s) and {max(0, scratch_blocks)} scratch tail block(s), expected " + f"{expected} (= ceil(prompt_len={req.prompt_len} / tokens_per_block=" + f"{tpb}) - {anchor}); refusing to transfer misaligned KV blocks" + ) + if beam0.size == 0: + return np.array([], dtype=np.int64), 0 + pages = np.concatenate([beam0, tails]) if tails.size > 0 else beam0 + return self._to_pool_indices(pages, window_size), anchor def commit_blocks_for_reuse(self, req: LlmRequest) -> None: if not self.enable_block_reuse: @@ -152,18 +234,86 @@ def _global_cached_token_count(self, req: LlmRequest) -> int: tpb = self.tokens_per_block return (kv_cache.num_committed_tokens // tpb) * tpb - def get_block_ids(self, req, group_idx, lg): # noqa: ARG002 + def get_transfer_span(self, req, group_idx, lg): + kv_cache = self._mgr.kv_cache_map[req.py_request_id] # V2 already returns per-cache-level pool slot indices (not logical block # IDs), and active sequences GPU-lock their pages (_UniqPageLock enforces # cache_level==GPU), so the slot_ids yielded here are already the right # offsets for primary-pool pointer arithmetic. No translation is needed, - # unlike V1 (see _CacheReuseAdapterV1.get_block_ids). - return np.fromiter( - self._mgr.kv_cache_map[req.py_request_id].get_aggregated_page_indices( - group_idx, valid_only=True - ), + # unlike V1 (see _CacheReuseAdapterV1._to_pool_indices). + # + # valid_only=False keeps one entry per block ordinal (BAD_PAGE_INDEX for + # unbacked ordinals), so index == ordinal and the anchor can be derived + # instead of guessed from the list length. + pages = np.fromiter( + kv_cache.get_aggregated_page_indices(group_idx, valid_only=False), dtype=np.int64, ) + if self._mgr.mapping.cp_size > 1: + # Helix CP: the list holds this rank's strided local blocks, so the + # global ordinal math below (prompt cap, stale/scratch ranges) does + # not apply. Mirrors _CacheReuseAdapterV1's helix pass-through. + if bool((pages == BAD_PAGE_INDEX).any()): + raise RuntimeError( + f"request {req.py_request_id} layer group {group_idx}: unbacked " + "block ordinals under helix CP cannot be anchored (local block " + "lists have no global stale/scratch bookkeeping)" + ) + return pages, 0 + tpb = self.tokens_per_block + prompt_blocks = (req.prompt_len + tpb - 1) // tpb + # Blocks past the prompt (speculative scratch tail) are never transferred. + pages = pages[: min(pages.size, prompt_blocks)] + invalid = np.flatnonzero(pages == BAD_PAGE_INDEX) + # The transferable span is the contiguous backed run ending at the last + # prompt block; ordinals at or before an unbacked one hold no readable KV + # (stale-unlocked, or written to rotating shared scratch slots). + anchor = int(invalid[-1]) + 1 if invalid.size > 0 else 0 + span = pages[anchor:] + if span.size == 0: + return np.array([], dtype=np.int64), 0 + if invalid.size > 0: + # A single anchor cannot represent a backed sink prefix followed by + # an unbacked hole — it would silently drop the sink blocks, which + # are required KV (unlike a held-for-commit stale prefix, which is + # droppable). PyExecutor never enables sinks (its only + # AttentionLayerConfig construction passes num_sink_tokens=None), + # so fail loud if a sink-configured life cycle ever gets here. + # Best-effort guard: only the Python reference backend exposes + # _life_cycles; the default C++ backend + # (TLLM_KV_CACHE_MANAGER_V2_BACKEND=cpp) has no public life-cycle + # accessor, so the check is skipped there. + life_cycles = getattr(kv_cache.manager, "_life_cycles", None) + if life_cycles is not None: + life_cycle = life_cycles[group_idx] + if getattr(life_cycle, "num_sink_blocks", 0): + raise RuntimeError( + f"request {req.py_request_id} layer group {group_idx}: KV " + "transfer of anchored spans does not support token sinks " + f"(num_sink_blocks={life_cycle.num_sink_blocks})" + ) + # Consistency check against manager truths: every unbacked ordinal + # must be explained by SWA staleness or scratch placement. + window_size = lg.sliding_window_size + stale_end = ( + max(0, (kv_cache.history_length + 1 - window_size) // tpb) + if window_size is not None + else 0 + ) + explained = invalid < stale_end + scratch = kv_cache.get_scratch_desc(group_idx) + if scratch is not None: + beg, end = scratch.range + explained |= (invalid >= beg) & (invalid < end) + if not bool(explained.all()): + raise RuntimeError( + f"request {req.py_request_id} layer group {group_idx}: unbacked " + f"block ordinals {invalid[~explained].tolist()} are neither SWA-" + f"stale (stale_end={stale_end}) nor scratch " + f"({scratch.range if scratch is not None else None}); refusing " + "to transfer misaligned KV blocks" + ) + return span, anchor def commit_blocks_for_reuse(self, req: LlmRequest) -> None: self._mgr.try_commit_blocks(req) diff --git a/tensorrt_llm/_torch/disaggregation/transceiver.py b/tensorrt_llm/_torch/disaggregation/transceiver.py index c8bea047f730..455ec73ccb1d 100644 --- a/tensorrt_llm/_torch/disaggregation/transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/transceiver.py @@ -28,7 +28,6 @@ KVSlice, RxSessionBase, SessionStatus, - TokenRange, TxSessionBase, WaitResult, get_unique_rid, @@ -41,6 +40,7 @@ from tensorrt_llm._torch.disaggregation.resource.cache_reuse import ( CacheReuseAdapter, create_cache_reuse_adapter, + split_packed_beam_block_ids, ) from tensorrt_llm._torch.disaggregation.resource.page import CacheKind from tensorrt_llm._torch.disaggregation.resource.utils import get_physical_pool @@ -285,8 +285,31 @@ def _get_mamba_slot_for_request(self, req: LlmRequest) -> Optional[int]: return self._kv_cache_manager.mamba_cache_index[req.py_request_id] return None - def _create_kv_slice(self, req: LlmRequest) -> KVSlice: - """Create a KV slice covering the request's whole prompt.""" + @staticmethod + def _defers_to_final_slice(lg, prompt_len: int) -> bool: + """Whether a layer group's SWA window forces whole-window transfer. + + Pipelined senders skip such groups on intermediate chunks and send the + complete final active window with the last slice. Single source for the + predicate so producer and chunker can never disagree on a group. + """ + window_size = getattr(lg, "sliding_window_size", None) + return window_size is not None and window_size < prompt_len + + def _create_kv_slice(self, req: LlmRequest, include_window_groups: bool = True) -> KVSlice: + """Create a KV slice covering the request's whole prompt. + + Each layer group carries an explicit block-ordinal anchor + (``KVSlice.first_ordinals``). The adapter reports the manager-truth + span; transfer-policy trims (SWA pre-window skip, gen-side reuse skip) + are applied here as explicit head-slices that advance the anchor. + + Args: + req: The request whose KV blocks to describe. + include_window_groups: When False, layer groups whose sliding + window is smaller than the prompt get empty lists. Pipelined + senders defer those groups to the final chunk. + """ adapter = self._reuse_adapter tpb = adapter.tokens_per_block assert self._page_table is not None @@ -307,8 +330,12 @@ def _create_kv_slice(self, req: LlmRequest) -> KVSlice: if is_gen_only else [0] * len(layer_groups) ) + # Helix CP block lists are strided local subsets of the global blocks; + # anchors are local (0) and global-position trims do not apply. + is_helix = self._mapping.cp_size > 1 - groups = [] + groups: List[np.ndarray] = [] + first_ordinals: List[int] = [] for idx, lg in enumerate(layer_groups): if lg.kind == CacheKind.STATE: slot = self._get_mamba_slot_for_request(req) @@ -317,66 +344,47 @@ def _create_kv_slice(self, req: LlmRequest) -> KVSlice: if slot is not None else np.array([], dtype=np.int64) ) + first_ordinals.append(0) continue - block_ids = adapter.get_block_ids(req, idx, lg) window_size = lg.sliding_window_size + if not include_window_groups and self._defers_to_final_slice(lg, req.prompt_len): + groups.append(np.array([], dtype=np.int64)) + first_ordinals.append(0) + continue - if window_size is not None: - allocated_blocks = ( - req.prompt_len + self._kv_cache_manager.num_extra_kv_tokens + tpb - 1 - ) // tpb - beam0_block_ids, tail_block_ids = self._split_packed_beam_block_ids( - block_ids, - req.py_beam_width, - allocated_blocks, + pages, anchor = adapter.get_transfer_span(req, idx, lg) + + # Transfer-policy head-slice: skip pre-window prompt blocks (the + # receiver never reads them; bandwidth only) and the gen-side + # reused prefix. The anchor advances with the slice so block + # positions stay explicit. + target_start = anchor + if window_size is not None and not is_helix: + target_start = max(target_start, (req.prompt_len + 1 - window_size) // tpb) + if is_gen_only: + target_start = max(target_start, cached_per_lg[idx] // tpb) + skip = target_start - anchor + if skip > 0 and pages.size > 0: + beam0, tails = split_packed_beam_block_ids( + pages, req.py_beam_width, prompt_blocks - anchor ) - if beam0_block_ids.size > allocated_blocks: - beam0_block_ids = beam0_block_ids[:allocated_blocks] - block_ids = ( - np.concatenate([beam0_block_ids, tail_block_ids]) - if tail_block_ids.size > 0 - else beam0_block_ids - ) - # Current PyExecutor cache managers disable KV-cache token sinks, - # so SWA block lists contain an evictable prompt prefix followed - # by the speculative scratch tail. If token sinks are enabled, - # this must use block-ordinal metadata to preserve the sink prefix. - # Remove scratch before trimming stale prompt blocks; otherwise a - # boundary-crossing allocation can displace initialized prompt KV. - scratch_blocks = max(0, allocated_blocks - prompt_blocks) - if scratch_blocks > 0: - if req.py_beam_width != 1: - raise ValueError("speculative scratch blocks require beam_width == 1") - block_ids = ( - block_ids[:-scratch_blocks] - if scratch_blocks < block_ids.size - else np.array([], dtype=np.int64) - ) - # Drop stale blocks the manager may still expose (V1 pre-eviction). - stale_end = max(0, (req.prompt_len + 1 - window_size) // tpb) - expected_valid = max(0, prompt_blocks - stale_end) - # Skip reused blocks that remain after stale-prefix pruning. - cache_skip = max(0, cached_per_lg[idx] // tpb - stale_end) - else: - # Drop the speculative scratch tail; only prompt_len is transferred. - if block_ids.size > prompt_blocks: - block_ids = block_ids[:prompt_blocks] - expected_valid = prompt_blocks - cache_skip = cached_per_lg[idx] // tpb - - block_ids = self._trim_packed_beam_block_ids( - block_ids, - beam_width=req.py_beam_width, - total_blocks=prompt_blocks, - expected_valid=expected_valid, - cache_skip=cache_skip, - ) + beam0 = beam0[skip:] if skip < beam0.size else beam0[:0] + if beam0.size == 0: + pages = np.array([], dtype=np.int64) + anchor = 0 + else: + pages = np.concatenate([beam0, tails]) if tails.size > 0 else beam0 + anchor = target_start + elif pages.size == 0: + anchor = 0 - groups.append(block_ids) + groups.append(pages) + first_ordinals.append(anchor) return KVSlice( is_last_slice=True, block_ids_per_layer_groups=groups, + first_ordinals=first_ordinals, ) def _slice_num_bytes(self, slice: KVSlice) -> int: @@ -410,56 +418,6 @@ def _slice_num_bytes(self, slice: KVSlice) -> int: total += n * pool.slot_bytes return total - @staticmethod - def _split_packed_beam_block_ids( - block_ids: np.ndarray, - beam_width: int, - total_blocks: int, - ) -> tuple[np.ndarray, np.ndarray]: - """Split 1-D block IDs into beam-0 prefix and appended beam-tail blocks.""" - if beam_width <= 1 or block_ids.size <= total_blocks: - return block_ids, np.array([], dtype=np.int64) - tail_count = min(beam_width - 1, block_ids.size - total_blocks) - if tail_count <= 0: - return block_ids, np.array([], dtype=np.int64) - return block_ids[:-tail_count], block_ids[-tail_count:] - - @staticmethod - def _trim_packed_beam_block_ids( - block_ids: np.ndarray, - beam_width: int, - total_blocks: int, - expected_valid: int, - cache_skip: int, - ) -> np.ndarray: - """Trim/skip beam-0 blocks while preserving packed beam-tail blocks.""" - if expected_valid <= 0: - return np.array([], dtype=np.int64) - - beam0_block_ids, tail_block_ids = KvCacheTransceiverV2._split_packed_beam_block_ids( - block_ids, beam_width, total_blocks - ) - - if beam0_block_ids.size > expected_valid: - beam0_block_ids = ( - beam0_block_ids[-expected_valid:] - if expected_valid > 0 - else np.array([], dtype=np.int64) - ) - if beam0_block_ids.size == 0: - tail_block_ids = np.array([], dtype=np.int64) - - if cache_skip > 0: - if cache_skip < beam0_block_ids.size: - beam0_block_ids = beam0_block_ids[cache_skip:] - else: - beam0_block_ids = np.array([], dtype=np.int64) - tail_block_ids = np.array([], dtype=np.int64) - - if tail_block_ids.size == 0: - return beam0_block_ids - return np.concatenate([beam0_block_ids, tail_block_ids]) - @staticmethod def _need_aux_transfer(req: LlmRequest) -> bool: params = req.py_disaggregated_params @@ -761,32 +719,46 @@ def _build_prefill_chunk( if chunk_end <= chunk_start and not is_last_chunk: return None - base_slice = self._create_kv_slice(req) - all_block_ids = base_slice.block_ids_per_layer_groups + base_slice = self._create_kv_slice(req, include_window_groups=is_last_chunk) chunk_block_ids = [] + chunk_first_ordinals = [] assert self._page_table is not None - for lg, block_ids in zip(self._page_table.layer_groups, all_block_ids): - window_size = getattr(lg, "sliding_window_size", None) - if window_size is not None and window_size < req.prompt_len: + for lg, block_ids, anchor in zip( + self._page_table.layer_groups, + base_slice.block_ids_per_layer_groups, + base_slice.first_ordinals, + ): + if self._defers_to_final_slice(lg, req.prompt_len): # SWA pages can leave the active window between chunks. Defer - # the group and send its complete final active window at once. + # the group and send its complete final active window at once, + # anchored at its true window-start ordinal. chunk_block_ids.append(block_ids if is_last_chunk else block_ids[:0]) + chunk_first_ordinals.append(anchor if is_last_chunk else 0) else: - # _build_kv_write_meta derives the chunk's start token from list length - # (total_blocks - len), which only agrees with the slice below - # once the group covers the chunk end. - assert block_ids.size >= chunk_end, ( - f"layer group holds {block_ids.size} blocks, fewer than the " - f"chunk end {chunk_end}; cannot address chunk " - f"[{chunk_start}, {chunk_end}) by position" + if anchor + block_ids.size < chunk_end: + # Not an assert: must fail under python -O too, or a short + # span would silently transfer misaligned chunk blocks. + raise ValueError( + f"layer group span [{anchor}, {anchor + block_ids.size}) does " + f"not cover the chunk end {chunk_end}; cannot address chunk " + f"[{chunk_start}, {chunk_end}) by block ordinal" + ) + # Intersect the chunk's global block range with the group span; + # a span starting past the chunk end contributes nothing. + start = max(chunk_start, anchor) + sub = ( + block_ids[start - anchor : chunk_end - anchor] + if start < chunk_end + else block_ids[:0] ) - chunk_block_ids.append(block_ids[chunk_start:chunk_end]) + chunk_block_ids.append(sub) + chunk_first_ordinals.append(start if sub.size > 0 else 0) if not is_last_chunk and not any(block_ids.size for block_ids in chunk_block_ids): return None return KVSlice( is_last_slice=is_last_chunk, block_ids_per_layer_groups=chunk_block_ids, - token_range=TokenRange(start=chunk_start * tpb, end=chunk_end * tpb), + first_ordinals=chunk_first_ordinals, ) @nvtx_range("KvCacheTransceiverV2.respond_and_send_async") diff --git a/tests/integration/defs/disaggregated/test_disaggregated.py b/tests/integration/defs/disaggregated/test_disaggregated.py index 106bc9d6790f..6fbd6254f7e3 100644 --- a/tests/integration/defs/disaggregated/test_disaggregated.py +++ b/tests/integration/defs/disaggregated/test_disaggregated.py @@ -1498,7 +1498,7 @@ def _verify_python_transceiver_under_host_offload(server_url: str, model: str): """End-to-end check: Python transceiver + ctx-side host offload. The fix translates logical block IDs to primary-pool slot indices in - `_CacheReuseAdapterV1.get_block_ids` before they reach the disagg + `_CacheReuseAdapterV1.get_transfer_span` before they reach the disagg sender. Without that translation, once host offload moves blocks around, the sender computes pool pointers from stale block IDs and either reads garbage memory or aborts. This test stresses that path diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index 2f0a0efd5b34..569c88b9c217 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -900,6 +900,7 @@ def test_v2_disagg_slice_skips_state_index_on_mamba_free_pp_rank(): transceiver._kv_cache_manager = manager transceiver._reuse_adapter = SimpleNamespace(tokens_per_block=32) transceiver._page_table = SimpleNamespace(layer_groups=[]) + transceiver._mapping = SimpleNamespace(cp_size=1) request = SimpleNamespace( is_generation_only_request=lambda: False, prompt_len=0, @@ -927,6 +928,7 @@ def test_v2_disagg_slice_reads_state_index_without_refreshing_batch_mask(): transceiver._kv_cache_manager = manager transceiver._reuse_adapter = SimpleNamespace(tokens_per_block=32) transceiver._page_table = SimpleNamespace(layer_groups=[mamba_lg]) + transceiver._mapping = SimpleNamespace(cp_size=1) request = SimpleNamespace( is_generation_only_request=lambda: False, prompt_len=0, diff --git a/tests/unittest/disaggregated/test_cache_reuse_adapter.py b/tests/unittest/disaggregated/test_cache_reuse_adapter.py index f8b1f773925f..83c2b7767d25 100644 --- a/tests/unittest/disaggregated/test_cache_reuse_adapter.py +++ b/tests/unittest/disaggregated/test_cache_reuse_adapter.py @@ -12,7 +12,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Tests for CacheReuseAdapter, _create_kv_slice SWA trim, and Sender token-start derivation.""" +"""Tests for adapter transfer spans, anchor trims, and anchored alignment.""" from types import SimpleNamespace from unittest.mock import MagicMock @@ -24,10 +24,13 @@ from tensorrt_llm._torch.disaggregation.resource.cache_reuse import ( CacheReuseAdapter, _CacheReuseAdapterV1, + _CacheReuseAdapterV2, + split_packed_beam_block_ids, ) from tensorrt_llm._torch.disaggregation.resource.page import AttentionLayerGroup, LocalLayer from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2 from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager +from tensorrt_llm.runtime.kv_cache_manager_v2 import BAD_PAGE_INDEX pytestmark = pytest.mark.cpu_only @@ -109,38 +112,81 @@ def test_dst_extra_draft_block(self): # --------------------------------------------------------------------------- +class _FakeReq: + def __init__(self, prompt_len: int, beam_width: int = 1, request_id: int = 0): + self.prompt_len = prompt_len + self.py_beam_width = beam_width + self.py_request_id = request_id + + +class _FakeSamplingConfig: + def __init__(self, beam_width: int): + self.beam_width = beam_width + + +class _FakeV1Mgr: + """C++-manager stand-in for _CacheReuseAdapterV1.get_transfer_span. + + Pool translation is identity but records what it was asked to translate, + so tests can assert dangling (evicted) or scratch block IDs never reach + the pool-pointer arithmetic. + """ + + enable_block_reuse = True + + def __init__( + self, + block_ids, + *, + tokens_per_block: int = 8, + num_extra_kv_tokens: int = 0, + front_blocks_removed: int = 0, + cp_size: int = 1, + ): + self.tokens_per_block = tokens_per_block + self.num_extra_kv_tokens = num_extra_kv_tokens + self.mapping = SimpleNamespace(cp_size=cp_size) + self._block_ids = list(block_ids) + self._front_blocks_removed = front_blocks_removed + self.requested_beam_width = None + self.translated_ids = None + self.translated_window = None + + def get_batch_cache_indices(self, request_ids, layer_idx=None, beam_width=1): + self.requested_beam_width = beam_width + return [list(self._block_ids)] + + def get_num_front_blocks_removed(self, request_id, window_size): + return self._front_blocks_removed + + def get_memory_pool_block_indices(self, block_ids, window_size): + self.translated_ids = list(block_ids) + self.translated_window = window_size + return block_ids + + +def _lg(window=None): + return AttentionLayerGroup( + pool_group_idx=0, + sliding_window_size=window, + local_layers=[LocalLayer(local_layer_id=0, global_layer_id=0)], + ) + + class TestPackedBeamBlockLayout: """Verify beam search block IDs stay 1-D with only final tail blocks appended.""" def test_v1_adapter_uses_request_py_beam_width(self): - class _FakeMgr: - enable_block_reuse = True - tokens_per_block = 32 - - def __init__(self): - self.beam_width = None - self.pool_indices_window = None - - def get_batch_cache_indices(self, request_ids, layer_idx=None, beam_width=1): - self.beam_width = beam_width - return [[10, 11, 12, 13]] - - def get_memory_pool_block_indices(self, block_ids, window_size): - # Identity translation: nothing offloaded, block_id == pool slot. - self.pool_indices_window = window_size - return block_ids - - req = _FakeReq(prompt_len=7) - req.py_request_id = 1 - req.py_beam_width = 4 + req = _FakeReq(prompt_len=7, beam_width=4, request_id=1) req.sampling_config = _FakeSamplingConfig(beam_width=1) - mgr = _FakeMgr() + mgr = _FakeV1Mgr([10, 11, 12, 13], tokens_per_block=32) - block_ids = _CacheReuseAdapterV1(mgr).get_block_ids(req, 0, _lg(window=512)) + pages, anchor = _CacheReuseAdapterV1(mgr).get_transfer_span(req, 0, _lg(window=512)) - assert mgr.beam_width == 4 - assert mgr.pool_indices_window == 512 - np.testing.assert_array_equal(block_ids, [10, 11, 12, 13]) + assert mgr.requested_beam_width == 4 + assert mgr.translated_window == 512 + assert anchor == 0 + np.testing.assert_array_equal(pages, [10, 11, 12, 13]) def test_pack_beam_cache_indices_single_block_prompt_keeps_all_beams(self): packed = KVCacheManager._pack_beam_cache_indices([[10], [10], [10], [10]]) @@ -170,232 +216,519 @@ def test_pack_beam_cache_indices_skips_shared_final_blocks(self): assert packed == [10, 11, 12, 13] - def test_beam0_block_count_for_full_packed_prompt(self): - block_ids = np.array([10, 11, 12, 13, 14, 15], dtype=np.int64) + def test_split_single_block_prompt_separates_beam_tails(self): + beam0, tails = split_packed_beam_block_ids( + np.array([10, 11, 12, 13], dtype=np.int64), beam_width=4, beam0_len=1 + ) + + np.testing.assert_array_equal(beam0, [10]) + np.testing.assert_array_equal(tails, [11, 12, 13]) + + def test_split_long_prompt_separates_beam_tails(self): + beam0, tails = split_packed_beam_block_ids( + np.array([10, 11, 12, 13, 14, 15], dtype=np.int64), beam_width=4, beam0_len=3 + ) + + np.testing.assert_array_equal(beam0, [10, 11, 12]) + np.testing.assert_array_equal(tails, [13, 14, 15]) + + def test_split_with_fewer_shared_tails_than_beam_width(self): + # The packer appends only UNSHARED final blocks, so a beam_width=4 list + # can carry fewer than 3 tails. The tail count is size - beam0_len, + # never a beam_width guess (which would misclassify beam-0 blocks + # as tails here). + beam0, tails = split_packed_beam_block_ids( + np.array([10, 11, 12, 14, 15], dtype=np.int64), beam_width=4, beam0_len=3 + ) + + np.testing.assert_array_equal(beam0, [10, 11, 12]) + np.testing.assert_array_equal(tails, [14, 15]) + + def test_split_anchored_span_beam0_len(self): + # beam0_len is span_end - first_ordinal: a 4-entry list whose beam-0 + # span holds 1 block (prompt 3 blocks, anchored at 2) has 3 tails. + beam0, tails = split_packed_beam_block_ids( + np.array([12, 13, 14, 15], dtype=np.int64), beam_width=4, beam0_len=1 + ) - assert Sender._beam0_block_count(block_ids, total_blocks=3, beam_width=4) == 3 + np.testing.assert_array_equal(beam0, [12]) + np.testing.assert_array_equal(tails, [13, 14, 15]) - def test_beam0_block_count_after_cached_prefix_skip(self): - block_ids = np.array([12, 13, 14, 15], dtype=np.int64) + def test_split_within_beam0_len_has_no_tails(self): + beam0, tails = split_packed_beam_block_ids( + np.array([10, 11], dtype=np.int64), beam_width=4, beam0_len=3 + ) - assert Sender._beam0_block_count(block_ids, total_blocks=3, beam_width=4) == 1 + np.testing.assert_array_equal(beam0, [10, 11]) + assert tails.size == 0 - def test_beam0_block_count_for_single_beam_unchanged(self): - block_ids = np.array([10, 11, 12], dtype=np.int64) + def test_split_single_beam_passthrough(self): + beam0, tails = split_packed_beam_block_ids( + np.array([10, 11, 12, 13], dtype=np.int64), beam_width=1, beam0_len=1 + ) - assert Sender._beam0_block_count(block_ids, total_blocks=3, beam_width=1) == 3 + np.testing.assert_array_equal(beam0, [10, 11, 12, 13]) + assert tails.size == 0 def test_align_packed_single_block_prompt_keeps_all_beam_blocks(self): + # Both spans anchored at 0 with a single beam-0 block each; tails ride + # behind the intersected beam-0 spans (see _build_kv_write_meta). src_block_ids = np.array([10, 10, 10, 10], dtype=np.int64) dst_block_ids = np.array([20, 21, 22, 23], dtype=np.int64) - total_blocks = 1 tpb = 32 - src_start = (total_blocks - Sender._beam0_block_count(src_block_ids, total_blocks, 4)) * tpb - dst_start = (total_blocks - Sender._beam0_block_count(dst_block_ids, total_blocks, 4)) * tpb + src_beam0, src_tail = split_packed_beam_block_ids(src_block_ids, beam_width=4, beam0_len=1) + dst_beam0, dst_tail = split_packed_beam_block_ids(dst_block_ids, beam_width=4, beam0_len=1) src, dst = Sender._align_kv_blocks( - src_block_ids, - dst_block_ids, - src_token_start=src_start, - dst_token_start=dst_start, + src_beam0, + dst_beam0, + src_token_start=0, + dst_token_start=0, tokens_per_block=tpb, ) + src = np.concatenate([src, src_tail]) + dst = np.concatenate([dst, dst_tail]) np.testing.assert_array_equal(src, [10, 10, 10, 10]) np.testing.assert_array_equal(dst, [20, 21, 22, 23]) - def test_trim_single_block_prompt_preserves_packed_beam_tails(self): - block_ids = np.array([10, 11, 12, 13], dtype=np.int64) - trimmed = KvCacheTransceiverV2._trim_packed_beam_block_ids( - block_ids, - beam_width=4, - total_blocks=1, - expected_valid=1, - cache_skip=0, +# --------------------------------------------------------------------------- +# Windowed layer group where only the generation side runs speculative +# decoding: the receiver keeps a larger window, so its span starts earlier. +# Anchors resolve the asymmetry by plain interval intersection — the deleted +# _trim_receiver_window_head count-based head trim has no successor mechanism +# beyond this. +# --------------------------------------------------------------------------- + + +class TestReceiverLargerWindowViaAnchors: + TPB = 128 + + def _align(self, src, src_anchor, dst, dst_anchor): + return Sender._align_kv_blocks( + np.array(src, dtype=np.int64), + np.array(dst, dtype=np.int64), + src_token_start=src_anchor * self.TPB, + dst_token_start=dst_anchor * self.TPB, + tokens_per_block=self.TPB, ) - np.testing.assert_array_equal(trimmed, [10, 11, 12, 13]) + def test_receiver_extra_head_block_gets_no_source(self): + src, dst = self._align([10], src_anchor=1224, dst=[20, 21], dst_anchor=1223) + + np.testing.assert_array_equal(src, [10]) + np.testing.assert_array_equal(dst, [21]) + + def test_last_prompt_block_maps_exactly(self): + # Regression (formerly against _trim_receiver_window_head): dropping the + # receiver's tail instead of its head paired dst block 20 with src block + # 10 one position early, so the last prompt block was never written. + # With anchors the pairing is by equal block ordinal by construction. + src, dst = self._align([10], src_anchor=1224, dst=[20, 21], dst_anchor=1223) - def test_trim_long_prompt_preserves_valid_beam0_and_tails(self): - block_ids = np.array([10, 11, 12, 13, 14, 15], dtype=np.int64) + assert (dst.tolist(), src.tolist()) == ([21], [10]) - trimmed = KvCacheTransceiverV2._trim_packed_beam_block_ids( - block_ids, - beam_width=4, - total_blocks=3, - expected_valid=3, - cache_skip=0, + def test_equal_windows_are_untouched(self): + src, dst = self._align([10, 11], src_anchor=1223, dst=[20, 21], dst_anchor=1223) + + np.testing.assert_array_equal(src, [10, 11]) + np.testing.assert_array_equal(dst, [20, 21]) + + def test_smaller_receiver_resolves_via_its_own_anchor(self): + # Generation prefix-cache reuse: dst declares a later-starting span. + src, dst = self._align([10, 11, 12], src_anchor=0, dst=[20], dst_anchor=2) + + np.testing.assert_array_equal(src, [12]) + np.testing.assert_array_equal(dst, [20]) + + +# --------------------------------------------------------------------------- +# _CacheReuseAdapterV1.get_transfer_span: manager-truth span and anchor. +# --------------------------------------------------------------------------- + + +class TestV1AdapterTransferSpan: + TPB = 8 + + def _span(self, mgr, *, prompt_len, beam_width=1, window=None): + req = _FakeReq(prompt_len=prompt_len, beam_width=beam_width) + window = window if window is not None else 1 << 30 # full attention + return _CacheReuseAdapterV1(mgr).get_transfer_span(req, 0, _lg(window=window)) + + def test_span_covers_ceil_prompt_blocks(self): + # prompt 17 + extra 7 = 24 tokens = 3 blocks allocated; prompt needs + # ceil(17/8) = 3, so nothing is scratch. + mgr = _FakeV1Mgr([0, 1, 2], num_extra_kv_tokens=7) + + pages, anchor = self._span(mgr, prompt_len=17) + + np.testing.assert_array_equal(pages, [0, 1, 2]) + assert anchor == 0 + + def test_extra_tokens_crossing_block_boundary_are_stripped(self): + # prompt 16 = 2 blocks; extra 7 pushes allocation to 3 blocks. The + # scratch block holds no prompt KV and must not be transferred. + mgr = _FakeV1Mgr([0, 1, 2], num_extra_kv_tokens=7) + + pages, anchor = self._span(mgr, prompt_len=16) + + np.testing.assert_array_equal(pages, [0, 1]) + assert anchor == 0 + assert mgr.translated_ids == [0, 1] + + def test_no_extra_defaults_to_prompt_len(self): + mgr = _FakeV1Mgr([0, 1, 2]) + + pages, anchor = self._span(mgr, prompt_len=17) + + np.testing.assert_array_equal(pages, [0, 1, 2]) + assert anchor == 0 + + def test_draft_over_allocation_is_capped(self): + # Draft-token allocation can extend the list past the allocated bound. + mgr = _FakeV1Mgr([100, 101, 102, 103, 104]) + + pages, anchor = self._span(mgr, prompt_len=32) + + np.testing.assert_array_equal(pages, [100, 101, 102, 103]) + assert anchor == 0 + + def test_scratch_requires_single_beam(self): + mgr = _FakeV1Mgr([100, 101, 102, 103, 104], num_extra_kv_tokens=2) + + with pytest.raises(ValueError, match="speculative scratch blocks require beam_width == 1"): + self._span(mgr, prompt_len=32, beam_width=4) + + def test_beam_tails_split_exactly_at_the_allocated_bound(self): + # 4 allocated beam-0 blocks + 3 packed tails: the tail count comes from + # size - allocated, so all tails survive. + mgr = _FakeV1Mgr([100, 101, 102, 103, 200, 201, 202]) + + pages, anchor = self._span(mgr, prompt_len=32, beam_width=4) + + np.testing.assert_array_equal(pages, [100, 101, 102, 103, 200, 201, 202]) + assert anchor == 0 + + def test_fewer_shared_tails_than_beam_width_split_correctly(self): + # The exact bug the beam0_len rename fixed: with beam_width=4 but only + # 2 UNSHARED tail blocks packed, the old beam_width-1 guess would have + # misclassified beam-0 block 103 as a tail. + mgr = _FakeV1Mgr([100, 101, 102, 103, 200, 201]) + + pages, anchor = self._span(mgr, prompt_len=32, beam_width=4) + + np.testing.assert_array_equal(pages, [100, 101, 102, 103, 200, 201]) + assert anchor == 0 + + def test_scratch_crossing_boundary_with_evicted_front_blocks(self): + """PR-17619 regression: boundary-crossing scratch with evicted fronts. + + prompt 32 = 4 blocks; extra 8 tokens allocate a 5th (scratch) block + holding garbage; the manager already detached the first 2 blocks, whose + IDs are dangling. The span must be honestly anchored at 2, exclude both + the dangling head and the garbage scratch tail, and never hand either + to pool translation. + """ + mgr = _FakeV1Mgr( + [100, 101, 102, 103, 104], # 104 = scratch garbage; 100, 101 = dangling + num_extra_kv_tokens=8, + front_blocks_removed=2, ) - np.testing.assert_array_equal(trimmed, [10, 11, 12, 13, 14, 15]) + pages, anchor = self._span(mgr, prompt_len=32, window=16) - def test_trim_swa_prefix_preserves_packed_beam_tails(self): - block_ids = np.array([10, 11, 12, 13, 14, 15, 16], dtype=np.int64) + np.testing.assert_array_equal(pages, [102, 103]) + assert anchor == 2 + # Dangling and scratch IDs are stripped BEFORE pool translation: a + # detached block may have been reused or offloaded by now. + assert mgr.translated_ids == [102, 103] - trimmed = KvCacheTransceiverV2._trim_packed_beam_block_ids( - block_ids, - beam_width=4, - total_blocks=4, - expected_valid=2, - cache_skip=0, + def test_eviction_counter_exceeding_the_stale_formula_is_honored(self): + """The eviction counter is authoritative over the stale formula. + + Extra/draft tokens can push eviction one block past the + (prompt_len + 1 - window) formula; the counter wins and the + dangling IDs are stripped before pool translation. + """ + tpb = 128 + prompt_len = 1150 # 9 blocks + mgr = _FakeV1Mgr( + list(range(200, 210)), # 9 prompt blocks + 1 scratch block + tokens_per_block=tpb, + num_extra_kv_tokens=5, + front_blocks_removed=8, # formula would give (1150+1-133)//128 = 7 ) + req = _FakeReq(prompt_len=prompt_len) - np.testing.assert_array_equal(trimmed, [12, 13, 14, 15, 16]) + pages, anchor = _CacheReuseAdapterV1(mgr).get_transfer_span(req, 0, _lg(window=133)) - def test_cache_skip_drops_tail_blocks_when_beam0_fully_cached(self): - block_ids = np.array([10, 11, 12, 13], dtype=np.int64) + np.testing.assert_array_equal(pages, [208]) + assert anchor == 8 + assert mgr.translated_ids == [208] - trimmed = KvCacheTransceiverV2._trim_packed_beam_block_ids( - block_ids, - beam_width=4, - total_blocks=1, - expected_valid=1, - cache_skip=1, + @pytest.mark.parametrize("prompt_len", (1150, 1151)) + def test_dspark_disagg_boundary_keeps_only_initialized_swa(self, prompt_len): + # 9 prompt blocks + 1 scratch block; 7 fronts already evicted. + tpb = 128 + mgr = _FakeV1Mgr( + list(range(200, 210)), + tokens_per_block=tpb, + num_extra_kv_tokens=5, + front_blocks_removed=7, ) + req = _FakeReq(prompt_len=prompt_len) + + pages, anchor = _CacheReuseAdapterV1(mgr).get_transfer_span(req, 0, _lg(window=133)) + + np.testing.assert_array_equal(pages, [207, 208]) + assert anchor == 7 + + def test_reconciliation_mismatch_raises(self): + """Anchor + len(beam0) != ceil(prompt_len / tpb) refuses to transfer.""" + mgr = _FakeV1Mgr([100, 101, 102]) # 3 blocks for a 4-block prompt + + with pytest.raises(RuntimeError, match="refusing to transfer misaligned KV blocks"): + self._span(mgr, prompt_len=32) + + def test_reconciliation_mismatch_after_eviction_raises(self): + # Counter says 1 evicted but the surviving list is still one short. + mgr = _FakeV1Mgr([100, 101, 102], front_blocks_removed=1) + + with pytest.raises(RuntimeError, match="refusing to transfer misaligned KV blocks"): + self._span(mgr, prompt_len=32) + + def test_everything_evicted_returns_empty_anchor_zero(self): + mgr = _FakeV1Mgr([100, 101, 102, 103], front_blocks_removed=4) - assert trimmed.size == 0 + pages, anchor = self._span(mgr, prompt_len=32) + + assert pages.size == 0 + assert anchor == 0 + + def test_empty_manager_list_returns_empty_anchor_zero(self): + mgr = _FakeV1Mgr([]) + + pages, anchor = self._span(mgr, prompt_len=32) + + assert pages.size == 0 + assert anchor == 0 + + def test_helix_cp_returns_local_ordinals_anchor_zero(self): + # Helix lists are strided local subsets; global eviction/scratch + # bookkeeping does not apply and the anchor stays 0. + mgr = _FakeV1Mgr([100, 102], num_extra_kv_tokens=8, front_blocks_removed=2, cp_size=2) + + pages, anchor = self._span(mgr, prompt_len=32) + + np.testing.assert_array_equal(pages, [100, 102]) + assert anchor == 0 # --------------------------------------------------------------------------- -# Windowed layer group where only the generation side runs speculative decoding. +# _CacheReuseAdapterV2.get_transfer_span: anchor derived from unbacked +# ordinals (valid_only=False keeps index == ordinal). # --------------------------------------------------------------------------- -class TestTrimReceiverWindowHead: - """Sender._trim_receiver_window_head drops the receiver's extra head blocks. +class _FakeKvCacheV2: + def __init__(self, pages, history_length, scratch_range=None, num_sink_blocks=0): + self._pages = list(pages) + self.history_length = history_length + self._scratch_range = scratch_range + # The sink guard reads the manager-internal life cycle (no public API). + self.manager = SimpleNamespace( + _life_cycles=[SimpleNamespace(num_sink_blocks=num_sink_blocks)] + ) + + def get_aggregated_page_indices(self, group_idx, valid_only=True): + # index == ordinal only holds for valid_only=False; the adapter must + # never ask for the compacted view. + assert valid_only is False + return list(self._pages) - The receiver keeps a larger window when only it runs speculative decoding, - so its suffix starts earlier. Both token starts are derived from list - length, so the extra blocks must come off the head. - """ + def get_scratch_desc(self, group_idx): + if self._scratch_range is None: + return None + return SimpleNamespace(range=self._scratch_range) - WINDOW = 128 - def test_extra_receiver_blocks_come_off_the_head(self): - src_block_ids = np.array([10], dtype=np.int64) - dst_block_ids = np.array([20, 21], dtype=np.int64) +class _FakeV2Mgr: + enable_block_reuse = False - trimmed = Sender._trim_receiver_window_head( - src_block_ids, dst_block_ids, peer_window_size=self.WINDOW, beam_width=1 - ) + def __init__(self, kv_cache, tokens_per_block=8, cp_size=1): + self.tokens_per_block = tokens_per_block + self.kv_cache_map = {0: kv_cache} + self.mapping = SimpleNamespace(cp_size=cp_size) - np.testing.assert_array_equal(trimmed, [21]) - def test_trimmed_receiver_maps_onto_the_last_prompt_block(self): - # Regression: trimming the tail leaves [20], which _align_kv_blocks then - # pairs with src block 10 -- one block early, so the last prompt block - # is never written. - src_block_ids = np.array([10], dtype=np.int64) - dst_block_ids = np.array([20, 21], dtype=np.int64) - total_blocks = 1225 - tpb = 128 +class TestV2AdapterTransferSpan: + TPB = 8 + PROMPT = 32 # 4 blocks - dst_block_ids = Sender._trim_receiver_window_head( - src_block_ids, dst_block_ids, peer_window_size=self.WINDOW, beam_width=1 + def _span(self, kv_cache, window=None, prompt_len=PROMPT, cp_size=1): + adapter = _CacheReuseAdapterV2( + _FakeV2Mgr(kv_cache, tokens_per_block=self.TPB, cp_size=cp_size) ) - src_start = (total_blocks - Sender._beam0_block_count(src_block_ids, total_blocks, 1)) * tpb - dst_start = (total_blocks - Sender._beam0_block_count(dst_block_ids, total_blocks, 1)) * tpb + req = _FakeReq(prompt_len=prompt_len) + return adapter.get_transfer_span(req, 0, _lg(window=window)) - src, dst = Sender._align_kv_blocks( - src_block_ids, - dst_block_ids, - src_token_start=src_start, - dst_token_start=dst_start, - tokens_per_block=tpb, + def test_fully_backed_prompt_anchor_zero(self): + kv = _FakeKvCacheV2([10, 11, 12, 13], history_length=self.PROMPT) + + pages, anchor = self._span(kv) + + np.testing.assert_array_equal(pages, [10, 11, 12, 13]) + assert anchor == 0 + + def test_swa_stale_holes_anchor_the_backed_run(self): + # window 16 → stale_end = (32 + 1 - 16) // 8 = 2: ordinals 0-1 unbacked. + kv = _FakeKvCacheV2([BAD_PAGE_INDEX, BAD_PAGE_INDEX, 12, 13], history_length=self.PROMPT) + + pages, anchor = self._span(kv, window=16) + + np.testing.assert_array_equal(pages, [12, 13]) + assert anchor == 2 + + def test_scratch_hole_is_explained(self): + # Ordinal 0 was written to a rotating shared scratch slot. + kv = _FakeKvCacheV2( + [BAD_PAGE_INDEX, 11, 12, 13], history_length=self.PROMPT, scratch_range=(0, 1) ) - np.testing.assert_array_equal(src, [10]) - np.testing.assert_array_equal(dst, [21]) + pages, anchor = self._span(kv) + + np.testing.assert_array_equal(pages, [11, 12, 13]) + assert anchor == 1 - def test_equal_counts_are_untouched(self): - src_block_ids = np.array([10, 11], dtype=np.int64) - dst_block_ids = np.array([20, 21], dtype=np.int64) + def test_unexplained_hole_raises(self): + """An unexplained unbacked ordinal means the span cannot be trusted.""" + kv = _FakeKvCacheV2([10, BAD_PAGE_INDEX, 12, 13], history_length=self.PROMPT) - trimmed = Sender._trim_receiver_window_head( - src_block_ids, dst_block_ids, peer_window_size=self.WINDOW, beam_width=1 + with pytest.raises(RuntimeError, match="refusing to transfer misaligned KV blocks"): + self._span(kv) + + def test_hole_past_swa_stale_end_raises(self): + # window 16 explains ordinals < 2 only; a hole at ordinal 2 is a bug. + kv = _FakeKvCacheV2( + [BAD_PAGE_INDEX, BAD_PAGE_INDEX, BAD_PAGE_INDEX, 13], history_length=self.PROMPT ) - np.testing.assert_array_equal(trimmed, [20, 21]) + with pytest.raises(RuntimeError, match="unbacked"): + self._span(kv, window=16) + + def test_scratch_tail_past_prompt_is_trimmed(self): + # A 5th page past the prompt (speculative scratch) is never transferred. + kv = _FakeKvCacheV2([10, 11, 12, 13, 99], history_length=self.PROMPT) + + pages, anchor = self._span(kv) + + np.testing.assert_array_equal(pages, [10, 11, 12, 13]) + assert 99 not in pages + assert anchor == 0 + + def test_all_unbacked_returns_empty_anchor_zero(self): + kv = _FakeKvCacheV2([BAD_PAGE_INDEX] * 4, history_length=self.PROMPT) + + pages, anchor = self._span(kv, window=16) - def test_smaller_receiver_is_untouched(self): - # Generation prefix-cache reuse: handled downstream via dst_start. - src_block_ids = np.array([10, 11, 12], dtype=np.int64) - dst_block_ids = np.array([20], dtype=np.int64) + assert pages.size == 0 + assert anchor == 0 - trimmed = Sender._trim_receiver_window_head( - src_block_ids, dst_block_ids, peer_window_size=self.WINDOW, beam_width=1 + def test_hole_at_last_prompt_block_returns_empty(self): + # The contiguous backed run must end at the last prompt block. + kv = _FakeKvCacheV2([10, 11, 12, BAD_PAGE_INDEX], history_length=self.PROMPT) + + pages, anchor = self._span(kv) + + assert pages.size == 0 + assert anchor == 0 + + def test_sink_blocks_with_holes_raise(self): + # A single anchor cannot represent a backed sink prefix followed by an + # unbacked hole; a sink-configured life cycle must fail loud. + kv = _FakeKvCacheV2( + [BAD_PAGE_INDEX, BAD_PAGE_INDEX, 12, 13], + history_length=self.PROMPT, + num_sink_blocks=1, ) - np.testing.assert_array_equal(trimmed, [20]) + with pytest.raises(RuntimeError, match="token sinks"): + self._span(kv, window=16) + + def test_sink_blocks_without_holes_pass(self): + # Sinks only conflict with the anchor when there are unbacked ordinals. + kv = _FakeKvCacheV2([10, 11, 12, 13], history_length=self.PROMPT, num_sink_blocks=1) - def test_non_windowed_group_still_raises(self): - src_block_ids = np.array([10], dtype=np.int64) - dst_block_ids = np.array([20, 21], dtype=np.int64) + pages, anchor = self._span(kv) - with pytest.raises(ValueError, match="block count mismatch"): - Sender._trim_receiver_window_head( - src_block_ids, dst_block_ids, peer_window_size=None, beam_width=1 - ) + np.testing.assert_array_equal(pages, [10, 11, 12, 13]) + assert anchor == 0 - def test_multi_beam_still_raises(self): - src_block_ids = np.array([10], dtype=np.int64) - dst_block_ids = np.array([20, 21], dtype=np.int64) + def test_helix_cp_returns_local_list_anchor_zero(self): + # Helix lists are strided local subsets; the prompt cap and global + # stale/scratch bookkeeping do not apply and the anchor stays 0. + kv = _FakeKvCacheV2([10, 12, 14, 16, 18], history_length=self.PROMPT) - with pytest.raises(ValueError, match="block count mismatch"): - Sender._trim_receiver_window_head( - src_block_ids, dst_block_ids, peer_window_size=self.WINDOW, beam_width=4 - ) + pages, anchor = self._span(kv, window=16, cp_size=2) + + np.testing.assert_array_equal(pages, [10, 12, 14, 16, 18]) + assert anchor == 0 + + def test_helix_cp_with_unbacked_ordinals_raises(self): + # Local block lists carry no global stale/scratch bookkeeping, so an + # unbacked ordinal cannot be anchored under helix. + kv = _FakeKvCacheV2([BAD_PAGE_INDEX, 12, 14], history_length=self.PROMPT) + + with pytest.raises(RuntimeError, match="helix"): + self._span(kv, window=16, cp_size=2) # --------------------------------------------------------------------------- -# _create_kv_slice: the block list spans prompt_len, excluding the extra KV -# slots speculative decoding reserves. +# _create_kv_slice: transfer-policy head trims (SWA bandwidth skip, gen-side +# reuse skip) are explicit head-slices that advance the anchor. # --------------------------------------------------------------------------- def _build_transceiver_for_kv_slice( - num_extra_kv_tokens: int, - prompt_len: int, *, + prompt_len: int, + block_ids, + span_anchor: int = 0, tokens_per_block: int = 8, - block_ids=None, sliding_window_size=None, cached_tokens: int = 0, is_generation_only: bool = False, beam_width: int = 1, + cp_size: int = 1, ): """Stub a KvCacheTransceiverV2 so _create_kv_slice runs without dist setup. Wires only the attributes the method touches: - - reuse adapter: tokens_per_block, per-layer-group cached count, block ids + - reuse adapter: tokens_per_block, per-layer-group cached count, and the + anchored (pages, first_block_ordinal) span - page table: layer groups - - cache manager: num_extra_kv_tokens (read in this code path) + - mapping: cp_size (helix skips global-position trims) """ layer_group = AttentionLayerGroup( pool_group_idx=0, kv_head_num_per_rank=1, sliding_window_size=sliding_window_size, ) - total_blocks = (prompt_len + num_extra_kv_tokens + tokens_per_block - 1) // tokens_per_block - if block_ids is None: - block_ids = np.arange(total_blocks, dtype=np.int64) - else: - block_ids = np.asarray(block_ids, dtype=np.int64) + block_ids = np.asarray(block_ids, dtype=np.int64) reuse_adapter = SimpleNamespace( tokens_per_block=tokens_per_block, get_cached_token_count_per_layer_group=lambda req, layer_groups: [cached_tokens] * len(layer_groups), - get_block_ids=lambda req, idx, lg: block_ids, + get_transfer_span=lambda req, idx, lg: (block_ids, span_anchor), ) - page_table = SimpleNamespace(layer_groups=[layer_group]) - cache_manager = SimpleNamespace(num_extra_kv_tokens=num_extra_kv_tokens) transceiver = object.__new__(KvCacheTransceiverV2) transceiver._reuse_adapter = reuse_adapter - transceiver._page_table = page_table - transceiver._kv_cache_manager = cache_manager + transceiver._page_table = SimpleNamespace(layer_groups=[layer_group]) + transceiver._kv_cache_manager = SimpleNamespace() + transceiver._mapping = SimpleNamespace(cp_size=cp_size) req = SimpleNamespace( prompt_len=prompt_len, @@ -406,149 +739,207 @@ def _build_transceiver_for_kv_slice( return transceiver, req -class TestCreateKvSliceBlockSpan: - """The block list must span prompt_len, not prompt_len + num_extra_kv_tokens. - - A monolithic slice carries no extent of its own: the sender's suffix - arithmetic anchors on the session's prompt_len and assumes the list is the - tail of ceil(prompt_len / tpb) blocks. An extra block would shift every - per-layer token start. - """ +class TestCreateKvSliceAnchorTrims: + """Anchor-trim tests: tpb=8, prompt_len=32 (4 blocks), window start 2.""" - def test_excludes_num_extra_kv_tokens(self): - prompt_len = 17 - num_extra_kv_tokens = 7 - transceiver, req = _build_transceiver_for_kv_slice(num_extra_kv_tokens, prompt_len) - tpb = transceiver._reuse_adapter.tokens_per_block + def _slice(self, **kwargs): + include = kwargs.pop("include_window_groups", True) + transceiver, req = _build_transceiver_for_kv_slice(**kwargs) + return transceiver._create_kv_slice(req, include_window_groups=include) - kv_slice = transceiver._create_kv_slice(req) + def test_full_attention_span_passes_through(self): + kv_slice = self._slice(prompt_len=32, block_ids=[100, 101, 102, 103]) - assert kv_slice.block_ids_per_layer_groups[0].size == (prompt_len + tpb - 1) // tpb + np.testing.assert_array_equal(kv_slice.block_ids_per_layer_groups[0], [100, 101, 102, 103]) + assert kv_slice.first_ordinals == [0] + assert kv_slice.is_last_slice is True - def test_extra_tokens_do_not_cross_block_boundary(self): - prompt_len = 16 - num_extra_kv_tokens = 7 - transceiver, req = _build_transceiver_for_kv_slice(num_extra_kv_tokens, prompt_len) - tpb = transceiver._reuse_adapter.tokens_per_block + def test_swa_policy_trim_advances_anchor(self): + # Pre-window blocks are a bandwidth-only skip; the head-slice advances + # the anchor so the remaining blocks stay honestly positioned. + kv_slice = self._slice( + prompt_len=32, block_ids=[100, 101, 102, 103], sliding_window_size=16 + ) - # Setup must actually exercise a boundary crossing: prompt_len ends on a - # block boundary and the extra tokens would otherwise add a block. - assert prompt_len % tpb == 0 - assert (prompt_len + num_extra_kv_tokens + tpb - 1) // tpb == prompt_len // tpb + 1 + np.testing.assert_array_equal(kv_slice.block_ids_per_layer_groups[0], [102, 103]) + assert kv_slice.first_ordinals == [2] - kv_slice = transceiver._create_kv_slice(req) + def test_swa_trim_respects_adapter_anchor(self): + # V2-style span: the adapter already excluded the stale prefix. + kv_slice = self._slice( + prompt_len=32, block_ids=[102, 103], span_anchor=2, sliding_window_size=16 + ) - assert kv_slice.block_ids_per_layer_groups[0].size == prompt_len // tpb + np.testing.assert_array_equal(kv_slice.block_ids_per_layer_groups[0], [102, 103]) + assert kv_slice.first_ordinals == [2] - def test_defaults_to_prompt_len_when_no_extra(self): - prompt_len = 17 - transceiver, req = _build_transceiver_for_kv_slice( - num_extra_kv_tokens=0, prompt_len=prompt_len + def test_adapter_anchor_beyond_policy_target_is_kept(self): + # The manager evicted more than the policy would trim; the manager wins. + kv_slice = self._slice( + prompt_len=32, block_ids=[103], span_anchor=3, sliding_window_size=16 ) - tpb = transceiver._reuse_adapter.tokens_per_block - kv_slice = transceiver._create_kv_slice(req) + np.testing.assert_array_equal(kv_slice.block_ids_per_layer_groups[0], [103]) + assert kv_slice.first_ordinals == [3] + + def test_gen_reuse_skip_advances_anchor(self): + kv_slice = self._slice( + prompt_len=32, + block_ids=[100, 101, 102, 103], + cached_tokens=16, + is_generation_only=True, + ) - assert kv_slice.block_ids_per_layer_groups[0].size == (prompt_len + tpb - 1) // tpb + np.testing.assert_array_equal(kv_slice.block_ids_per_layer_groups[0], [102, 103]) + assert kv_slice.first_ordinals == [2] - def test_swa_caps_oversized_non_speculative_v1_list_before_window_trim(self): - transceiver, req = _build_transceiver_for_kv_slice( - num_extra_kv_tokens=0, + def test_gen_reuse_and_swa_take_the_max(self): + # window start 2 vs cached 3 blocks → skip to 3. + kv_slice = self._slice( prompt_len=32, - block_ids=[100, 101, 102, 103, 104], + block_ids=[100, 101, 102, 103], sliding_window_size=16, + cached_tokens=24, + is_generation_only=True, ) - kv_slice = transceiver._create_kv_slice(req) + np.testing.assert_array_equal(kv_slice.block_ids_per_layer_groups[0], [103]) + assert kv_slice.first_ordinals == [3] - np.testing.assert_array_equal( - kv_slice.block_ids_per_layer_groups[0], - np.array([102, 103], dtype=np.int64), + def test_gen_reuse_inside_stale_region_adds_no_skip(self): + # Regression: reuse-hit (1 block) below the window start (2) must not + # skip anything beyond the SWA trim itself. + kv_slice = self._slice( + prompt_len=32, + block_ids=[102, 103], + span_anchor=2, + sliding_window_size=16, + cached_tokens=8, + is_generation_only=True, ) - def test_swa_allocation_cap_preserves_packed_beam_tails(self): - transceiver, req = _build_transceiver_for_kv_slice( - num_extra_kv_tokens=0, + np.testing.assert_array_equal(kv_slice.block_ids_per_layer_groups[0], [102, 103]) + assert kv_slice.first_ordinals == [2] + + def test_ctx_side_ignores_cached_tokens(self): + # get_cached_token_count_per_layer_group is a gen-side concern; the ctx + # sender transfers its full window. + kv_slice = self._slice( prompt_len=32, - block_ids=[100, 101, 102, 103, 104, 200, 201, 202], + block_ids=[100, 101, 102, 103], sliding_window_size=16, - beam_width=4, + cached_tokens=24, + is_generation_only=False, ) - kv_slice = transceiver._create_kv_slice(req) + np.testing.assert_array_equal(kv_slice.block_ids_per_layer_groups[0], [102, 103]) + assert kv_slice.first_ordinals == [2] - np.testing.assert_array_equal( - kv_slice.block_ids_per_layer_groups[0], - np.array([102, 103, 200, 201, 202], dtype=np.int64), + def test_window_offset_skip_subtracts_the_adapter_anchor(self): + # window=24 → window start 1; cached 2 blocks → skip only 1 more block + # from the anchored span, not cached//tpb blocks from its head. + kv_slice = self._slice( + prompt_len=32, + block_ids=[10, 11, 12], + span_anchor=1, + sliding_window_size=24, + cached_tokens=16, + is_generation_only=True, ) - @pytest.mark.parametrize( - "block_ids", - ( - pytest.param([100, 101, 102, 103, 104], id="v1-pre-eviction"), - pytest.param([102, 103, 104], id="v2-valid-only"), - ), - ) - def test_swa_trims_speculative_tail_before_stale_prompt_blocks(self, block_ids): - transceiver, req = _build_transceiver_for_kv_slice( - num_extra_kv_tokens=2, + np.testing.assert_array_equal(kv_slice.block_ids_per_layer_groups[0], [11, 12]) + assert kv_slice.first_ordinals == [2] + + def test_window_covering_prompt_behaves_like_full_attention(self): + kv_slice = self._slice( prompt_len=32, - block_ids=block_ids, - sliding_window_size=16, - cached_tokens=16, + block_ids=[10, 11, 12, 13], + sliding_window_size=32, + cached_tokens=8, is_generation_only=True, ) - kv_slice = transceiver._create_kv_slice(req) + np.testing.assert_array_equal(kv_slice.block_ids_per_layer_groups[0], [11, 12, 13]) + assert kv_slice.first_ordinals == [1] - np.testing.assert_array_equal( - kv_slice.block_ids_per_layer_groups[0], - np.array([102, 103], dtype=np.int64), + def test_skip_covering_the_span_yields_empty_anchor_zero(self): + kv_slice = self._slice( + prompt_len=32, + block_ids=[100, 101, 102, 103], + cached_tokens=32, + is_generation_only=True, ) - def test_swa_speculative_tail_requires_single_beam(self): - transceiver, req = _build_transceiver_for_kv_slice( - num_extra_kv_tokens=2, + assert kv_slice.block_ids_per_layer_groups[0].size == 0 + assert kv_slice.first_ordinals == [0] + + def test_head_slice_preserves_packed_beam_tails(self): + kv_slice = self._slice( prompt_len=32, - block_ids=[100, 101, 102, 103, 104], + block_ids=[100, 101, 102, 103, 200, 201, 202], sliding_window_size=16, beam_width=4, ) - with pytest.raises(ValueError, match="speculative scratch blocks require beam_width == 1"): - transceiver._create_kv_slice(req) + np.testing.assert_array_equal( + kv_slice.block_ids_per_layer_groups[0], [102, 103, 200, 201, 202] + ) + assert kv_slice.first_ordinals == [2] - @pytest.mark.parametrize("prompt_len", (1150, 1151)) - def test_dspark_disagg_boundary_keeps_only_initialized_swa(self, prompt_len): - tokens_per_block = 128 - total_blocks = (prompt_len + tokens_per_block - 1) // tokens_per_block - sliding_window_size = 128 + 5 - stale_end = max( - 0, - (prompt_len + 1 - sliding_window_size) // tokens_per_block, - ) - valid_prompt_blocks = total_blocks - stale_end - block_ids = np.arange( - 200, - 200 + valid_prompt_blocks + 1, - dtype=np.int64, - ) - transceiver, req = _build_transceiver_for_kv_slice( - num_extra_kv_tokens=5, - prompt_len=prompt_len, - tokens_per_block=tokens_per_block, - block_ids=block_ids, - sliding_window_size=sliding_window_size, + def test_skip_covering_beam0_drops_tails_too(self): + kv_slice = self._slice( + prompt_len=8, # 1 block + block_ids=[10, 11, 12, 13], # 1 beam-0 block + 3 tails + beam_width=4, + cached_tokens=8, is_generation_only=True, ) - kv_slice = transceiver._create_kv_slice(req) + assert kv_slice.block_ids_per_layer_groups[0].size == 0 + assert kv_slice.first_ordinals == [0] - np.testing.assert_array_equal( - kv_slice.block_ids_per_layer_groups[0], - block_ids[:-1], + def test_helix_skips_global_position_trims(self): + # Helix spans are strided local subsets; the SWA trim is a global- + # ordinal concept and must not apply. + kv_slice = self._slice( + prompt_len=32, + block_ids=[100, 102], + sliding_window_size=16, + cp_size=2, + ) + + np.testing.assert_array_equal(kv_slice.block_ids_per_layer_groups[0], [100, 102]) + assert kv_slice.first_ordinals == [0] + + def test_deferred_window_group_is_empty_with_anchor_zero(self): + # Pipelined senders defer SWA groups to the final chunk. + kv_slice = self._slice( + prompt_len=32, + block_ids=[100, 101, 102, 103], + sliding_window_size=16, + include_window_groups=False, + ) + + assert kv_slice.block_ids_per_layer_groups[0].size == 0 + assert kv_slice.first_ordinals == [0] + + def test_window_at_least_prompt_is_not_deferred(self): + kv_slice = self._slice( + prompt_len=32, + block_ids=[100, 101, 102, 103], + sliding_window_size=32, + include_window_groups=False, ) + np.testing.assert_array_equal(kv_slice.block_ids_per_layer_groups[0], [100, 101, 102, 103]) + assert kv_slice.first_ordinals == [0] + + def test_empty_span_stays_empty_with_anchor_zero(self): + kv_slice = self._slice(prompt_len=32, block_ids=[]) + + assert kv_slice.block_ids_per_layer_groups[0].size == 0 + assert kv_slice.first_ordinals == [0] + # --------------------------------------------------------------------------- # CacheReuseAdapter.get_cached_token_count_per_layer_group: SWA clamp. @@ -574,31 +965,13 @@ def tokens_per_block(self) -> int: def _global_cached_token_count(self, req): # noqa: ARG002 return self._scalar - def get_block_ids(self, req, group_idx, lg): # noqa: ARG002 - return np.array([], dtype=np.int64) + def get_transfer_span(self, req, group_idx, lg): # noqa: ARG002 + return np.array([], dtype=np.int64), 0 def commit_blocks_for_reuse(self, req): # noqa: ARG002 pass -class _FakeReq: - def __init__(self, prompt_len: int): - self.prompt_len = prompt_len - - -class _FakeSamplingConfig: - def __init__(self, beam_width: int): - self.beam_width = beam_width - - -def _lg(window=None): - return AttentionLayerGroup( - pool_group_idx=0, - sliding_window_size=window, - local_layers=[LocalLayer(local_layer_id=0, global_layer_id=0)], - ) - - class TestAdapterPerLayerGroup: """Per-layer cached prefix: adapter reports only the reuse-hit scalar.""" @@ -629,7 +1002,7 @@ def test_swa_passthrough_above_stale(self): def test_swa_passthrough_below_stale(self): # scalar=8 is below stale_end*tpb=16; adapter still returns the raw - # scalar — the call site reconciles with stale_end via max(0, ...). + # scalar — the call site reconciles with the window start via max(). ad = _StubAdapter(scalar=8, tpb=self.TPB) out = ad.get_cached_token_count_per_layer_group(_FakeReq(32), [_lg(window=16)]) assert out == [8] @@ -644,206 +1017,61 @@ def test_mixed_groups(self): # --------------------------------------------------------------------------- -# _create_kv_slice SWA block trim: window-trim + cache-skip via per-layer cached. +# Sender anchored token starts: src/dst starts come straight from the +# per-group anchors (first_ordinal * tpb), replacing the deleted +# (total_blocks - len) * tpb suffix derivation and its SWA clamp. # --------------------------------------------------------------------------- -def _swa_trim(block_ids, prompt_len, tpb, window_size, cached_tokens, is_gen_only=True): - """Replicate the SWA branch of KvCacheTransceiverV2._create_kv_slice. - - Inputs: - block_ids: list possibly containing stale entries (V1 pre-eviction view). - cached_tokens: reuse-hit prefix reported by the adapter (token-aligned). - is_gen_only: True mirrors the gen-side path; False mirrors the ctx-side - path where ``cached_per_lg`` is synthetically 0. - """ - block_ids = np.array(block_ids, dtype=np.int64) - total_blocks = (prompt_len + tpb - 1) // tpb - stale_end = max(0, (prompt_len + 1 - window_size) // tpb) - expected_valid = max(0, total_blocks - stale_end) - if block_ids.size > expected_valid: - block_ids = ( - block_ids[-expected_valid:] if expected_valid > 0 else np.array([], dtype=np.int64) - ) - # Ctx side bypasses adapter (cached=0); gen side uses adapter scalar. - cached_lg = cached_tokens if is_gen_only else 0 - # Reuse-hit blocks beyond the already-pruned stale region. - cache_skip = max(0, cached_lg // tpb - stale_end) - if cache_skip > 0: - block_ids = ( - block_ids[cache_skip:] if cache_skip < block_ids.size else np.array([], dtype=np.int64) - ) - return block_ids - - -class TestSwaTrim: - """Window-trim + cache-skip in _create_kv_slice's SWA path. - - Setup: tpb=8, prompt_len=32 → total_blocks=4; window=16 → stale_end=2. - """ - +class TestSenderAnchoredStarts: TPB = 8 - PROMPT_LEN = 32 - WINDOW = 16 - - def _trim(self, blocks, scalar): - return _swa_trim(blocks, self.PROMPT_LEN, self.TPB, self.WINDOW, scalar) - - def test_no_cache(self): - np.testing.assert_array_equal(self._trim([20, 21], scalar=0), [20, 21]) - - def test_cache_entirely_stale(self): - # scalar=16=stale_end*tpb → cached_lg=16, cache_skip=0. - np.testing.assert_array_equal(self._trim([20, 21], scalar=16), [20, 21]) - - def test_cache_one_block_in_window(self): - # scalar=24 → cached_lg=24, cache_skip=24/8-2=1. - np.testing.assert_array_equal(self._trim([20, 21], scalar=24), [21]) - - def test_cache_covers_full_window(self): - # scalar=32 → cache_skip=2, list size=2 → empty. - assert self._trim([20, 21], scalar=32).size == 0 - - def test_window_offset_skip_subtracts_stale(self): - # window=24 → stale_end=1; scalar=16 (2 blocks) → cache_skip=2-1=1. - # Naive block_ids[scalar//tpb:] would skip 2 from a 3-block list and return 1 block. - out = _swa_trim([10, 11, 12], prompt_len=32, tpb=8, window_size=24, cached_tokens=16) - np.testing.assert_array_equal(out, [11, 12]) - - def test_window_covers_all_no_stale(self): - # window=prompt_len → stale_end=0; behaves like full-attn. - out = _swa_trim([10, 11, 12, 13], prompt_len=32, tpb=8, window_size=32, cached_tokens=8) - np.testing.assert_array_equal(out, [11, 12, 13]) - - def test_v1_pre_eviction_includes_stale(self): - # Pre-eviction list has all 4 blocks; window-trim keeps last expected_valid=2. - out = _swa_trim([10, 11, 12, 13], self.PROMPT_LEN, self.TPB, self.WINDOW, 0) - np.testing.assert_array_equal(out, [12, 13]) - - def test_ctx_side_no_adapter_no_skip(self): - # Ctx-side path: adapter not invoked, cached_per_lg synthetically 0. - # cache_skip = max(0, 0 - stale_end) = 0 — full valid window is sent. - out = _swa_trim([20, 21], self.PROMPT_LEN, self.TPB, self.WINDOW, 0, is_gen_only=False) - np.testing.assert_array_equal(out, [20, 21]) - - def test_ctx_side_v1_pre_eviction(self): - # Ctx-side path with V1 pre-eviction list: window-trim still drops stale - # blocks, cache_skip stays 0 so trimmed window is sent in full. - out = _swa_trim( - [10, 11, 12, 13], self.PROMPT_LEN, self.TPB, self.WINDOW, 0, is_gen_only=False - ) - np.testing.assert_array_equal(out, [12, 13]) - def test_gen_side_reuse_inside_stale_no_skip(self): - # gen side with reuse-hit fully inside the stale region: cache_skip = 0. - # Regression for SWA + reuse-hit < stale_end*tpb (no adapter clamp). - out = _swa_trim([20, 21], self.PROMPT_LEN, self.TPB, self.WINDOW, 8, is_gen_only=True) - np.testing.assert_array_equal(out, [20, 21]) - - -# --------------------------------------------------------------------------- -# Sender token-start derivation: (total_blocks - n_blocks) * tpb + SWA clamp. -# --------------------------------------------------------------------------- - - -def _derive_starts(prompt_len, tpb, window_size, n_src, n_dst, slice_end=None): - """Replicate _build_kv_write_meta's per-layer src/dst token-start derivation.""" - if slice_end is None: - slice_end = prompt_len - total_blocks = (slice_end + tpb - 1) // tpb - src_start = (total_blocks - n_src) * tpb - dst_start = (total_blocks - n_dst) * tpb - if window_size is not None: - stale_end = max(0, (prompt_len + 1 - window_size) // tpb) - src_start = max(stale_end * tpb, src_start) - dst_start = max(stale_end * tpb, dst_start) - return src_start, dst_start - - -class TestSenderTokenStarts: - """Verify (total_blocks - n) * tpb + SWA clamp produces correct src/dst starts.""" - - TPB = 8 - - def _align(self, src, dst, src_start, dst_start): + def _align(self, src, src_anchor, dst, dst_anchor): return Sender._align_kv_blocks( np.array(src, dtype=np.int64), np.array(dst, dtype=np.int64), - src_token_start=src_start, - dst_token_start=dst_start, + src_token_start=src_anchor * self.TPB, + dst_token_start=dst_anchor * self.TPB, tokens_per_block=self.TPB, ) def test_full_prompt_no_cache(self): - src_start, dst_start = _derive_starts( - prompt_len=32, tpb=self.TPB, window_size=None, n_src=4, n_dst=4 - ) - assert (src_start, dst_start) == (0, 0) - - def test_full_prompt_dst_cached(self): - # dst cached 2 blocks → dst sends 2 → dst_start=16. - src_start, dst_start = _derive_starts( - prompt_len=32, tpb=self.TPB, window_size=None, n_src=4, n_dst=2 - ) - assert (src_start, dst_start) == (0, 16) - - def test_full_prompt_src_cached(self): - src_start, dst_start = _derive_starts( - prompt_len=32, tpb=self.TPB, window_size=None, n_src=2, n_dst=4 - ) - assert (src_start, dst_start) == (16, 0) + src, dst = self._align([10, 11, 12, 13], 0, [20, 21, 22, 23], 0) + np.testing.assert_array_equal(src, [10, 11, 12, 13]) + np.testing.assert_array_equal(dst, [20, 21, 22, 23]) - def test_swa_no_cache_stale_present(self): - # window=16 → stale_end=2 → stale_end*tpb=16; both sides 2 blocks. - src_start, dst_start = _derive_starts( - prompt_len=32, tpb=self.TPB, window_size=16, n_src=2, n_dst=2 - ) - assert (src_start, dst_start) == (16, 16) + def test_dst_cached_prefix(self): + # dst reused 2 blocks → its span is anchored at 2; src head is trimmed. + src, dst = self._align([10, 11, 12, 13], 0, [22, 23], 2) + np.testing.assert_array_equal(src, [12, 13]) + np.testing.assert_array_equal(dst, [22, 23]) - def test_swa_dst_cache_in_stale_region(self): - # dst cached 2 blocks but all stale → dst still has 2 valid window blocks. - src_start, dst_start = _derive_starts( - prompt_len=32, tpb=self.TPB, window_size=16, n_src=2, n_dst=2 - ) - result_src, result_dst = self._align([10, 11], [20, 21], src_start, dst_start) - np.testing.assert_array_equal(result_src, [10, 11]) - np.testing.assert_array_equal(result_dst, [20, 21]) + def test_src_cached_prefix(self): + src, dst = self._align([12, 13], 2, [20, 21, 22, 23], 0) + np.testing.assert_array_equal(src, [12, 13]) + np.testing.assert_array_equal(dst, [22, 23]) - def test_swa_src_cache_inside_window(self): - # window=24 → stale_end=1 → stale_end*tpb=8. - # ctx cached 16 tokens (2 blocks), window-trim leaves 3 blocks, skip 1 → src has 2 blocks. - # dst no cache → window-trim leaves 3 blocks. - src_start, dst_start = _derive_starts( - prompt_len=32, tpb=self.TPB, window_size=24, n_src=2, n_dst=3 - ) - # total_blocks = 4. src_start = (4-2)*8 = 16. dst_start = (4-3)*8 = 8. SWA clamp keeps both. - assert (src_start, dst_start) == (16, 8) - result_src, result_dst = self._align([10, 11], [20, 21, 22], src_start, dst_start) - np.testing.assert_array_equal(result_src, [10, 11]) - np.testing.assert_array_equal(result_dst, [21, 22]) + def test_swa_both_sides_anchored_at_window_start(self): + src, dst = self._align([10, 11], 2, [20, 21], 2) + np.testing.assert_array_equal(src, [10, 11]) + np.testing.assert_array_equal(dst, [20, 21]) - def test_swa_window_covers_prompt_no_stale(self): - # window=prompt_len → stale_end=0; SWA clamp is a no-op. - src_start, dst_start = _derive_starts( - prompt_len=32, tpb=self.TPB, window_size=32, n_src=4, n_dst=4 - ) - assert (src_start, dst_start) == (0, 0) + def test_swa_asymmetric_anchors(self): + # ctx trimmed window start + reuse skip → anchor 2; gen window only → 1. + src, dst = self._align([10, 11], 2, [20, 21, 22], 1) + np.testing.assert_array_equal(src, [10, 11]) + np.testing.assert_array_equal(dst, [21, 22]) - def test_chunked_slice_end_below_prompt(self): - # Non-final slice: slice_end=16, prompt_len=32, no window. - # 2 blocks in slice; cache-free. - src_start, dst_start = _derive_starts( - prompt_len=32, tpb=self.TPB, window_size=None, n_src=2, n_dst=2, slice_end=16 - ) - assert (src_start, dst_start) == (0, 0) + def test_chunk_anchored_at_its_start(self): + # Non-final chunk covering blocks [0, 2) against an uncached receiver. + src, dst = self._align([10, 11], 0, [20, 21, 22, 23], 0) + np.testing.assert_array_equal(src, [10, 11]) + np.testing.assert_array_equal(dst, [20, 21]) - def test_chunked_slice_entirely_stale_for_swa(self): - # slice_end=16 ≤ stale_end*tpb=16 → SWA layer sends 0 blocks; clamp pushes start to 16. - src_start, dst_start = _derive_starts( - prompt_len=32, tpb=self.TPB, window_size=16, n_src=0, n_dst=0, slice_end=16 - ) - # total_blocks for slice = 2 → raw start = 16; clamped = max(16, 16) = 16. - assert (src_start, dst_start) == (16, 16) + def test_chunk_entirely_within_receiver_cached_prefix(self): + src, dst = self._align([10, 11], 0, [22, 23], 2) + assert src.size == 0 + assert dst.size == 0 # --------------------------------------------------------------------------- diff --git a/tests/unittest/disaggregated/test_chunked_transfer.py b/tests/unittest/disaggregated/test_chunked_transfer.py index 7cd2bad8b94b..df6e6310ba6c 100644 --- a/tests/unittest/disaggregated/test_chunked_transfer.py +++ b/tests/unittest/disaggregated/test_chunked_transfer.py @@ -18,28 +18,25 @@ TxSession/RxSession classes with lightweight stub sender/receiver objects. """ +import itertools from types import MethodType, SimpleNamespace from unittest.mock import MagicMock +import msgpack import numpy as np import pytest from tensorrt_llm import DisaggregatedParams -from tensorrt_llm._torch.disaggregation.base.transfer import ( - KVSlice, - SessionStatus, - TokenRange, - WaitResult, -) +from tensorrt_llm._torch.disaggregation.base.transfer import KVSlice, SessionStatus, WaitResult from tensorrt_llm._torch.disaggregation.native.transfer import ( AgentResult, KVSendTask, + Receiver, RecvReqInfo, RxSession, Sender, TaskStatus, TxSession, - project_blocks_to_global_chunk, ) from tensorrt_llm._torch.disaggregation.resource.page import CacheKind from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState, LlmRequestType @@ -112,80 +109,76 @@ def _make_rx_session(num_slices: int, rid: int = 42, prompt_len: int = 8) -> RxS # --------------------------------------------------------------------------- -# Global chunk projection tests +# Anchored span intersection tests +# +# The deleted project_blocks_to_global_chunk inferred a list's global position +# from its length (suffix convention). Anchors make the position explicit; +# _align_kv_blocks intersects two anchored spans directly. # --------------------------------------------------------------------------- +_ANCHOR_TPB = 8 + + +def _align_anchored(src_ids, src_anchor, dst_ids, dst_anchor, tpb=_ANCHOR_TPB): + return Sender._align_kv_blocks( + np.asarray(src_ids, dtype=np.int64), + np.asarray(dst_ids, dtype=np.int64), + src_token_start=src_anchor * tpb, + dst_token_start=dst_anchor * tpb, + tokens_per_block=tpb, + ) -def test_chunk_projection_noops_when_chunk_is_outside_short_layer_group(): - """A shared chunk cursor past a short layer group's resident range is a no-op.""" - block_ids = np.array([10, 11, 12], dtype=np.int64) - projected_ids = project_blocks_to_global_chunk( - block_ids, - chunk_block_offset=4, - chunk_block_count=4, - resident_block_end=3, +def test_chunk_anchored_past_short_group_span_is_a_noop(): + """A chunk anchored past a short layer group's declared span transfers nothing.""" + src, dst = _align_anchored( + src_ids=np.arange(4, 8, dtype=np.int64), # chunk covering block ordinals [4, 8) + src_anchor=4, + dst_ids=np.array([10, 11, 12], dtype=np.int64), # short group spans [0, 3) + dst_anchor=0, ) - assert projected_ids.size == 0 + assert src.size == 0 + assert dst.size == 0 @pytest.mark.parametrize( - "resident_block_end,chunk_block_offset,expected", + "chunk_anchor,chunk_ids,expected_dst", [ - (16, 0, np.arange(16, dtype=np.int64)), - (32, 16, np.arange(16, 32, dtype=np.int64)), + (0, np.arange(16, dtype=np.int64), np.arange(100, 116, dtype=np.int64)), + (16, np.arange(16, 32, dtype=np.int64), np.arange(116, 132, dtype=np.int64)), ], ids=["first_chunk", "later_chunk"], ) -def test_chunk_projection_maps_incrementally_allocated_source( - resident_block_end, chunk_block_offset, expected -): - """Source blocks end at the current chunk, not at the full prompt.""" - block_ids = np.arange(resident_block_end, dtype=np.int64) +def test_anchored_chunks_map_incrementally_allocated_source(chunk_anchor, chunk_ids, expected_dst): + """Each chunk's anchor addresses the receiver's whole-prompt list directly.""" + dst_ids = np.arange(100, 132, dtype=np.int64) # 32 blocks anchored at 0 - projected_ids = project_blocks_to_global_chunk( - block_ids, - chunk_block_offset=chunk_block_offset, - chunk_block_count=16, - resident_block_end=resident_block_end, - ) + src, dst = _align_anchored(chunk_ids, chunk_anchor, dst_ids, 0) - assert np.array_equal(projected_ids, expected) + assert np.array_equal(src, chunk_ids) + assert np.array_equal(dst, expected_dst) -def test_chunk_projection_maps_prefix_reuse_suffix_by_overlap(): - """Destination suffixes are matched by overlap, not by raw chunk-offset indexing.""" - block_ids = np.array([104, 105, 106, 107], dtype=np.int64) +def test_anchored_chunks_map_prefix_reuse_suffix_by_overlap(): + """A receiver suffix anchored at its true ordinal overlaps only later chunks.""" + dst_ids = np.array([104, 105, 106, 107], dtype=np.int64) # spans [4, 8) + dst_anchor = 4 - first_chunk = project_blocks_to_global_chunk( - block_ids, - chunk_block_offset=0, - chunk_block_count=4, - resident_block_end=8, - ) - second_chunk = project_blocks_to_global_chunk( - block_ids, - chunk_block_offset=4, - chunk_block_count=4, - resident_block_end=8, - ) + first_src, first_dst = _align_anchored(np.arange(4), 0, dst_ids, dst_anchor) + second_src, second_dst = _align_anchored(np.arange(4, 8), 4, dst_ids, dst_anchor) - assert first_chunk.size == 0 - assert np.array_equal(second_chunk, block_ids) + assert first_src.size == 0 and first_dst.size == 0 + assert np.array_equal(second_src, np.arange(4, 8)) + assert np.array_equal(second_dst, dst_ids) _PROJECTION_TPB = 8 _PROJECTION_PROMPT_TOKENS = 8 * _PROJECTION_TPB -def _projection_token_range(start_block: int, end_block: int) -> TokenRange: - """A chunk's block window as the block-aligned token range that rides the slice.""" - return TokenRange(start=start_block * _PROJECTION_TPB, end=end_block * _PROJECTION_TPB) - - -def _make_projection_sender() -> Sender: - """Create a Sender wired to a stub registrar with two full-attention layer groups. +def _make_projection_sender(num_groups: int = 2, tpb: int = _PROJECTION_TPB) -> Sender: + """Create a Sender wired to a stub registrar with full-attention layer groups. Full attention is spelled as a window the prompt never outgrows, which is how both extractor paths build it; they read max_attention_window_vec, so an @@ -203,10 +196,10 @@ def _make_projection_sender() -> Sender: extractor = MagicMock() extractor.page_table = SimpleNamespace( - tokens_per_block=_PROJECTION_TPB, + tokens_per_block=tpb, layer_groups=[ - SimpleNamespace(kind=CacheKind.PAGED, sliding_window_size=_PROJECTION_PROMPT_TOKENS), - SimpleNamespace(kind=CacheKind.PAGED, sliding_window_size=_PROJECTION_PROMPT_TOKENS), + SimpleNamespace(kind=CacheKind.PAGED, sliding_window_size=_PROJECTION_PROMPT_TOKENS) + for _ in range(num_groups) ], ) # extract(region_ids, layer_group_id, pool_idx) - the caller passes the @@ -230,10 +223,7 @@ def _make_projection_sender() -> Sender: registrar.get_peer_rank_info.return_value = peer_ri registrar.get_peer_overlap.return_value = SimpleNamespace(ranks=[0]) registrar.should_send_kv.return_value = True - registrar.get_pool_mapping.return_value = { - (0, 0): (0, 0), - (1, 0): (1, 0), - } + registrar.get_pool_mapping.return_value = {(lg, 0): (lg, 0) for lg in range(num_groups)} registrar.peer_extractor.return_value = extractor registrar.get_kv_map.return_value = mapper @@ -242,7 +232,9 @@ def _make_projection_sender() -> Sender: return sender -def _make_projection_task(slice_id: int = 1) -> KVSendTask: +def _make_projection_task(slice_id: int = 1, beam_width: int = 1) -> KVSendTask: + # Group 0 carries the chunk covering block ordinals [4, 8); group 1 is a + # shorter (asymmetric) group whose resident suffix spans [5, 8). return KVSendTask( KVSlice( is_last_slice=True, @@ -250,11 +242,12 @@ def _make_projection_task(slice_id: int = 1) -> KVSendTask: np.array([4, 5, 6, 7], dtype=np.int64), np.array([10, 11, 12], dtype=np.int64), ], - token_range=_projection_token_range(4, 8), + first_ordinals=[4, 5], ), _make_params(), slice_id=slice_id, prompt_len=_PROJECTION_PROMPT_TOKENS, + beam_width=beam_width, ) @@ -268,6 +261,7 @@ def _make_projection_req_info(slice_id=None) -> RecvReqInfo: np.array([200, 201, 202], dtype=np.int64), ], unique_rid=42, + first_ordinals=[4, 5], slice_id=slice_id, ) @@ -305,7 +299,9 @@ def test_final_swa_slice_keeps_the_receivers_complete_active_window(): np.arange(2, 8, dtype=np.int64), np.array([], dtype=np.int64), ], - token_range=_projection_token_range(6, 8), + # The deferred SWA group is anchored at its true window start (2), + # not at the final context chunk's start (6). + first_ordinals=[2, 0], ), _make_params(), slice_id=1, @@ -320,6 +316,7 @@ def test_final_swa_slice_keeps_the_receivers_complete_active_window(): np.array([], dtype=np.int64), ], unique_rid=42, + first_ordinals=[2, 0], ) write_meta = sender._build_kv_write_meta(task, req_info) @@ -328,34 +325,55 @@ def test_final_swa_slice_keeps_the_receivers_complete_active_window(): assert np.array_equal(write_meta.dst_ptrs, np.arange(102, 108, dtype=np.int64)) -def test_whole_prompt_chunk_addresses_like_a_monolithic_slice(): - """A whole-prompt chunk has monolithic addressing.""" +def test_chunked_slices_reassemble_to_the_monolithic_transfer(): + """Anchored chunks concatenate to exactly the monolithic slice's pairs.""" sender = _make_projection_sender() - src_per_group = [ - np.arange(8, dtype=np.int64), - np.array([10, 11, 12], dtype=np.int64), - ] + src_group0 = np.arange(8, dtype=np.int64) # full-attention, anchored at 0 + src_group1 = np.array([10, 11, 12], dtype=np.int64) # short group, spans [5, 8) - def task_for(token_range): - return KVSendTask( + monolithic_task = KVSendTask( + KVSlice( + is_last_slice=True, + block_ids_per_layer_groups=[src_group0, src_group1], + first_ordinals=[0, 5], + ), + _make_params(), + slice_id=0, + prompt_len=64, + ) + chunk_tasks = [ + KVSendTask( KVSlice( - is_last_slice=True, - block_ids_per_layer_groups=src_per_group, - token_range=token_range, + is_last_slice=False, + block_ids_per_layer_groups=[src_group0[:4], src_group1[:0]], + first_ordinals=[0, 0], ), _make_params(), slice_id=0, prompt_len=64, - ) + ), + KVSendTask( + KVSlice( + is_last_slice=True, + block_ids_per_layer_groups=[src_group0[4:], src_group1], + first_ordinals=[4, 5], + ), + _make_params(), + slice_id=1, + prompt_len=64, + ), + ] - chunked = sender._build_kv_write_meta( - task_for(_projection_token_range(0, 8)), _make_projection_req_info() - ) - monolithic = sender._build_kv_write_meta(task_for(None), _make_projection_req_info()) + monolithic = sender._build_kv_write_meta(monolithic_task, _make_projection_req_info()) + chunked = [sender._build_kv_write_meta(t, _make_projection_req_info()) for t in chunk_tasks] - assert np.array_equal(chunked.src_ptrs, monolithic.src_ptrs) - assert np.array_equal(chunked.dst_ptrs, monolithic.dst_ptrs) - assert np.array_equal(chunked.sizes, monolithic.sizes) + chunked_src = np.concatenate([m.src_ptrs for m in chunked]) + chunked_dst = np.concatenate([m.dst_ptrs for m in chunked]) + # Pair sets are equal: every (src, dst) block pair is transferred exactly once. + assert sorted(zip(chunked_src.tolist(), chunked_dst.tolist())) == sorted( + zip(monolithic.src_ptrs.tolist(), monolithic.dst_ptrs.tolist()) + ) + assert np.concatenate([m.sizes for m in chunked]).sum() == monolithic.sizes.sum() def test_build_kv_write_meta_tracks_sender_and_receiver_slice_ids(): @@ -370,6 +388,417 @@ def test_build_kv_write_meta_tracks_sender_and_receiver_slice_ids(): assert write_meta.receiver_slice_id == 3 +def test_helix_receiver_requires_zero_block_anchors(): + """A helix peer rejects anchored (trimmed/reused) spans loudly.""" + sender = _make_projection_sender() + sender._registrar.get_peer_rank_info.return_value.cp_size = 2 + sender._registrar.get_peer_rank_info.return_value.cp_rank = 0 + + # slice_id=0 and is_last_slice=True: not a partial chunk, so the pipelined + # CP guard passes and the anchor guard is what must fire. + with pytest.raises(ValueError, match="helix CP requires zero block anchors"): + sender._build_kv_write_meta(_make_projection_task(slice_id=0), _make_projection_req_info()) + + +def test_beam_tails_pair_positionally_after_the_beam0_spans(): + """Packed beam-tail blocks ride behind the anchored beam-0 intersection.""" + sender = _make_projection_sender(num_groups=1) + beam_width = 4 + # 8 prompt blocks + 3 per-beam tail blocks on each side, both anchored at 0. + task = KVSendTask( + KVSlice( + is_last_slice=True, + block_ids_per_layer_groups=[ + np.concatenate([np.arange(8), np.array([70, 71, 72])]).astype(np.int64) + ], + first_ordinals=[0], + ), + _make_params(), + slice_id=0, + prompt_len=_PROJECTION_PROMPT_TOKENS, + beam_width=beam_width, + ) + req_info = RecvReqInfo( + sender_req_id=42, + instance_name="decode", + instance_rank=0, + block_ids_per_layer_groups=[ + np.concatenate([np.arange(100, 108), np.array([170, 171, 172])]).astype(np.int64) + ], + unique_rid=42, + first_ordinals=[0], + ) + + write_meta = sender._build_kv_write_meta(task, req_info) + + assert np.array_equal( + write_meta.src_ptrs, np.concatenate([np.arange(8), np.array([70, 71, 72])]) + ) + assert np.array_equal( + write_meta.dst_ptrs, np.concatenate([np.arange(100, 108), np.array([170, 171, 172])]) + ) + + +def test_fewer_shared_tails_than_beam_width_pair_correctly(): + """Anchored beam-0 length splits partially-shared beam tails exactly. + + The packer appends only UNSHARED final blocks, so beam_width=4 lists can + carry just 2 tails; the old beam_width-1 guess misclassified the last + beam-0 block as a tail here. + """ + sender = _make_projection_sender(num_groups=1) + beam_width = 4 + # 8 prompt blocks + 2 unshared tail blocks on each side, both anchored at 0. + task = KVSendTask( + KVSlice( + is_last_slice=True, + block_ids_per_layer_groups=[ + np.concatenate([np.arange(8), np.array([70, 71])]).astype(np.int64) + ], + first_ordinals=[0], + ), + _make_params(), + slice_id=0, + prompt_len=_PROJECTION_PROMPT_TOKENS, + beam_width=beam_width, + ) + req_info = RecvReqInfo( + sender_req_id=42, + instance_name="decode", + instance_rank=0, + block_ids_per_layer_groups=[ + np.concatenate([np.arange(100, 108), np.array([170, 171])]).astype(np.int64) + ], + unique_rid=42, + first_ordinals=[0], + ) + + write_meta = sender._build_kv_write_meta(task, req_info) + + # All 8 beam-0 blocks pair by ordinal; exactly the 2 tails ride behind. + assert np.array_equal(write_meta.src_ptrs, np.concatenate([np.arange(8), np.array([70, 71])])) + assert np.array_equal( + write_meta.dst_ptrs, np.concatenate([np.arange(100, 108), np.array([170, 171])]) + ) + + +def test_beam_tail_count_mismatch_raises(): + """Unequal packed beam-tail counts cannot pair positionally and must fail.""" + sender = _make_projection_sender(num_groups=1) + task = KVSendTask( + KVSlice( + is_last_slice=True, + block_ids_per_layer_groups=[ + np.concatenate([np.arange(8), np.array([70, 71, 72])]).astype(np.int64) + ], + first_ordinals=[0], + ), + _make_params(), + slice_id=0, + prompt_len=_PROJECTION_PROMPT_TOKENS, + beam_width=4, + ) + # dst has no tail blocks (its final blocks are shared across beams). + req_info = RecvReqInfo( + sender_req_id=42, + instance_name="decode", + instance_rank=0, + block_ids_per_layer_groups=[np.arange(100, 108, dtype=np.int64)], + unique_rid=42, + first_ordinals=[0], + ) + + with pytest.raises(ValueError, match="packed beam-tail count mismatch"): + sender._build_kv_write_meta(task, req_info) + + +def test_receiver_with_larger_window_resolves_via_anchors(): + """The receiver's earlier-starting window is aligned by anchors, not trimmed. + + Old behavior (_trim_receiver_window_head, deleted): a receiver keeping a + larger SWA window (only it runs speculative decoding) held extra leading + blocks, which the sender dropped by count so that suffix-derived starts + lined up. With explicit anchors the sender simply intersects the two + declared spans: the receiver's extra head block gets no source and the + last prompt block maps exactly (the old off-by-one regression). + """ + tpb = 128 + total_blocks = 1225 + sender = _make_projection_sender(num_groups=1, tpb=tpb) + task = KVSendTask( + KVSlice( + is_last_slice=True, + block_ids_per_layer_groups=[np.array([10], dtype=np.int64)], + first_ordinals=[total_blocks - 1], # sender window: last block only + ), + _make_params(), + slice_id=0, + prompt_len=total_blocks * tpb, + ) + req_info = RecvReqInfo( + sender_req_id=42, + instance_name="decode", + instance_rank=0, + block_ids_per_layer_groups=[np.array([20, 21], dtype=np.int64)], + unique_rid=42, + first_ordinals=[total_blocks - 2], # receiver window starts one block earlier + ) + + write_meta = sender._build_kv_write_meta(task, req_info) + + # src block 10 lands on dst block 21 (same ordinal); dst block 20 has no + # source and is not written. + assert np.array_equal(write_meta.src_ptrs, np.array([10], dtype=np.int64)) + assert np.array_equal(write_meta.dst_ptrs, np.array([21], dtype=np.int64)) + + +# --------------------------------------------------------------------------- +# Anchored alignment property test +# +# For a grid over (tpb, prompt length, per-side window sizes, gen cached +# prefix, chunk boundaries, beam width): every (src, dst) pointer pair maps +# equal block ordinals, and the transferred set equals the intersection of the +# declared spans. Block IDs encode their ordinal so the mocked extractors +# (identity ptrs) expose the mapping directly. +# --------------------------------------------------------------------------- + +_SRC_ID_BASE = 1_000 +_DST_ID_BASE = 5_000 +_SRC_TAIL_BASE = 3_000 +_DST_TAIL_BASE = 7_000 + + +def _swa_start_block(prompt_len, window, tpb): + if window is None or window >= prompt_len: + return 0 + return (prompt_len + 1 - window) // tpb + + +def _grid_req_info(dst_ids: np.ndarray, dst_anchor: int) -> RecvReqInfo: + return RecvReqInfo( + sender_req_id=42, + instance_name="decode", + instance_rank=0, + block_ids_per_layer_groups=[dst_ids], + unique_rid=42, + first_ordinals=[dst_anchor], + ) + + +def test_anchored_alignment_property_grid(): + tpb_values = (4, 16) + prompt_block_counts = (1, 3, 8) + window_block_options = (None, 2, 5) # window sizes in blocks (None = full attention) + gen_cached_block_options = (0, 1, 4) + beam_widths = (1, 2) + + for tpb, prompt_blocks, src_wblocks, dst_wblocks, gen_cached, beam_width in itertools.product( + tpb_values, + prompt_block_counts, + window_block_options, + window_block_options, + gen_cached_block_options, + beam_widths, + ): + prompt_len = prompt_blocks * tpb - 1 # exercise the ceil in prompt_blocks + src_window = src_wblocks * tpb + 1 if src_wblocks is not None else None + dst_window = dst_wblocks * tpb + 1 if dst_wblocks is not None else None + + src_anchor = _swa_start_block(prompt_len, src_window, tpb) + dst_anchor = max( + _swa_start_block(prompt_len, dst_window, tpb), min(gen_cached, prompt_blocks) + ) + # Beam search never combines with eviction/reuse trims in production + # (front-block eviction and chunked transfer assert beam_width == 1). + if beam_width > 1 and (src_anchor > 0 or dst_anchor > 0): + continue + + case = ( + f"tpb={tpb} prompt_blocks={prompt_blocks} src_w={src_wblocks} " + f"dst_w={dst_wblocks} cached={gen_cached} beam={beam_width}" + ) + src_span = np.arange(_SRC_ID_BASE + src_anchor, _SRC_ID_BASE + prompt_blocks) + dst_span = np.arange(_DST_ID_BASE + dst_anchor, _DST_ID_BASE + prompt_blocks) + n_tails = beam_width - 1 + src_ids = np.concatenate([src_span, np.arange(_SRC_TAIL_BASE, _SRC_TAIL_BASE + n_tails)]) + dst_ids = np.concatenate([dst_span, np.arange(_DST_TAIL_BASE, _DST_TAIL_BASE + n_tails)]) + + sender = _make_projection_sender(num_groups=1, tpb=tpb) + + # --- monolithic slice --- + task = KVSendTask( + KVSlice( + is_last_slice=True, + block_ids_per_layer_groups=[src_ids.astype(np.int64)], + first_ordinals=[src_anchor], + ), + _make_params(), + slice_id=0, + prompt_len=prompt_len, + beam_width=beam_width, + ) + meta = sender._build_kv_write_meta( + task, _grid_req_info(dst_ids.astype(np.int64), dst_anchor) + ) + + expected_ordinals = list(range(max(src_anchor, dst_anchor), prompt_blocks)) + got_pairs = list(zip(meta.src_ptrs.tolist(), meta.dst_ptrs.tolist())) + beam0_pairs = [p for p in got_pairs if p[0] < _SRC_TAIL_BASE] + tail_pairs = [p for p in got_pairs if p[0] >= _SRC_TAIL_BASE] + # Every pair maps equal ordinals on both sides. + assert all(s - _SRC_ID_BASE == d - _DST_ID_BASE for s, d in beam0_pairs), case + # The transferred range is exactly the intersection of the two spans. + assert [s - _SRC_ID_BASE for s, _ in beam0_pairs] == expected_ordinals, case + if expected_ordinals and n_tails: + assert tail_pairs == [ + (_SRC_TAIL_BASE + i, _DST_TAIL_BASE + i) for i in range(n_tails) + ], case + else: + assert tail_pairs == [], case + + # --- chunked slices (full-attention source only; SWA groups are + # deferred whole to the final chunk by _build_prefill_chunk) --- + if src_wblocks is not None or beam_width > 1: + continue + for chunk_blocks in (1, 2, prompt_blocks): + seen_ordinals = [] + bounds = list(range(0, prompt_blocks, chunk_blocks)) + [prompt_blocks] + for slice_id, (c0, c1) in enumerate(zip(bounds[:-1], bounds[1:])): + if c0 >= c1: + continue + chunk_ids = np.arange(_SRC_ID_BASE + c0, _SRC_ID_BASE + c1) + chunk_task = KVSendTask( + KVSlice( + is_last_slice=(c1 == prompt_blocks), + block_ids_per_layer_groups=[chunk_ids.astype(np.int64)], + first_ordinals=[c0], + ), + _make_params(), + slice_id=slice_id, + prompt_len=prompt_len, + ) + chunk_meta = sender._build_kv_write_meta( + chunk_task, _grid_req_info(dst_ids.astype(np.int64), dst_anchor) + ) + for s, d in zip(chunk_meta.src_ptrs.tolist(), chunk_meta.dst_ptrs.tolist()): + assert s - _SRC_ID_BASE == d - _DST_ID_BASE, f"{case} chunk={chunk_blocks}" + seen_ordinals.append(s - _SRC_ID_BASE) + # Chunks tile the intersection exactly once, in order. + assert seen_ordinals == expected_ordinals, f"{case} chunk={chunk_blocks}" + + +# --------------------------------------------------------------------------- +# RecvReqInfo wire format +# --------------------------------------------------------------------------- + + +def test_recv_req_info_first_ordinals_roundtrip(): + """first_ordinals survive to_bytes/from_bytes unchanged.""" + info = RecvReqInfo( + sender_req_id=7, + instance_name="decode", + instance_rank=3, + block_ids_per_layer_groups=[ + np.array([104, 105], dtype=np.int64), + np.array([], dtype=np.int64), + ], + unique_rid=42, + first_ordinals=[6, 0], + aux_slot=1, + slice_id=2, + bounce_dst_base=None, + ) + + restored = RecvReqInfo.from_bytes(info.to_bytes()) + + assert restored.first_ordinals == [6, 0] + assert restored.sender_req_id == 7 + assert restored.instance_name == "decode" + assert restored.instance_rank == 3 + assert restored.unique_rid == 42 + assert restored.aux_slot == 1 + assert restored.slice_id == 2 + assert np.array_equal(restored.block_ids_per_layer_groups[0], [104, 105]) + assert restored.block_ids_per_layer_groups[1].size == 0 + + +def test_recv_req_info_rejects_pre_anchor_wire_layout(): + """Bytes from a peer still sending dst_start_token fail loudly, never misalign.""" + old_layout = msgpack.packb( + { + "sender_req_id": 7, + "instance_name": "decode", + "instance_rank": 0, + "block_ids_per_layer_groups": [np.array([104], dtype=np.int64).tobytes()], + "unique_rid": 42, + "dst_start_token": 128, # pre-anchor field, no first_ordinals + "aux_slot": None, + "slice_id": 0, + "bounce_dst_base": None, + } + ) + + with pytest.raises(TypeError): + RecvReqInfo.from_bytes(old_layout) + + +# --------------------------------------------------------------------------- +# Receiver-side anchor validation +# --------------------------------------------------------------------------- + + +def _make_recv_task(kv_slice: KVSlice) -> SimpleNamespace: + return SimpleNamespace( + _unique_rid=42, + _params=SimpleNamespace(ctx_request_id=7, disagg_request_id=42), + _kv_slice=kv_slice, + _aux_slot=None, + slice_id=0, + ) + + +def _make_bare_receiver() -> Receiver: + receiver = Receiver.__new__(Receiver) + receiver._registrar = SimpleNamespace( + self_rank_info=SimpleNamespace(instance_name="decode", instance_rank=0) + ) + return receiver + + +def test_recv_req_info_requires_one_anchor_per_layer_group(): + """An unanchored slice never reaches the wire.""" + receiver = _make_bare_receiver() + kv_slice = KVSlice( + is_last_slice=True, + block_ids_per_layer_groups=[ + np.array([1, 2], dtype=np.int64), + np.array([3], dtype=np.int64), + ], + first_ordinals=[0], # one anchor missing + ) + + with pytest.raises(ValueError, match="first_ordinals"): + receiver._build_recv_req_info(_make_recv_task(kv_slice)) + + +def test_recv_req_info_carries_the_slice_anchors(): + receiver = _make_bare_receiver() + kv_slice = KVSlice( + is_last_slice=True, + block_ids_per_layer_groups=[ + np.array([1, 2], dtype=np.int64), + np.array([], dtype=np.int64), + ], + first_ordinals=[3, 0], + ) + + info = receiver._build_recv_req_info(_make_recv_task(kv_slice)) + + assert info.first_ordinals == [3, 0] + assert info.unique_rid == 42 + assert info.sender_req_id == 7 + + def test_process_kv_agent_result_resolves_task_by_receiver_slice_id(): """The receiver resolves the task using its own slice id.""" session = _make_rx_session(1) @@ -935,12 +1364,13 @@ def test_pipelined_multiple_chunks_use_real_builder_and_tx_session(): transceiver._transfer_worker = SimpleNamespace(create_tx_session=lambda _req: session) transceiver._reuse_adapter = SimpleNamespace( tokens_per_block=tokens_per_block, - get_block_ids=lambda _req, _idx, _lg: source_block_ids, + get_transfer_span=lambda _req, _idx, _lg: (source_block_ids, 0), ) transceiver._page_table = SimpleNamespace( layer_groups=[SimpleNamespace(kind=CacheKind.PAGED, sliding_window_size=None)] ) transceiver._kv_cache_manager = SimpleNamespace(tokens_per_block=tokens_per_block) + transceiver._mapping = SimpleNamespace(cp_size=1) transceiver._dp_rank = 0 transceiver._context_info_endpoint = "ctx" @@ -965,10 +1395,7 @@ def test_pipelined_multiple_chunks_use_real_builder_and_tx_session(): request.context_remaining_length = 0 transceiver.respond_and_send_async(request) - assert [task._slice.token_range for task in session.kv_tasks] == [ - TokenRange(start=0, end=2 * tokens_per_block), - TokenRange(start=2 * tokens_per_block, end=4 * tokens_per_block), - ] + assert [task._slice.first_ordinals for task in session.kv_tasks] == [[0], [2]] assert [task._slice.block_ids_per_layer_groups[0].tolist() for task in session.kv_tasks] == [ [0, 1], [2, 3], @@ -989,11 +1416,6 @@ def test_pipelined_multiple_chunks_use_real_builder_and_tx_session(): _REUSE_TOTAL_BLOCKS = 8 -def _reuse_token_range(start_block: int, end_block: int) -> TokenRange: - """A chunk's block window as the block-aligned token range that rides the slice.""" - return TokenRange(start=start_block * _REUSE_TPB, end=end_block * _REUSE_TPB) - - def _build_prefill_chunk_tokens_for( prepopulated_tokens, chunk_start_pos, @@ -1001,11 +1423,13 @@ def _build_prefill_chunk_tokens_for( resident_blocks=None, sliding_window_size=_REUSE_TOTAL_BLOCKS * _REUSE_TPB, source_block_ids=None, + span_anchor=0, ): """Drive the real _build_prefill_chunk for one chunk, in token coordinates. ``resident_blocks`` is how many blocks the mocked ``_create_kv_slice`` hands - back, and defaults to the block holding ``chunk_end_pos``. + back, and defaults to the block holding ``chunk_end_pos``. ``span_anchor`` + is the block ordinal the mocked base slice is anchored at. ``sliding_window_size`` defaults to a full-attention layer as the V1 extractor actually builds one: its groups come from max_attention_window_vec @@ -1022,9 +1446,13 @@ def _build_prefill_chunk_tokens_for( source_block_ids = np.arange(resident_blocks, dtype=np.int64) base_slice = KVSlice( block_ids_per_layer_groups=[np.asarray(source_block_ids, dtype=np.int64)], + first_ordinals=[span_anchor], ) transceiver = MagicMock() + # Bind the real predicate: a bare MagicMock's truthy auto-mock would defer + # every group (even full-attention) to the final slice. + transceiver._defers_to_final_slice = KvCacheTransceiverV2._defers_to_final_slice transceiver._kv_cache_manager.tokens_per_block = _REUSE_TPB transceiver._create_kv_slice.return_value = base_slice transceiver._page_table = SimpleNamespace( @@ -1069,7 +1497,7 @@ def test_build_prefill_chunk_rounds_unaligned_non_final_end_down(): ) assert kv_slice.is_last_slice is False - assert kv_slice.token_range == _reuse_token_range(0, 1) + assert kv_slice.first_ordinals == [0] assert np.array_equal(kv_slice.block_ids_per_layer_groups[0], np.arange(1, dtype=np.int64)) @@ -1086,7 +1514,8 @@ def test_unaligned_chunk_boundaries_tile_block_space_exactly(): ] block_spans = [ - (s.token_range.start // _REUSE_TPB, s.token_range.end // _REUSE_TPB) for s in slices + (s.first_ordinals[0], s.first_ordinals[0] + s.block_ids_per_layer_groups[0].size) + for s in slices ] assert block_spans == [(0, 1), (1, 3), (3, _REUSE_TOTAL_BLOCKS)] @@ -1129,6 +1558,7 @@ def test_swa_blocks_are_deferred_until_the_complete_final_window(): resident_blocks=_REUSE_TOTAL_BLOCKS, sliding_window_size=16, source_block_ids=np.arange(4, _REUSE_TOTAL_BLOCKS), + span_anchor=4, ) assert first_slice is None @@ -1137,6 +1567,9 @@ def test_swa_blocks_are_deferred_until_the_complete_final_window(): final_slice.block_ids_per_layer_groups[0], np.arange(4, _REUSE_TOTAL_BLOCKS, dtype=np.int64), ) + # The deferred window rides anchored at its true window-start ordinal, not + # at the final chunk's start. + assert final_slice.first_ordinals == [4] @pytest.mark.parametrize( @@ -1163,6 +1596,7 @@ def test_window_covering_the_whole_prompt_streams_like_full_attention(window_tok assert kv_slice is not None assert kv_slice.is_last_slice is False + assert kv_slice.first_ordinals == [0] assert np.array_equal(kv_slice.block_ids_per_layer_groups[0], np.arange(4, dtype=np.int64)) @@ -1174,7 +1608,7 @@ def test_unaligned_reuse_prefix_still_extends_first_chunk_to_block_zero(): chunk_end_pos=14, ) - assert kv_slice.token_range == _reuse_token_range(0, 3) + assert kv_slice.first_ordinals == [0] assert np.array_equal(kv_slice.block_ids_per_layer_groups[0], np.arange(3, dtype=np.int64)) @@ -1186,7 +1620,7 @@ def test_first_chunk_covers_ctx_prefix_reuse(): chunk_end_block=6, ) - assert kv_slice.token_range == _reuse_token_range(0, 6) + assert kv_slice.first_ordinals == [0] assert np.array_equal(kv_slice.block_ids_per_layer_groups[0], np.arange(6, dtype=np.int64)) assert kv_slice.is_last_slice is False @@ -1211,7 +1645,7 @@ def test_only_the_first_chunk_extends_to_block_zero( resident_blocks=_REUSE_TOTAL_BLOCKS, ) - assert kv_slice.token_range == _reuse_token_range(expected_start_block, chunk_end_block) + assert kv_slice.first_ordinals == [expected_start_block] assert np.array_equal( kv_slice.block_ids_per_layer_groups[0], np.arange(expected_start_block, chunk_end_block, dtype=np.int64), @@ -1221,9 +1655,9 @@ def test_only_the_first_chunk_extends_to_block_zero( def test_single_chunk_with_reuse_degenerates_to_monolithic_slice(): """One chunk plus a reuse hit yields the same slice shape a monolithic send would. - The chunk still spans [0, total_blocks), which _build_kv_write_meta addresses - exactly as an unpipelined write — see - test_whole_prompt_chunk_addresses_like_a_monolithic_slice. + The chunk still spans [0, total_blocks) anchored at 0, which + _build_kv_write_meta addresses exactly as an unpipelined write — see + test_chunked_slices_reassemble_to_the_monolithic_transfer. """ kv_slice = _build_prefill_chunk_for( prepopulated_blocks=3, @@ -1233,13 +1667,65 @@ def test_single_chunk_with_reuse_degenerates_to_monolithic_slice(): ) assert kv_slice.is_last_slice is True - assert kv_slice.token_range == _reuse_token_range(0, _REUSE_TOTAL_BLOCKS) + assert kv_slice.first_ordinals == [0] assert np.array_equal( kv_slice.block_ids_per_layer_groups[0], np.arange(_REUSE_TOTAL_BLOCKS, dtype=np.int64), ) +def test_chunk_before_an_anchored_group_span_contributes_nothing(): + """A chunk entirely before a group's anchored span sends no blocks for it.""" + # The group's span covers block ordinals [4, 8); the chunk covers [0, 4). + kv_slice = _build_prefill_chunk_tokens_for( + prepopulated_tokens=0, + chunk_start_pos=0, + chunk_end_pos=4 * _REUSE_TPB, + resident_blocks=4, + source_block_ids=np.array([104, 105, 106, 107], dtype=np.int64), + span_anchor=4, + ) + + # Nothing to send yet: the only group contributes no blocks. + assert kv_slice is None + + +def test_chunk_overlapping_an_anchored_group_span_starts_at_the_anchor(): + """The chunk's sub-slice starts at max(chunk_start, span anchor).""" + kv_slice = _build_prefill_chunk_tokens_for( + prepopulated_tokens=0, + chunk_start_pos=3 * _REUSE_TPB, + chunk_end_pos=6 * _REUSE_TPB, + resident_blocks=4, + source_block_ids=np.array([104, 105, 106, 107], dtype=np.int64), + span_anchor=4, + ) + + # Span [4, 8) ∩ chunk [3, 6) = [4, 6): blocks 104, 105 anchored at 4. + assert kv_slice.first_ordinals == [4] + assert np.array_equal( + kv_slice.block_ids_per_layer_groups[0], np.array([104, 105], dtype=np.int64) + ) + + +def test_span_not_covering_the_chunk_end_raises(): + """A short streaming-group span cannot address the chunk by ordinal. + + A ValueError (not a bare assert) so the guard survives python -O: a short + span silently transferring misaligned chunk blocks is the failure mode the + anchors exist to prevent. + """ + with pytest.raises(ValueError, match="does not cover the chunk end"): + _build_prefill_chunk_tokens_for( + prepopulated_tokens=0, + chunk_start_pos=0, + chunk_end_pos=_REUSE_TOTAL_BLOCKS * _REUSE_TPB, # final chunk: end block 8 + resident_blocks=2, + source_block_ids=np.array([104, 105], dtype=np.int64), # span [4, 6) + span_anchor=4, + ) + + # --------------------------------------------------------------------------- # Transfer activity as a dimension owned by the transceiver # --------------------------------------------------------------------------- diff --git a/tests/unittest/disaggregated/test_kv_transfer.py b/tests/unittest/disaggregated/test_kv_transfer.py index 81ced03755da..abcc8f587784 100644 --- a/tests/unittest/disaggregated/test_kv_transfer.py +++ b/tests/unittest/disaggregated/test_kv_transfer.py @@ -43,7 +43,6 @@ KVSlice, LayerRange, SessionStatus, - TokenRange, WaitResult, ) from tensorrt_llm._torch.disaggregation.native.transfer import TransferWorker, TransferWorkerConfig @@ -100,31 +99,28 @@ class KvCacheConfigV2: @pytest.mark.cpu_only -def test_token_range_valid(): - tr = TokenRange(start=0, end=10) - assert tr.start == 0 - assert tr.end == 10 - - -@pytest.mark.cpu_only -def test_token_range_allows_empty(): - """A chunk clamped past the end of the prompt covers no tokens.""" - tr = TokenRange(start=8, end=8) - assert tr.start == tr.end == 8 - - -@pytest.mark.cpu_only -def test_token_range_invalid_negative(): - with pytest.raises(ValueError, match="non-negative"): - TokenRange(start=-1, end=5) - with pytest.raises(ValueError, match="non-negative"): - TokenRange(start=0, end=-1) +def test_kv_slice_first_ordinals_parallel_to_layer_groups(): + """Each layer group's anchor rides next to its block list.""" + s = KVSlice( + block_ids_per_layer_groups=[ + np.array([7, 8], dtype=np.int64), + np.array([], dtype=np.int64), + ], + first_ordinals=[3, 0], + ) + assert s.first_ordinals == [3, 0] + assert len(s.first_ordinals) == len(s.block_ids_per_layer_groups) @pytest.mark.cpu_only -def test_token_range_invalid_start_gt_end(): - with pytest.raises(ValueError, match="Invalid range"): - TokenRange(start=10, end=3) +def test_kv_slice_empty_group_uses_anchor_zero(): + """Empty lists (deferred SWA groups, STATE with no slot) anchor at 0.""" + s = KVSlice( + block_ids_per_layer_groups=[np.array([], dtype=np.int64)], + first_ordinals=[0], + ) + assert s.first_ordinals == [0] + assert s.block_ids_per_layer_groups[0].size == 0 @pytest.mark.cpu_only @@ -157,19 +153,19 @@ def test_kv_slice_construction(): layer_range=lr, block_ids_per_layer_groups=[[1, 2, 3]], is_last_slice=True, - token_range=TokenRange(start=0, end=3), + first_ordinals=[0], ) assert s.layer_range == lr assert s.block_ids_per_layer_groups == [[1, 2, 3]] assert s.is_last_slice is True - assert s.token_range == TokenRange(start=0, end=3) + assert s.first_ordinals == [0] # Test defaults s2 = KVSlice() assert s2.layer_range is None assert s2.block_ids_per_layer_groups == [] assert s2.is_last_slice is False - assert s2.token_range is None + assert s2.first_ordinals == [] @pytest.mark.cpu_only @@ -205,10 +201,17 @@ def _send_prefill_chunks( """Build and optionally send slices through the real prefill-chunk path.""" all_block_ids = [np.asarray(ids, dtype=np.int64) for ids in all_block_ids] total_blocks = max((len(ids) for ids in all_block_ids), default=0) - base_slice = KVSlice(block_ids_per_layer_groups=all_block_ids) + base_slice = KVSlice( + block_ids_per_layer_groups=all_block_ids, + first_ordinals=[0] * len(all_block_ids), + ) session = sender_session if sender_session is not None else MagicMock() session.kv_tasks = [] transceiver = MagicMock() + # Bind the real predicate: on a bare MagicMock it returns a truthy + # auto-mock, which would defer every group (even full-attention) to the + # final slice. + transceiver._defers_to_final_slice = KvCacheTransceiverV2._defers_to_final_slice transceiver._get_or_create_send_session.return_value = session transceiver._create_kv_slice = MagicMock(return_value=base_slice) transceiver._reuse_adapter.tokens_per_block = tokens_per_block @@ -274,6 +277,9 @@ def test_build_prefill_chunk_slices_chunk_window_from_whole_prompt(): chunk_blocks = 16 transceiver = MagicMock() + # Bind the real predicate so full-attention groups stream (a bare + # MagicMock's truthy auto-mock would defer every group). + transceiver._defers_to_final_slice = KvCacheTransceiverV2._defers_to_final_slice transceiver._reuse_adapter.tokens_per_block = tokens_per_block transceiver._kv_cache_manager.tokens_per_block = tokens_per_block transceiver._kv_cache_manager.kv_cache_map = {} @@ -282,7 +288,8 @@ def test_build_prefill_chunk_slices_chunk_window_from_whole_prompt(): # A full-attention group keeps the whole prompt resident from block 0, so # _create_kv_slice describes every block regardless of prefill progress. transceiver._create_kv_slice.return_value = KVSlice( - block_ids_per_layer_groups=[np.arange(prompt_blocks, dtype=np.int64)] + block_ids_per_layer_groups=[np.arange(prompt_blocks, dtype=np.int64)], + first_ordinals=[0], ) req = MagicMock() @@ -307,20 +314,19 @@ def test_build_prefill_chunk_slices_chunk_window_from_whole_prompt(): kv_slice.block_ids_per_layer_groups[0], np.arange(chunk_start, chunk_end, dtype=np.int64), ) - assert kv_slice.token_range == TokenRange( - start=chunk_start * tokens_per_block, end=chunk_end * tokens_per_block - ) + # The chunk slice is anchored at its own first block ordinal. + assert kv_slice.first_ordinals == [chunk_start] @pytest.mark.parametrize( - "source_block_ids", + "source_block_ids,span_anchor", [ - np.arange(16, dtype=np.int64), - np.arange(9, 13, dtype=np.int64), + (np.arange(16, dtype=np.int64), 0), + (np.arange(9, 13, dtype=np.int64), 12), ], ids=["v1_full_prompt_allocation", "v2_incremental_allocation"], ) -def test_build_prefill_chunk_defers_partial_swa_chunk(source_block_ids): +def test_build_prefill_chunk_defers_partial_swa_chunk(source_block_ids, span_anchor): """A partial SWA chunk is deferred for both cache-manager allocation strategies.""" tokens_per_block = 8 prompt_blocks = 16 @@ -333,8 +339,9 @@ def test_build_prefill_chunk_defers_partial_swa_chunk(source_block_ids): transceiver._reuse_adapter = SimpleNamespace( tokens_per_block=tokens_per_block, get_cached_token_count_per_layer_group=lambda req, layer_groups: [0], - get_block_ids=lambda req, idx, lg: source_block_ids, + get_transfer_span=lambda req, idx, lg: (source_block_ids, span_anchor), ) + transceiver._mapping = SimpleNamespace(cp_size=1) transceiver._page_table = SimpleNamespace(layer_groups=[layer_group]) transceiver._kv_cache_manager = SimpleNamespace( tokens_per_block=tokens_per_block, @@ -404,10 +411,7 @@ def test_send_prefill_chunks_unaligned_boundary_splits_on_a_block(): boundary_offset_tokens=2, ) - assert [s.token_range for s in slices] == [ - TokenRange(start=0, end=4 * tokens_per_block), - TokenRange(start=4 * tokens_per_block, end=8 * tokens_per_block), - ] + assert [s.first_ordinals for s in slices] == [[0], [4]] assert np.array_equal(slices[0].block_ids_per_layer_groups[0], np.arange(4)) assert np.array_equal(slices[1].block_ids_per_layer_groups[0], np.arange(4, 8)) @@ -421,9 +425,9 @@ def test_send_prefill_chunks_multiple_layer_groups(): assert np.array_equal(slices[1].block_ids_per_layer_groups[0], np.array([4, 5, 6, 7])) assert np.array_equal(slices[0].block_ids_per_layer_groups[1], np.array([100, 101, 102, 103])) assert np.array_equal(slices[1].block_ids_per_layer_groups[1], np.array([104, 105, 106, 107])) - # Default tokens_per_block is 1, so the token range doubles as block coords. - assert slices[0].token_range == TokenRange(start=0, end=4) - assert slices[1].token_range == TokenRange(start=4, end=8) + # Each layer group carries its own explicit anchor for each chunk. + assert slices[0].first_ordinals == [0, 0] + assert slices[1].first_ordinals == [4, 4] def create_transfer_worker_setup( @@ -818,11 +822,23 @@ def get_block_data( def get_block_ids_per_layer_groups( - kv_cache_manager, transfer_worker, request_id: int, use_v2: bool, tokens_per_block: int -) -> List[List[int]]: - """Get block_ids for each layer group with window_size filtering.""" + kv_cache_manager, + transfer_worker, + request_id: int, + use_v2: bool, + tokens_per_block: int, + request_len: int, +): + """Get (block_ids, first_block_ordinal) for each layer group. + + Windowed groups are trimmed to the active window; the anchor makes the + trimmed list's position explicit (KVSlice.first_ordinals) instead of + leaving it to be inferred from the list length. + """ page_table = transfer_worker._rank_info.page_table - block_ids_per_layer_groups: List[List[int]] = [] + block_ids_per_layer_groups: List[np.ndarray] = [] + first_ordinals: List[int] = [] + prompt_blocks = (request_len + tokens_per_block - 1) // tokens_per_block for group_id, group_meta in enumerate(page_table.layer_groups): if use_v2: @@ -845,8 +861,11 @@ def get_block_ids_per_layer_groups( block_ids = block_ids[-max_blocks_in_window:] block_ids_per_layer_groups.append(np.asarray(block_ids, dtype=np.int64)) + # The lists above are resident suffixes of [0, prompt_blocks); anchor + # them explicitly at their first block ordinal. + first_ordinals.append(max(0, prompt_blocks - len(block_ids))) - return block_ids_per_layer_groups + return block_ids_per_layer_groups, first_ordinals def add_and_verify_request( @@ -962,32 +981,38 @@ def add_and_verify_request( [(gen_request.py_request_id, gen_request.prompt_len, 1)], [gen_request] ) - # Get block_ids per layer_group with window_size filtering - ctx_block_ids_per_groups = [ + # Get (block_ids, anchors) per layer_group with window_size filtering + ctx_spans_per_rank = [ get_block_ids_per_layer_groups( ctx_kv_cache_manager, ctx_transfer_worker, ctx_request.py_request_id, use_v2, tokens_per_block, + request_len, ) for ctx_kv_cache_manager, ctx_transfer_worker in zip( valid_ctx_kv_cache_managers, valid_ctx_transfer_workers ) ] + ctx_block_ids_per_groups = [blocks for blocks, _ in ctx_spans_per_rank] + ctx_first_ordinals = [anchors for _, anchors in ctx_spans_per_rank] - gen_block_ids_per_groups = [ + gen_spans_per_rank = [ get_block_ids_per_layer_groups( gen_kv_cache_manager, gen_transfer_worker, gen_request.py_request_id, use_v2, tokens_per_block, + request_len, ) for gen_kv_cache_manager, gen_transfer_worker in zip( valid_gen_kv_cache_managers, valid_gen_transfer_workers ) ] + gen_block_ids_per_groups = [blocks for blocks, _ in gen_spans_per_rank] + gen_first_ordinals = [anchors for _, anchors in gen_spans_per_rank] # Determine number of layer_groups num_layer_groups = len(ctx_block_ids_per_groups[0]) if ctx_block_ids_per_groups else 1 @@ -1002,8 +1027,11 @@ def add_and_verify_request( KVSlice( is_last_slice=True, block_ids_per_layer_groups=ctx_block_ids_per_group, + first_ordinals=ctx_anchors, + ) + for ctx_block_ids_per_group, ctx_anchors in zip( + ctx_block_ids_per_groups, ctx_first_ordinals ) - for ctx_block_ids_per_group in ctx_block_ids_per_groups ] for sender_session, send_kv_slice in zip(sender_sessions, send_kv_slices): sender_session.send(send_kv_slice) @@ -1019,8 +1047,11 @@ def add_and_verify_request( KVSlice( is_last_slice=True, block_ids_per_layer_groups=gen_block_ids_per_group, + first_ordinals=gen_anchors, + ) + for gen_block_ids_per_group, gen_anchors in zip( + gen_block_ids_per_groups, gen_first_ordinals ) - for gen_block_ids_per_group in gen_block_ids_per_groups ] for receiver_session, recv_kv_slice in zip(receiver_sessions, recv_kv_slices): receiver_session.receive(recv_kv_slice) @@ -1034,8 +1065,11 @@ def add_and_verify_request( KVSlice( is_last_slice=True, block_ids_per_layer_groups=gen_block_ids_per_group, + first_ordinals=gen_anchors, + ) + for gen_block_ids_per_group, gen_anchors in zip( + gen_block_ids_per_groups, gen_first_ordinals ) - for gen_block_ids_per_group in gen_block_ids_per_groups ] for receiver_session, recv_kv_slice in zip(receiver_sessions, recv_kv_slices): receiver_session.receive(recv_kv_slice) @@ -1056,8 +1090,11 @@ def add_and_verify_request( KVSlice( is_last_slice=True, block_ids_per_layer_groups=ctx_block_ids_per_group, + first_ordinals=ctx_anchors, + ) + for ctx_block_ids_per_group, ctx_anchors in zip( + ctx_block_ids_per_groups, ctx_first_ordinals ) - for ctx_block_ids_per_group in ctx_block_ids_per_groups ] for sender_session, send_kv_slice in zip(sender_sessions, send_kv_slices): sender_session.send(send_kv_slice) @@ -1363,8 +1400,9 @@ def test_transfer_with_gen_prefix_offset(use_v2, chunk_size_blocks): """Verify that only suffix blocks are transferred when gen has a prefix offset. Simulates gen-side prefix cache: ctx sends all blocks for [0, request_len), - gen only provides the suffix block list. The receiver-side prefix is - implicit in the block count; the sender derives dst_start from it. + gen only provides the suffix block list, anchored at its first block + ordinal via KVSlice.first_ordinals (positions are explicit, never derived + from the block count). """ tensorrt_llm.logger.set_level("info") tokens_per_block = 8 @@ -1439,22 +1477,29 @@ def test_transfer_with_gen_prefix_offset(use_v2, chunk_size_blocks): ) # Get block IDs - ctx_block_ids = get_block_ids_per_layer_groups(ctx_mgr, ctx_tw, 0, use_v2, tokens_per_block) - gen_block_ids = get_block_ids_per_layer_groups(gen_mgr, gen_tw, 1, use_v2, tokens_per_block) + ctx_block_ids, ctx_anchors = get_block_ids_per_layer_groups( + ctx_mgr, ctx_tw, 0, use_v2, tokens_per_block, request_len + ) + gen_block_ids, gen_anchors = get_block_ids_per_layer_groups( + gen_mgr, gen_tw, 1, use_v2, tokens_per_block, request_len + ) - # Gen: only provide suffix block IDs (skip prefix_blocks) + # Gen: only provide suffix block IDs (skip prefix_blocks); the skipped + # prefix is declared explicitly by advancing each anchor. gen_suffix_block_ids = [ np.asarray(bids[prefix_blocks:], dtype=np.int64) for bids in gen_block_ids ] + gen_suffix_anchors = [anchor + prefix_blocks for anchor in gen_anchors] try: tx = ctx_tw.create_tx_session(ctx_request) - # Gen receives only the suffix list; dst_start is derived from block count. + # Gen receives only the suffix list, anchored at its true block ordinal. rx = gen_tw.create_rx_session(gen_request) recv_slice = KVSlice( is_last_slice=True, block_ids_per_layer_groups=gen_suffix_block_ids, + first_ordinals=gen_suffix_anchors, ) rx.receive(recv_slice) @@ -1463,6 +1508,7 @@ def test_transfer_with_gen_prefix_offset(use_v2, chunk_size_blocks): KVSlice( is_last_slice=True, block_ids_per_layer_groups=ctx_block_ids, + first_ordinals=ctx_anchors, ) ) else: @@ -1603,7 +1649,11 @@ def test_session_cancel_after_send(): # changed shape across versions). page_table = ctx_transfer_worker._rank_info.page_table block_ids_per_groups = [np.array([], dtype=np.int64) for _ in page_table.layer_groups] - kv_slice = KVSlice(is_last_slice=True, block_ids_per_layer_groups=block_ids_per_groups) + kv_slice = KVSlice( + is_last_slice=True, + block_ids_per_layer_groups=block_ids_per_groups, + first_ordinals=[0] * len(block_ids_per_groups), + ) tx_session.send(kv_slice) # No receiver registered yet; task is INIT. @@ -1682,12 +1732,16 @@ def _setup_chunked_request(setup, ctx_request_id, gen_request_id, request_len): [(gen_request.py_request_id, request_len, 1)], [gen_request] ) - ctx_block_ids = [ - get_block_ids_per_layer_groups(mgr, tw, ctx_request.py_request_id, use_v2, tokens_per_block) + ctx_spans = [ + get_block_ids_per_layer_groups( + mgr, tw, ctx_request.py_request_id, use_v2, tokens_per_block, request_len + ) for mgr, tw in zip(ctx_kv_cache_managers, ctx_transfer_workers, strict=True) ] - gen_block_ids = [ - get_block_ids_per_layer_groups(mgr, tw, gen_request.py_request_id, use_v2, tokens_per_block) + gen_spans = [ + get_block_ids_per_layer_groups( + mgr, tw, gen_request.py_request_id, use_v2, tokens_per_block, request_len + ) for mgr, tw in zip(gen_kv_cache_managers, gen_transfer_workers, strict=True) ] @@ -1696,8 +1750,10 @@ def _setup_chunked_request(setup, ctx_request_id, gen_request_id, request_len): "gen_request": gen_request, "ctx_kv_caches": ctx_kv_caches, "gen_kv_caches": gen_kv_caches, - "ctx_block_ids": ctx_block_ids, - "gen_block_ids": gen_block_ids, + "ctx_block_ids": [blocks for blocks, _ in ctx_spans], + "ctx_first_ordinals": [anchors for _, anchors in ctx_spans], + "gen_block_ids": [blocks for blocks, _ in gen_spans], + "gen_first_ordinals": [anchors for _, anchors in gen_spans], } @@ -1852,7 +1908,11 @@ def test_session_has_transferring_tasks_false(): # TxSession: after send(), task is INIT (no receiver → not yet dispatched) page_table = ctx_transfer_worker._rank_info.page_table block_ids_per_groups = [np.array([], dtype=np.int64) for _ in page_table.layer_groups] - kv_slice = KVSlice(is_last_slice=True, block_ids_per_layer_groups=block_ids_per_groups) + kv_slice = KVSlice( + is_last_slice=True, + block_ids_per_layer_groups=block_ids_per_groups, + first_ordinals=[0] * len(block_ids_per_groups), + ) tx_session.send(kv_slice) assert not tx_session.has_transferring_tasks() tx_session.close() @@ -1920,6 +1980,7 @@ def make_gen_request(request_id): empty_slice = KVSlice( is_last_slice=True, block_ids_per_layer_groups=[np.array([], dtype=np.int64) for _ in page_table.layer_groups], + first_ordinals=[0] * len(page_table.layer_groups), ) validate_calls = [] @@ -2004,10 +2065,13 @@ def add_and_verify_pipelined_request( receiver_sessions = [ tw.create_rx_session(ctx_info["gen_request"]) for tw in gen_transfer_workers ] - for recv_session, block_ids_per_groups in zip(receiver_sessions, gen_block_ids, strict=True): + for recv_session, block_ids_per_groups, anchors in zip( + receiver_sessions, gen_block_ids, ctx_info["gen_first_ordinals"], strict=True + ): full_slice = KVSlice( is_last_slice=True, block_ids_per_layer_groups=block_ids_per_groups, + first_ordinals=anchors, ) recv_session.receive(full_slice) diff --git a/tests/unittest/disaggregated/test_kv_transfer_mp.py b/tests/unittest/disaggregated/test_kv_transfer_mp.py index f5aa124054a7..10a074c4847e 100644 --- a/tests/unittest/disaggregated/test_kv_transfer_mp.py +++ b/tests/unittest/disaggregated/test_kv_transfer_mp.py @@ -345,6 +345,7 @@ def process_and_verify_request( send_kv_slice = KVSlice( is_last_slice=True, block_ids_per_layer_groups=[block_ids], + first_ordinals=[0], ) sender_session.send(send_kv_slice) @@ -390,6 +391,7 @@ def process_and_verify_request( recv_kv_slice = KVSlice( is_last_slice=True, block_ids_per_layer_groups=[block_ids], + first_ordinals=[0], ) receiver_session.receive(recv_kv_slice)